Optimizing a loop body expressed with tvm.extern

Questions
1

Hi all,
As an alternative approach to expressing a computation with tvm.compute, I’m using a tvm.extern call to a Python implementation of the computation with @tvm.register_func.
While both tvm.compute and tvm.extern return a tvm.tensor.Tensor object, a schedule s created with the latter does not give access to loop axes of the computation for optimization.

Eg.

B = tvm.extern(args)
s = tvm.create_schedule(B.op)
x = s[B].op.axis

AttributeError: '<class 'tvm.tensor.ExternOp'>' object has no attribute 'axis'

Is there a way to access computation axes in this case?

Thanks in advance,

2

The documentation states the following, and I believe axis and reduce_axis are not defined when the tensor is anything other than tvm.tensor.ComputeOp.

class tvm.tensor.ComputeOp(handle) Bases: tvm.tensor.Operation

Compute operation.

axis Represent axis of IterVar, only defined when it is a ComputeOp

reduce_axis Represent axis of reductions, only defined when it is a ComputeOp

3

extern op cannot be scheduled because we get no control over the internal implementation of the op.

3 Likes
4

May I ask, in this case, the only way to optimize the extern op is to do it by hand (with tir intrinsics) right ? For example like topi.sparse.csrmm implementation, no schedule configs could be associated with it.