Hzfengsy commented on code in PR #14931:
URL: https://github.com/apache/tvm/pull/14931#discussion_r1205512727
##########
tests/python/relax/test_transform_fuse_tir.py:
##########
@@ -1003,5 +1003,86 @@ def main(
_check(Module, Expected)
+def test_tir_expression_in_shape():
+ @I.ir_module
+ class Module:
+ @R.function
+ def fused_transpose_matmul(
+ x: R.Tensor((3, 4), dtype="float32"),
+ y: R.Tensor(("n - 1", 4), dtype="float32"),
+ tir_vars: R.Shape(["n"]),
+ ) -> R.Tensor(("n - 1", 3), dtype="float32"):
+ R.func_attr({"Primitive": 1})
+ with R.dataflow():
+ lv = R.emit_te(topi.transpose, x)
+ gv = R.emit_te(topi.matmul, y, lv)
+ R.output(gv)
+ return gv
+
+ @R.function
+ def main(
+ x: R.Tensor((3, 4), dtype="float32"),
+ y: R.Tensor(("n - 1", 4), dtype="float32"),
+ tir_vars: R.Shape(["n"]),
+ ) -> R.Tensor(("n - 1", 3), dtype="float32"):
+ cls = Module
+ with R.dataflow():
+ lv = cls.fused_transpose_matmul(x, y, tir_vars)
+ R.output(lv)
+ return lv
+
+ @I.ir_module
+ class Expected:
+ @T.prim_func
+ def fused_transpose_matmul(
+ x: T.Buffer((T.int64(3), T.int64(4)), "float32"),
+ p_y: T.handle,
+ p_output0: T.handle,
+ n: T.int64,
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ y = T.match_buffer(p_y, (n - T.int64(1), T.int64(4)))
+ var_T_matmul_intermediate = T.match_buffer(p_output0, (n -
T.int64(1), T.int64(3)))
Review Comment:
Please simplify the tir function (var naming and necessary
`reads`/`writes`). Testcase is kind of document and let's make it clear.
##########
include/tvm/tir/schedule/block_scope.h:
##########
@@ -41,7 +41,6 @@ namespace tir {
* - Parent sref: The parent reference of an sref is the block or loop
reference to the closest
schedulable statement. We define closest to be the nearest schedulable
statement of an ancestor in
the AST.
- * schedulable statement of its ancestors on the TensorIR AST.
Review Comment:
Would be great to keep the non-relax part consistent with the mainline, in
order to reduce the cost of sync.
You can send it to the mainline if you want. Otherwise, we can skip it as
it's only a nit.
--
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]
