[GitHub] [tvm] Hzfengsy commented on a diff in pull request #14189: [Unity][TVMScript] Fix prim_func lost issue in relax.emit_te

["Hzfengsy (via GitHub)" <gi...@apache.org>]

Hzfengsy commented on code in PR #14189:
URL: https://github.com/apache/tvm/pull/14189#discussion_r1125417440


##########
tests/python/relax/test_tvmscript_parser.py:
##########
@@ -184,6 +184,42 @@ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
     _check(TestModule, bb.get())
 
 
+def test_emit_te():
+    @I.ir_module
+    class EmitTE:
+        @R.function
+        def main(x: R.Tensor((10, 20), "float32")) -> R.Tensor((10, 20), dtype="float32"):
+            gv = R.emit_te(topi.add, x, x)
+            return gv
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func
+        def add(
+            rxplaceholder: T.Buffer((T.int64(10), T.int64(20)), "float32"),
+            rxplaceholder_1: T.Buffer((T.int64(10), T.int64(20)), "float32"),
+            T_add: T.Buffer((T.int64(10), T.int64(20)), "float32"),
+        ):
+            T.func_attr({"tir.noalias": True})
+            # with T.block("root"):
+            for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
+                with T.block("T_add"):
+                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+                    T.reads(rxplaceholder[v_ax0, v_ax1], rxplaceholder_1[v_ax0, v_ax1])
+                    T.writes(T_add[v_ax0, v_ax1])
+                    T_add[v_ax0, v_ax1] = (
+                        rxplaceholder[v_ax0, v_ax1] + rxplaceholder_1[v_ax0, v_ax1]
+                    )
+
+        @R.function
+        def main(x: R.Tensor((10, 20), dtype="float32")) -> R.Tensor((10, 20), dtype="float32"):
+            gv = R.call_tir(add, (x, x), out_sinfo=R.Tensor((10, 20), dtype="float32"))
+            out: R.Tensor((10, 20), dtype="float32") = gv

Review Comment:
   It would be a bit trikcy if we have such a indirection



##########
include/tvm/script/ir_builder/relax/ir.h:
##########
@@ -69,6 +69,23 @@ TVM_DLL void FuncRetStructInfo(const tvm::relax::StructInfo& ret_sinfo);
  */
 TVM_DLL void FuncRetValue(const tvm::relax::Expr& value);
 
+/*!
+ * \brief Add a Relax function or a TIR PrimFunc to the IRModuleFrame.
+ * \param func The function to be added.
+ * \param func_name_hint The name hint of the function to be added.
+ * \note If the function to be added already exists, return its
+ * GlobalVar directly.
+ * \return The global var bound to the added function.
+ */
+TVM_DLL GlobalVar AddFunction(const BaseFunc& func, String func_name_hint);

Review Comment:
   I wonder why these two functions are under `relax` folder. Should it be in `ir` folder?



-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org