Hook a multi-ouputs Python Function as Extern problem


Following the tutorial of External Tensor Functions, I find it helpful to customize my own functions. However, when I want to add a function with multiple inputs and ouputs, like:


def my_tvm_mulsome2(x1,x2, y1, y2):
print(x1, x2)
tvm.nd.array(np.matmul(x1.asnumpy(), x2.asnumpy())).copyto(y1)

I fount it hard to specify the op. I try to write the code as follows:

ctx = tvm.cpu(0)
A = tvm.placeholder((3,2), name=‘A’)
B = tvm.placeholder((2,3), name=‘B’)
C = tvm.extern(((3,3),(3,2)), [A, B], lambda ins, outs: tvm.call_packed(
“tvm.contrib.my_tvm_mulsome2”, ins[0],ins[1], outs[0], outs[1]), name=“C”)
s = tvm.create_schedule(C.op)

But it doesn’t treat it like an op and shows:

AttributeError Traceback (most recent call last)
7 # D = tvm.extern((3, 2), [A, B], lambda ins, outs: tvm.call_packed(
8 # “tvm.contrib.my_tvm_mulsome2”, ins[0],ins[1], outs[1]), name=“D”)
----> 9 s = tvm.create_schedule(C.op)

AttributeError: ‘list’ object has no attribute ‘op’

anyone can help? thank you!