[Unity] Support simple dynamic-shape-aware fusion

This PR adds support for simple dynamic-shape-aware fusion, which is the first step towards supporting dynamic shapes. The main changes are as follows:

- Fix FuncStructInfo in well-formed checks
- Renew symbolic var defs in fuse_ops to prevent malformed functions

Author: Siyuan Feng

3rdparty/cutlass

@@ -402,6 +402,18 @@ class WellFormedChecker : public relax::ExprVisitor,
     }
   }
 
+  void VisitStructInfo_(const FuncStructInfoNode* op) final {
+    if (op->params.defined()) {
+      WithMode(VisitMode::kMatchVarDef, [&]() {
+        ICHECK(mode_ == VisitMode::kMatchVarDef);
+        for (StructInfo param : op->params.value()) {
+          this->VisitStructInfo(param);
+        }
+      });
+    }
+    this->VisitStructInfo(op->ret);
+  }
+
   void VisitStructInfoExprField(const Expr& expr) final {
     if (mode_ == VisitMode::kMatchVarDef) {
       // populate symbolic var in first occurrence

src/relax/analysis/well_formed.cc

@@ -752,6 +752,7 @@ Var ExprMutator::VisitVarDef(const Var& var) {
   } else {
     LOG(FATAL) << "TypeError: Invalid type: " << var->GetTypeKey();
   }
+  this->var_remap_[var->vid] = ret;
   return ret;
 }
 

src/relax/ir/expr_functor.cc

@@ -33,6 +33,7 @@
 #include <tvm/relax/expr_functor.h>
 #include <tvm/relax/struct_info.h>
 #include <tvm/relax/transform.h>
+#include <tvm/tir/expr_functor.h>
 #include <tvm/tir/function.h>
 
 #include <optional>
@@ -344,6 +345,42 @@ class GraphCreator : public ExprVisitor {
   std::unordered_set<IndexedForwardGraph::Node*> initialized_nodes_;
 };
 
+/*!
+ * \brief Renew the definition of symbolic vars in Relax.
+ * \details This mutator is used to prevent the same symbolic var from being used in different
+ *          functions, which is malformed.
+ */
+class SymbolicVarRenewMutator : public ExprMutator, tir::ExprMutator {
+ public:
+  static Function Renew(const Function& function) {
+    SymbolicVarRenewMutator mutator;
+    return Downcast<Function>(mutator.VisitExpr(function));
+  }
+
+ private:
+  SymbolicVarRenewMutator() = default;
+  using relax::ExprMutator::VisitExpr;
+  using relax::ExprMutator::VisitExpr_;
+  using tir::ExprMutator::VisitExpr_;
+
+  PrimExpr VisitPrimExpr(const PrimExpr& expr) final { return tir::ExprMutator::VisitExpr(expr); }
+
+  PrimExpr VisitExpr_(const tir::VarNode* op) final {
+    auto it = var_map_.find(GetRef<tir::Var>(op));
+    if (it != var_map_.end()) {
+      return (*it).second;
+    } else {
+      auto n = make_object<tir::VarNode>(*op);
+      tir::Var v(n);
+      var_map_.Set(GetRef<tir::Var>(op), v);
+      return v;
+    }
+  }
+
+ private:
+  Map<tir::Var, tir::Var> var_map_;
+};
+
 /*!
  * \brief The ExprMutator used to create a new grouped function
  * \details The workflow of this ExprMutator is:
@@ -466,10 +503,10 @@ class FunctionCreator : public ExprMutator {
       body = builder_->Normalize(body);
       body = builder_->Normalize(SeqExpr({new_block}, body));
       group_attrs.Set(tvm::relax::attr::kPrimitive, Integer(1));
-      function_ = Function(/*params=*/params_,           //
-                           /*body=*/body,                //
-                           /*ret_struct_info=*/NullOpt,  //
-                           /*attrs=*/DictAttrs(group_attrs));
+      function_ = SymbolicVarRenewMutator::Renew(Function(/*params=*/params_,           //
+                                                          /*body=*/body,                //
+                                                          /*ret_struct_info=*/NullOpt,  //
+                                                          /*attrs=*/DictAttrs(group_attrs)));
     }
   }
 

src/relax/transform/fuse_ops.cc

@@ -520,5 +520,18 @@ def foo(x: R.Tensor(("m", "n"), dtype="float32")) -> R.Tensor(("m1", "n1"), dtyp
     assert not rx.analysis.well_formed(mod)
 
 
+def test_func_sinfo_well_formed():
+    @R.function
+    def foo():
+        @R.function
+        def local(x: R.Tensor(["m", "n"], "float32")):
+            return x
+
+        return local
+
+    mod = rx.transform.Normalize()(tvm.IRModule.from_expr(foo))
+    assert rx.analysis.well_formed(mod)
+
+
 if __name__ == "__main__":
     tvm.testing.main()

tests/python/relax/test_analysis_well_formed.py

@@ -1254,5 +1254,42 @@ def main(inp_0: R.Tensor((1, 784), dtype="float32"), inp_1: R.Tensor((1, 128), d
     _check(mod, Expected)
 
 
+def test_symbolic_shape_aware_fuse():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(x: R.Tensor(["n", "m"], "float32")):
+            with R.dataflow():
+                lv0 = R.emit_te(topi.add, x, R.const(1, "float32"))
+                lv1 = R.emit_te(topi.exp, lv0)
+                gv = R.emit_te(topi.squeeze, lv1)
+                R.output(gv)
+            return gv
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def fused_add_exp_squeeze(
+            x: R.Tensor(["n", "m"], "float32"), p0: R.Tensor([], "float32")
+        ) -> R.Tensor(["n", "m"], dtype="float32"):
+            R.func_attr({"Primitive": 1})
+            with R.dataflow():
+                lv0 = R.emit_te(topi.add, x, p0)
+                lv1 = R.emit_te(topi.exp, lv0)
+                gv = R.emit_te(topi.squeeze, lv1)
+                R.output(gv)
+            return gv
+
+        @R.function
+        def main(x: R.Tensor(["n", "m"], "float32")) -> R.Tensor(["n", "m"], dtype="float32"):
+            cls = Expected
+            with R.dataflow():
+                gv = cls.fused_add_exp_squeeze(x, R.const(1, "float32"))
+                R.output(gv)
+            return gv
+
+    _check(Before, Expected)
+
+
 if __name__ == "__main__":
     tvm.testing.main()