How to return a DLTensor in external function

#1

I’m interested in writing some custom external functions that return DLTensors, however I’m having issues with the python interface as the function always returns a c_void_p object. Here’s a simple sample of some code that demonstrates the problem.

 void simple_test(TVMArgs args, TVMRetValue* rv) {
  DLTensor* input = args[0]'
  *rv = input;
}

TVM_REGISTER_GLOBAL("tvm.contrib.test.simple_test").set_body(simple_test)

From python I call the function as follows:

my_func = tvm.get_global_func("tvm.contrib.test.simple_test")
test_input = tvm.nd.array(np.array([1, 2, 3]))
test_output = my_func(test_input)

Examining test_output shows its a void pointer rather than an ndarray like I expect. What am I doing wrong here?

#2

If you want to just use the customized function from python side, try to return a NDArray instead and it should work out of the box

1 Like
#3

The main reason is we need to return a managed tensor(NDArray) while DLTensor pointer is a managed pointer

1 Like
#4

Oh wow that does make it much easier, thanks for the quick response! I have one more hopefully quick question while I have your attention: is it possible to return multiple NDArrays? I couldn’t find any examples with more than one return value.

#5

We cannot atm. A quick workaround is to return a PackedFunc that maps index to NDArray

#6

Ah clever, I really appreciate the help.