[Unity][Pass] Canonicalize Bindings (#14079)

YuchenJin · slyubomirsky · tqchen · commit 4ad8d6415554 · 2023-03-05T08:38:11.000-05:00
It may be useful for some passes to collapse chains of definitions, particularly after other compiler transformations that may reduce or simplify some expressions. This pass will take chains of definitions and replace references to later definitions to the original one. It works by checking `LookupBinding` for each var use-site and replacing the var with its definition if the definition was another var. Additionally, `MatchCast` bindings where the LHS and the RHS are guaranteed to match at compile time are canonicalized into ordinary `VarBinding`s. Example: ```python y = x z = y w = z o = w p = o ``` Will be replaced with ```python y = x z = x w = x o = x p = x ``` Original PR: tlc-pack/relax#233 Co-authored-by: Steven S. Lyubomirsky <slyubomirsky@gmail.com>
diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h
@@ -127,6 +127,15 @@ TVM_DLL Pass AttachGlobalSymbol();
  */
 TVM_DLL Pass Normalize();
 
+/*!
+ * \brief Simplify a Relax module by folding var bindings and match shape nodes.
+ * May include other forms of expression simplification in the future.
+ * Best used alongside constant folding and eliminating unused bindings.
+ *
+ * \return The Pass.
+ */
+TVM_DLL Pass CanonicalizeBindings();
+
 /*!
  * \brief Bind params of function of the module to constant tensors.
  *
diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py
@@ -80,6 +80,20 @@ def Normalize() -> tvm.ir.transform.Pass:
     return _ffi_api.Normalize()  # type: ignore
 
 
+def CanonicalizeBindings() -> tvm.ir.transform.Pass:
+    """
+    Canonicalizes variable definitions
+    (e.g., if there is y = x and z = y, it replaces uses of y and z with x).
+
+    Best combined with constant folding and the elimination of unused definitions.
+
+    Returns
+    -------
+    ret: tvm.ir.transform.Pass
+    """
+    return _ffi_api.CanonicalizeBindings()  # type: ignore
+
+
 def RewriteDataflowReshape() -> tvm.ir.transform.Pass:
     """Convert all reshape-like call_tir to VM reshape operator call.
     The VM reshape operator calls will be further lowered to a CreateView
diff --git a/src/relax/transform/canonicalize_bindings.cc b/src/relax/transform/canonicalize_bindings.cc
@@ -0,0 +1,135 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file src/relax/transform/canonicalize_bindings.cc
+ * \brief Pass for simplifying modules by folding var bindings and match shape nodes.
+ *        May include other forms of simplification in the future.
+ *        Ideally should be used before constant folding and eliminating unused bindings.
+ */
+
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/relax/transform.h>
+
+namespace tvm {
+namespace relax {
+
+class BindingCanonicalizer : public ExprMutator {
+ public:
+  BindingCanonicalizer() {}
+
+  Expr VisitExpr_(const VarNode* op) override {
+    // remap first
+    Var v = Downcast<Var>(ExprMutator::VisitExpr_(op));
+    if (!CanCanonicalizeVar(v)) {
+      return Downcast<Expr>(v);
+    }
+    // visit again in case we need to do a substitution in the value
+    return ExprMutator::VisitExpr_(LookupBinding(v).as<VarNode>());
+  }
+
+  Expr VisitExpr_(const DataflowVarNode* op) override {
+    Var v = Downcast<Var>(ExprMutator::VisitExpr_(op));
+    if (!CanCanonicalizeVar(v)) {
+      return Downcast<Expr>(v);
+    }
+    return ExprMutator::VisitExpr_(LookupBinding(v).as<DataflowVarNode>());
+  }
+
+  void VisitBinding_(const VarBindingNode* binding) override {
+    // Unlike default visitor, we do not permit the checked type to change
+    // if the new value's checked type is different (this preserves user annotations)
+    Expr new_value = this->VisitExpr(binding->value);
+    Var new_var = this->VisitVarDef(binding->var);
+
+    if (new_var.same_as(binding->var) && new_value.same_as(binding->value)) {
+      this->builder_->EmitNormalized(GetRef<VarBinding>(binding));
+      return;
+    }
+
+    this->builder_->EmitNormalized(VarBinding(new_var, new_value));
+  }
+
+  void VisitBinding_(const MatchCastNode* binding) override {
+    // If we have a trivial shape check (the shape_ of LHS and RHS is the same),
+    // we can canonicalize to a var binding
+    Expr new_value = this->VisitExpr(binding->value);
+
+    // if the LHS and RHS have the same struct info, we canonicalize to a var binding instead
+    if (StructuralEqual()(binding->struct_info, GetStructInfo(new_value))) {
+      builder_->EmitNormalized(VarBinding(binding->var, new_value));
+    } else if (new_value.same_as(binding->value)) {
+      builder_->EmitNormalized(GetRef<MatchCast>(binding));
+    } else {
+      builder_->EmitNormalized(MatchCast(binding->var, new_value, binding->struct_info));
+    }
+  }
+
+ private:
+  bool AnnotationsDiffer(const ObjectRef& obj1, const ObjectRef& obj2,
+                         std::function<bool(const ObjectRef&, const ObjectRef&)> check_eq) {
+    // annotations differ if one is present but not the other
+    // or they're both present and they differ
+    bool both_present = obj1.defined() && obj2.defined();
+    bool neither_present = !obj1.defined() && !obj2.defined();
+    return !(both_present || neither_present) || (both_present && !check_eq(obj1, obj2));
+  }
+
+  bool CanCanonicalizeVar(Var v) {
+    Optional<Expr> value = LookupBinding(v);
+    // can replace only if the value is also a var
+    if (!value || !value.as<VarNode>()) {
+      return false;
+    }
+    Var parent_var = Downcast<Var>(value);
+
+    // Cases when we conservatively do not unify:
+    // 1. checked_type_ or shape_ of the child differs from that of the parent
+    //    In this case, we could be overriding user annotations.
+    // 2. If the child is a Var and the parent is a DataflowVar.
+    //    That could result in a DataflowVar leaving the current DataflowBlock.
+    bool annotations_differ = AnnotationsDiffer(v->struct_info_, parent_var->struct_info_,
+                                                [&](const ObjectRef& lhs, const ObjectRef& rhs) {
+                                                  return tvm::StructuralEqual()(lhs, rhs);
+                                                });
+    bool var_to_dataflow = (!v.as<DataflowVarNode>() && parent_var.as<DataflowVarNode>());
+    return !annotations_differ && !var_to_dataflow;
+  }
+};
+
+Expr CanonicalizeBindings(const Expr& e) { return BindingCanonicalizer().VisitExpr(e); }
+
+namespace transform {
+
+Pass CanonicalizeBindings() {
+  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
+      [=](Function f, IRModule m, PassContext pc) {
+        return Downcast<Function>(CanonicalizeBindings(f));
+      };
+  return CreateFunctionPass(pass_func, 1, "CanonicalizeBindings", {});
+}
+
+TVM_REGISTER_GLOBAL("relax.transform.CanonicalizeBindings").set_body_typed(CanonicalizeBindings);
+
+}  // namespace transform
+
+}  // namespace relax
+}  // namespace tvm
diff --git a/tests/python/relax/test_transform_canonicalize_bindings.py b/tests/python/relax/test_transform_canonicalize_bindings.py
@@ -0,0 +1,224 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import tvm
+import tvm.script
+import tvm.testing
+import pytest
+from tvm import relax
+from tvm.ir.base import assert_structural_equal
+from tvm.script import relax as R, tir as T
+
+
+def test_simple_assignments():
+    @tvm.script.ir_module
+    class TestChainAssignments:
+        @R.function
+        def main(x: R.Tensor):
+            y = x
+            z = y
+            q = z
+            p = q
+            o = p
+            return o
+
+    # a little annoying to have these unused bindings around
+    # but they can be eliminated in a separate pass
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor):
+            y = x
+            z = x
+            q = x
+            p = x
+            o = x
+            return x
+
+    new_mod = relax.transform.CanonicalizeBindings()(TestChainAssignments)
+    assert_structural_equal(new_mod, Expected)
+
+
+def test_dataflow_block():
+    @tvm.script.ir_module
+    class TestDataflowAssignments:
+        @R.function
+        def main(x: R.Tensor):
+            with R.dataflow():
+                y = R.const(1)
+                z = y
+                o = z
+                p = o
+                m = p
+                n = m
+                R.output(n)
+            return n
+
+    # a little annoying to have these unused bindings around
+    # but they can be eliminated in a separate pass
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor):
+            with R.dataflow():
+                y = R.const(1)
+                z = y
+                o = y
+                p = y
+                m = y
+                # we can't get rid of n because it leaves the block
+                n = y
+                R.output(n)
+            return n
+
+    new_mod = relax.transform.CanonicalizeBindings()(TestDataflowAssignments)
+    assert_structural_equal(new_mod, Expected)
+
+
+def test_ops():
+    @tvm.script.ir_module
+    class TestOps:
+        @R.function
+        def main(x: R.Tensor, y: R.Tensor):
+            w = y
+            q = x
+            z = R.add(w, q)
+            return R.add(q, z)
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor, y: R.Tensor):
+            w = y
+            q = x
+            z = R.add(y, x)
+            return R.add(x, z)
+
+    new_mod = relax.transform.CanonicalizeBindings()(TestOps)
+    assert_structural_equal(new_mod, Expected)
+
+
+@pytest.mark.xfail(reason="The lhs and rhs of an assignment should have the same struct info.")
+def test_casting():
+    @tvm.script.ir_module
+    class TestCasting:
+        @R.function
+        def main(x: R.Tensor) -> R.Object:
+            y = x
+            # z will be treated as object type even though it's a tensor
+            z: R.Object = y
+            return z
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor) -> R.Object:
+            y = x
+            # Cannot unify because the cast indicates user intent
+            z: R.Object = x
+            return z
+
+    new_mod = relax.transform.CanonicalizeBindings()(TestCasting)
+    assert_structural_equal(new_mod, Expected)
+
+
+def test_match_cast():
+    @tvm.script.ir_module
+    class TestMatchCast:
+        @R.function
+        def main(x: R.Tensor):
+            q = x
+            m, n = T.var("int64"), T.var("int64")
+            z = R.match_cast(q, R.Tensor((m, n)))
+            w = z
+            return w
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor):
+            q = x
+            # can't get rid of z because its shape_ is different from x's
+            m, n = T.var("int64"), T.var("int64")
+            z = R.match_cast(x, R.Tensor((m, n)))
+            w = z
+            return z
+
+    new_mod = relax.transform.CanonicalizeBindings()(TestMatchCast)
+    assert_structural_equal(new_mod, Expected)
+
+
+def test_same_shape():
+    @tvm.script.ir_module
+    class TestSameShape:
+        @R.function
+        def main(x: R.Tensor(("m", "n"), "float32")):
+            m, n = T.var("int64"), T.var("int64")
+            y = x
+            # trivial check
+            z = R.match_cast(x, R.Tensor((m, n), "float32"))
+            w = z
+            q = R.add(w, y)
+            return R.add(q, w)
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor(("m", "n"), "float32")):
+            m, n = T.var("int64"), T.var("int64")
+            y = x
+            # canonicalized into a var binding
+            z = x
+            w = x
+            q = R.add(x, x)
+            return R.add(q, x)
+
+    new_mod = relax.transform.CanonicalizeBindings()(TestSameShape)
+    assert_structural_equal(new_mod, Expected)
+
+
+def test_change_shape():
+    @tvm.script.ir_module
+    class TestChangeShape:
+        @R.function
+        def main(x: R.Tensor(("m", "n"))):
+            y = x
+            # not trivial: introduces new shape vars
+            o, p = T.var("int64"), T.var("int64")
+            z = R.match_cast(x, R.Tensor((o, p)))
+            w = z
+            q = R.add(w, y)
+            return R.add(q, w)
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor(("m", "n"))):
+            y = x
+            o, p = T.var("int64"), T.var("int64")
+            z = R.match_cast(x, R.Tensor((o, p)))
+            w = z
+            # the shape_ field on q will need to be updated
+            q = R.add(z, x)
+            return R.add(q, z)
+
+    new_mod = relax.transform.CanonicalizeBindings()(TestChangeShape)
+    assert_structural_equal(new_mod, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()
