[TVMScript] Ensure consistent struct info between assign lhs and rhs with sinfo annotation (apache#328)

* [TVMScript] Ensure consistent struct info between assign lhs and rhs with sinfo annotation

* fix

* fix

Author: Siyuan Feng

include/tvm/script/ir_builder/relax/ir.h

@@ -106,9 +106,12 @@ TVM_DLL void DataflowBlockOutput(const Array<tvm::relax::Var>& vars);
 /*!
  * \brief Emit a binding to the last binding block frame.
  * \param value The right side value of the bindings to be emitted.
+ * \param annotate_struct_info The optional struct info annotation for the emitted value.
  * \return The left side var of the emitted binding.
  */
-TVM_DLL tvm::relax::Var Emit(const tvm::relax::Expr& value);
+TVM_DLL tvm::relax::Var Emit(
+    const tvm::relax::Expr& value,
+    const Optional<tvm::relax::StructInfo>& annotate_struct_info = NullOpt);
 
 /*!
  * \brief Emit a match_cast binding to the last binding block frame.
@@ -119,18 +122,6 @@ TVM_DLL tvm::relax::Var Emit(const tvm::relax::Expr& value);
 TVM_DLL tvm::relax::Var EmitMatchCast(const tvm::relax::Expr& value,
                                       const tvm::relax::StructInfo& struct_info);
 
-///////////////////////////// Type Deduce //////////////////////////////
-
-/*!
- * \brief Annotate the struct info of a var.
- * \param var The input var to be annotated.
- * \param anno_struct_info The annotated struct info, which can be undefined.
- * \note This function will check if the type of var is compatible with the annotated type.
- * And we annotate to the var with more detailed type.
- */
-TVM_DLL void AnnotateStructInfo(const tvm::relax::Var& var,
-                                const tvm::relax::StructInfo& anno_struct_info);
-
 ///////////////////////////// If Then Else /////////////////////////////
 
 /*!

python/tvm/script/ir_builder/relax/ir.py

@@ -269,19 +269,22 @@ def wrapped(*args, **kwargs):
 ############################### Bindings ###############################
 
 
-def emit(value: Expr) -> Var:
+def emit(value: Expr, annotate_struct_info: Optional[StructInfo] = None) -> Var:
     """Emit a binding to the last binding block frame.
     Parameters
     ----------
     value: Expr
         The right side value of the bindings to be emitted.
 
+    annotate_struct_info: Optional[StructInfo]
+        The optional struct info annotation for the emitted value.
+
     Returns
     -------
     var: Var
         The left side var of the emitted binding.
     """
-    return _ffi_api.Emit(value)  # pylint: disable=no-member # type: ignore
+    return _ffi_api.Emit(value, annotate_struct_info)  # pylint: disable=no-member # type: ignore
 
 
 def emit_match_cast(value: Expr, struct_info: StructInfo) -> Var:
@@ -301,25 +304,6 @@ def emit_match_cast(value: Expr, struct_info: StructInfo) -> Var:
     return _ffi_api.EmitMatchCast(value, struct_info)  # type: ignore
 
 
-############################# Type Deduce ##############################
-
-
-def annotate_struct_info(var: Var, anno_struct_info: StructInfo) -> None:
-    """Annotate the struct info of relax var.
-
-    Parameters
-    ----------
-    var: Var
-        The input var to be annotated.
-
-
-    anno_struct_info: StructInfo
-        The annotated struct info
-
-    """
-    _ffi_api.AnnotateStructInfo(var, anno_struct_info)
-
-
 ############################# If Then Else #############################
 
 

python/tvm/script/parser/relax/parser.py

@@ -16,12 +16,13 @@
 # under the License.
 # pylint: disable=missing-docstring
 
+import functools
 import numbers
 from typing import Any, Optional, Tuple, Union
 
 from tvm import relax, tir
-from tvm.ir import Type
-from tvm.relax import StructInfo, Expr
+from tvm.ir import Type, structural_equal
+from tvm.relax import Expr, StructInfo
 from tvm.relax.utils import convert_to_expr
 from tvm.script.ir_builder.relax.frame import BlockFrame
 
@@ -32,7 +33,13 @@
 from .entry import MatchCastPair
 
 
-def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
+def bind_assign_value(
+    self: Parser,
+    node: doc.expr,
+    var_name: str,
+    value: Any,
+    anno_sinfo: Optional[StructInfo] = None,
+) -> Any:
     var_table = self.var_table.get()
 
     if isinstance(value, tir.Var):
@@ -64,19 +71,21 @@ def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -
         value = convert_to_expr(value)
     if isinstance(value, numbers.Number):
         value = R.const(value)
+
     if isinstance(value, relax.Expr):
-        var = R.emit(value)
-        # It's an internal check, so directly use assert here.
-        assert var is not None
-        IRBuilder.name(var_name, var)
-        return var
+        var = R.emit(value, anno_sinfo)
     elif isinstance(value, MatchCastPair):
+        if anno_sinfo is not None and not structural_equal(anno_sinfo, value.struct_info):
+            self.report_error(
+                node, "Cannot specify inconsistent annotation for a match cast pair. "
+            )
         var = R.emit_match_cast(value.value, value.struct_info)
-        IRBuilder.name(var_name, var)
-        return var
     else:
         raise TypeError(f"Unsupported type {type(value)} in assignment")
 
+    IRBuilder.name(var_name, var)
+    return var
+
 
 # pylint: disable=inconsistent-return-statements
 def eval_type_annotation(
@@ -213,16 +222,13 @@ def visit_assign(self: Parser, node: doc.Assign) -> None:
 def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
     lhs = node.target
     rhs = self.eval_expr(node.value)
-    ann_sinfo = self.visit_tvm_annotation(node.annotation)
+    anno_sinfo = self.visit_tvm_annotation(node.annotation)
     self.eval_assign(
         target=lhs,
         source=rhs,
-        bind_value=bind_assign_value,
+        bind_value=functools.partial(bind_assign_value, anno_sinfo=anno_sinfo),
         allow_shadowing=True,
     )
-    var = self.var_table.get().get(lhs.id)
-    assert isinstance(var, relax.Var)
-    R.ir.annotate_struct_info(var, ann_sinfo)
 
 
 @dispatch.register(token="relax", type_name="Return")

src/script/ir_builder/relax/ir.cc

@@ -17,6 +17,7 @@
  * under the License.
  */
 #include <tvm/relax/analysis.h>
+#include <tvm/relax/struct_info.h>
 #include <tvm/relax/type_analysis.h>
 #include <tvm/script/ir_builder/relax/ir.h>
 #include <tvm/tir/op.h>
@@ -195,11 +196,22 @@ TVM_REGISTER_GLOBAL("script.ir_builder.relax.DataflowBlockOutput")
 
 /////////////////////////////// Bindings ///////////////////////////////
 
-tvm::relax::Var Emit(const tvm::relax::Expr& expr) {
+tvm::relax::Var Emit(const tvm::relax::Expr& expr,
+                     const Optional<tvm::relax::StructInfo>& annotate_struct_info) {
+  using tvm::relax::GetStructInfo;
   BlockFrame block_frame = CheckBlockFrameExistAndUnended();
   const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
-  tvm::relax::Var var{nullptr};
-  var = block_builder->Emit(expr);
+  if (annotate_struct_info.defined()) {
+    const auto& sinfo = annotate_struct_info.value();
+    if (!expr->struct_info_.defined()) {
+      UpdateStructInfo(expr, sinfo);
+    } else {
+      CHECK(StructInfoBaseCheck(sinfo, GetStructInfo(expr)) != tvm::relax::BaseCheckResult::kFailL0)
+          << "Invalid annotation. Got rhs value struct info: " << GetStructInfo(expr)
+          << ", given struct info: " << sinfo;
+    }
+  }
+  tvm::relax::Var var = block_builder->Emit(expr);
   block_frame->emitted_vars.push_back(var);
   return var;
 }
@@ -217,17 +229,6 @@ tvm::relax::Var EmitMatchCast(const tvm::relax::Expr& value,
 TVM_REGISTER_GLOBAL("script.ir_builder.relax.Emit").set_body_typed(Emit);
 TVM_REGISTER_GLOBAL("script.ir_builder.relax.EmitMatchCast").set_body_typed(EmitMatchCast);
 
-///////////////////////////// Type Deduce //////////////////////////////
-
-void AnnotateStructInfo(const tvm::relax::Var& var,
-                        const tvm::relax::StructInfo& anno_struct_info) {
-  var->checked_type_ = GetStaticType(anno_struct_info);
-  var->struct_info_ = anno_struct_info;
-}
-
-TVM_REGISTER_GLOBAL("script.ir_builder.relax.AnnotateStructInfo")
-    .set_body_typed(AnnotateStructInfo);
-
 ///////////////////////////// If Then Else /////////////////////////////
 
 IfFrame If(tvm::relax::Expr condition) {

tests/python/relax/test_transform_canonicalize_bindings.py

@@ -112,6 +112,7 @@ def main(x: R.Tensor, y: R.Tensor):
     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:

tests/python/relax/test_tvmscript_parser.py

@@ -484,6 +484,7 @@ def foo(
 
     def _check_struct_info(binding, expected_sinfo):
         tvm.ir.assert_structural_equal(binding.var.struct_info, expected_sinfo)
+        tvm.ir.assert_structural_equal(binding.value.struct_info, expected_sinfo)
 
     # Cannot use block builder here because we need to check the annotated type,
     # which may be inconsistent with deduced type.
@@ -505,13 +506,31 @@ def test_annotate_override():
     def foo(x: R.Tensor):
         y = x
         # z will be treated as object type even though it's a tensor
-        z: R.Object = y
+        z: R.Object = R.add(x, y)
         return z
 
     assert isinstance(foo.ret_struct_info, relax.ObjectStructInfo)
     y_bind, z_bind = foo.body.blocks[0].bindings
-    assert isinstance(y_bind.var.checked_type, relax.DynTensorType)
-    assert isinstance(z_bind.var.checked_type, relax.ObjectType)
+    assert isinstance(y_bind.var.struct_info, relax.TensorStructInfo)
+    assert isinstance(z_bind.var.struct_info, relax.ObjectStructInfo)
+
+    with pytest.raises(tvm.error.DiagnosticError):
+
+        @R.function
+        def test(x: R.Tensor):
+            # Error: x is of Tensor StructInfo, which can not annotate to R.Shape.
+            z: R.Shape = x
+            return z
+
+    @R.function
+    def bar(x: R.Tensor):
+        # x is of Tensor StructInfo, the annotation of `z` is ignored.
+        z: R.Object = x
+        return z
+
+    assert isinstance(bar.ret_struct_info, relax.TensorStructInfo)
+    (z_bind,) = bar.body.blocks[0].bindings
+    assert isinstance(z_bind.var.struct_info, relax.TensorStructInfo)
 
 
 def test_empty_shape():