Unclear error from TE Inliner

Questions
#1

Hi, I’m trying to compile the following trivial function:

fn (%p0: Tensor[(3), bool], Primitive=1) -> Tensor[(3), int32] {
  where(%p0, 1 /* ty=int32 */, -1 /* ty=int32 */) /* ty=Tensor[(3), int32] */
}

Note that the second and third arg to where is a scalar. Since this is not supported in topi, I have updated the topi::where to handle this case (I used x(0) to turn a zero dim te::Tensor into a scalar, is this correct?):

  if (x->shape.size() == 0) {
    // x and y are scalar
    return compute(
        condition->shape,
        [&](const Array<Var>& indices) {
          Array<PrimExpr> condition_idx{indices[0]};
          return tvm::tir::Select(condition(condition_idx) != 0, x(0), y(0));
        },
        name, tag);
  } else if (condition->shape.size() != 1) {
     ...
  }

I expected this should work, but I got the following error from TE Inliner:

  File "/home/masa/projects/dev/tvm/src/te/schedule/operation_inline.cc", line 54
TVMError: Check failed: args_.size() == op->indices.size() (0 vs. 1) : 
During handling of the above exception, another exception occurred:

TVMError: Check failed: args_.size() == op->indices.size() (0 vs. 1) : 
Error during compile function
-----------------------------
#[version = "0.0.5"]
fn (%p0: Tensor[(3), bool], Primitive=1) -> Tensor[(3), int32] {
  where(%p0, 1 /* ty=int32 */, -1 /* ty=int32 */) /* ty=Tensor[(3), int32] */
}

The error is happening at the line below:

How do I solve this? I don’t understand why inlining is happening for the trivial function above (nothing can be inlined).

@tqchen

1 Like