MasterJH5574 opened a new pull request, #14632:
URL: https://github.com/apache/tvm/pull/14632
This PR changes the naming behavior of call_te. With this PR, all tensors
created by call_te will be named in alphabetical order (`A`, `B`, `C`...). When
the number of tensors exceeds 26, the trailing tensors will be named as
`input26`, `input27` and so on so forth.
### Background
Prior to this PR, call_te uses the default name "rxplaceholder" for every
tensor. On one hand, this name is too long and doesn't look clean in the
printed TVMScript. On the other hand, using this name for every tensor is
problematic from the perspective of TIR and will lead to scheduling error. For
example, previously, the following code snippet will print the TIR function
where two cache-read blocks having the same name, which is not a legal TIR
function.
```python
bb = relax.BlockBuilder()
x = relax.Var("x", R.Tensor((2, 3), "float32"))
y = relax.Var("y", R.Tensor((3, 4), "float32"))
with bb.function("main", [x, y]):
gv = bb.emit_te(topi.nn.matmul, x, y)
bb.emit_func_output(gv)
sch = tir.Schedule(bb.get())
sch.work_on("matmul")
sch.cache_read("T_matmul_NN", 0, "global")
sch.cache_read("T_matmul_NN", 1, "global")
print(sch.mod["matmul"].script())
## Output:
@T.prim_func
def matmul(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"),
rxplaceholder_1: T.Buffer((T.int64(3), T.int64(4)), "float32"), T_matmul_NN:
T.Buffer((T.int64(2), T.int64(4)), "float32")):
T.func_attr({"layout_free_buffers": [1], "tir.noalias": T.bool(True)})
# with T.block("root"):
rxplaceholder_global = T.alloc_buffer((T.int64(2), T.int64(3)))
rxplaceholder_global_1 = T.alloc_buffer((T.int64(3), T.int64(4)))
for ax0, ax1 in T.grid(T.int64(3), T.int64(4)):
with T.block("rxplaceholder_global"):
v0, v1 = T.axis.remap("SS", [ax0, ax1])
T.reads(rxplaceholder_1[v0, v1])
T.writes(rxplaceholder_global_1[v0, v1])
rxplaceholder_global_1[v0, v1] = rxplaceholder_1[v0, v1]
for ax0, ax1 in T.grid(T.int64(2), T.int64(3)):
with T.block("rxplaceholder_global"):
v0, v1 = T.axis.remap("SS", [ax0, ax1])
T.reads(rxplaceholder[v0, v1])
T.writes(rxplaceholder_global[v0, v1])
rxplaceholder_global[v0, v1] = rxplaceholder[v0, v1]
for i, j, k in T.grid(T.int64(2), T.int64(4), T.int64(3)):
with T.block("T_matmul_NN"):
v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
T.reads(rxplaceholder_global[v_i, v_k],
rxplaceholder_global_1[v_k, v_j])
T.writes(T_matmul_NN[v_i, v_j])
with T.init():
T_matmul_NN[v_i, v_j] = T.float32(0)
T_matmul_NN[v_i, v_j] = T_matmul_NN[v_i, v_j] +
rxplaceholder_global[v_i, v_k] * rxplaceholder_global_1[v_k, v_j]
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
