Is there a way to print elements in the TensorIR program itself?

For eg. I have a simple program like:
```
import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
import numpy as np

@T.prim_func
def main(A: T.Buffer[(128,), "float32"], 
        B: T.Buffer[(128,), "float32"], 
        C: T.Buffer[(128,), "float32"]):
    T.func_attr({"global_symbol": "main", "tir.noalias": True})
    for i in range(128):
        with T.block("C"):
            vi = T.axis.spatial(128, I)
            // TODO: Print A[vi]
            C[vi] = A[vi] + B[vi]

mod = tvm.IRModule.from_expr(main)

# Build for CPU
target = "llvm"
f = tvm.build(mod, target=target)

# Create numpy arrays
a_np = np.array([1.0] * 128, dtype="float32")
b_np = np.array([2.0] * 128, dtype="float32")
c_np = np.array([0.0] * 128, dtype="float32")

# Create TVM arrays
dev = tvm.cpu(0)
a = tvm.nd.array(a_np, dev)
b = tvm.nd.array(b_np, dev)
c = tvm.nd.array(c_np, dev)

# Run the function
f(a, b, c)
```

Can I print `A[vi]` like in the `TODO`?





---
[Visit 
Topic](https://discuss.tvm.apache.org/t/print-values-in-tensor-ir/18385/1) 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/bbcdbfa9a6ff861c0deff8a18122732c22f736e8a3057e7018d8d2ef378d96e7).

Reply via email to