[quote="reikdas, post:5, topic:18371"]
`out = R.call_dps_packed("foo", (a), out_sinfo=R.Tensor((n,), dtype="float32"))`
[/quote]
Hi @reikdas, your `foo` doesn't make any effective change to `out`, since the
line of `out = A.numpy()` is pointing the variable `out` to the object of
`A.numpy()`, rather than in-place mutating the input NDArray `out`.
There are two ways to make it properly work. The first way is to use
`out.copyfrom(A)`, in which way there will be a cuda memory copy from `A` to
`out`:
```python
@tvm.register_func("foo", override=True)
def foo(A: tvm.nd.NDArray, out: tvm.nd.NDArray):
out.copyfrom(A)
print(out)
```
The second approach is to use `call_pure_packed`, in which way your customized
`foo` needs to handle memory allocation for output. (Or like the example below,
`foo` just returns `A` instead of allocating new NDArray, where the return
value shares the same memory as `A`).
```python
import numpy as np
import tvm
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tir as T
@tvm.register_func("foo", override=True)
def foo(A):
return A
@I.ir_module
class Module:
@R.function
def main(a: R.Tensor(("n",), dtype="float32")):
with R.dataflow():
n = T.int64()
out = R.call_pure_packed("foo", a, sinfo_args=R.Tensor((n,),
dtype="float32"))
R.output(out)
return out
if __name__ == "__main__":
ex = tvm.relax.build(Module, target="llvm")
vm = tvm.relax.VirtualMachine(ex, tvm.cpu())
a = tvm.nd.array(np.array([1, 2, 3, 4, 5], dtype=np.float32),
device=tvm.cpu())
out = vm["main"](a)
print(out)
```
Let me know if this makes sense to you.
---
[Visit
Topic](https://discuss.tvm.apache.org/t/provide-external-functions-as-tvmscript/18371/6)
to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click
here](https://discuss.tvm.apache.org/email/unsubscribe/1e10e374b6ac25d649dc13c02c73ac25fda4c8984c3bf835ddaa9211280ae320).
