[REFACTOR][TIR] Rename LetStmt to Bind and flatten to sequential semantics

include/tvm/script/ir_builder/tir/frame.h

@@ -341,48 +341,6 @@ class AssertFrame : public TIRFrame {
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AssertFrame, TIRFrame, AssertFrameNode);
 };
 
-/*!
- * \brief A frame represents the let binding expression, which binds a var.
- *
- * \sa LetFrameNode
- */
-class LetFrameNode : public TIRFrameNode {
- public:
-  /*! \brief The variable we bind to */
-  tvm::tir::Var var;
-  /*! \brief The value we bind var to */
-  PrimExpr value;
-
-  static void RegisterReflection() {
-    namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<LetFrameNode>()
-        .def_ro("var", &LetFrameNode::var)
-        .def_ro("value", &LetFrameNode::value);
-  }
-  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tir.LetFrame", LetFrameNode, TIRFrameNode);
-
- public:
-  /*!
-   * \brief The method called when exiting RAII scope.
-   * \sa tvm::support::With
-   */
-  void ExitWithScope() final;
-};
-
-/*!
- * \brief Managed reference to LetFrameNode.
- *
- * \sa LetFrameNode
- */
-class LetFrame : public TIRFrame {
- public:
-  explicit LetFrame(ObjectPtr<LetFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
-    TVM_FFI_ICHECK(data != nullptr);
-    data_ = std::move(data);
-  }
-  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LetFrame, TIRFrame, LetFrameNode);
-};
-
 /*!
  * \brief The LaunchThreadFrameNode.
  * \note It is used only inside a PrimFunc.

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

@@ -294,16 +294,19 @@ AssertFrame Assert(PrimExpr condition, ffi::String error_kind,
                    ffi::Array<ffi::String> message_parts);
 
 /*!
- * \brief The let binding.
+ * \brief Create a Bind (variable binding).
+ *
+ * Emits a flat Bind statement to the current frame and returns the bound variable.
+ *
  * \param value The value to be bound.
- * \param type_annotation  The type annotation of the let binding.
+ * \param type_annotation  The type annotation of the binding.
  *                         Usually it is used for fine-grained var typing,
  *                         particularly, PointerType.
  * \param var The variable to be bound. If not specified, a new variable will be created.
- * \return The created LetFrame.
+ * \return The bound Var.
  */
-LetFrame LetStmt(PrimExpr value, ffi::Optional<Type> type_annotation = std::nullopt,
-                 ffi::Optional<Var> var = std::nullopt);
+Var Bind(PrimExpr value, ffi::Optional<Type> type_annotation = std::nullopt,
+         ffi::Optional<Var> var = std::nullopt);
 
 /*!
  * \brief The allocate node.

include/tvm/script/printer/ir_docsifier_functor.h

@@ -81,6 +81,7 @@ class IRDocsifierFunctor {
                           << runtime::Object::TypeIndex2Key(type_index) << " (token: " << token
                           << ")"
                           << ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
+    TVM_FFI_UNREACHABLE();
   }
 
   /*!

include/tvm/tir/stmt.h

@@ -68,37 +68,38 @@ class Stmt : public ObjectRef {
 };
 
 /*!
- * \brief Let binding, bind var to value, then run body.
+ * \brief Bind a variable to a value in the enclosing scope.
+ *
+ * BindNode has no body field. The bound variable is visible
+ * in all subsequent statements within the same enclosing scope (SeqStmt,
+ * ForNode.body, etc.). This enables flat (non-nested) IR sequences.
  */
-class LetStmtNode : public StmtNode {
+class BindNode : public StmtNode {
  public:
-  /*! \brief The variable. */
+  /*! \brief The variable being bound. */
   Var var;
-  /*! \brief The value to be bound. */
+  /*! \brief The value to bind to the variable. */
   PrimExpr value;
-  /*! \brief The body block. */
-  Stmt body;
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<LetStmtNode>()
-        .def_ro("var", &LetStmtNode::var, refl::AttachFieldFlag::SEqHashDef())
-        .def_ro("value", &LetStmtNode::value)
-        .def_ro("body", &LetStmtNode::body);
+    refl::ObjectDef<BindNode>()
+        .def_ro("var", &BindNode::var, refl::AttachFieldFlag::SEqHashDef())
+        .def_ro("value", &BindNode::value);
   }
-  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.LetStmt", LetStmtNode, StmtNode);
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.Bind", BindNode, StmtNode);
 };
 
 /*!
- * \brief Managed reference to LetStmtNode.
- * \sa LetStmtNode
+ * \brief Managed reference to BindNode.
+ * \sa BindNode
  */
-class LetStmt : public Stmt {
+class Bind : public Stmt {
  public:
-  TVM_DLL LetStmt(Var var, PrimExpr value, Stmt body, Span span = Span());
+  TVM_DLL Bind(Var var, PrimExpr value, Span span = Span());
 
-  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(LetStmt, Stmt, LetStmtNode);
-  TVM_DEFINE_OBJECT_REF_COW_METHOD(LetStmtNode);
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Bind, Stmt, BindNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(BindNode);
 };
 
 /*!
@@ -978,6 +979,7 @@ inline const char* ForKind2String(ForKind t) {
       return "thread_binding";
   }
   TVM_FFI_THROW(InternalError) << "Unknown ForKind" << t;
+  TVM_FFI_UNREACHABLE();
 }
 
 }  // namespace tir

include/tvm/tir/stmt_functor.h

@@ -84,7 +84,7 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
     return vtable(n, this, std::forward<Args>(args)...);
   }
   // Functions that can be overriden by subclass
-  virtual R VisitStmt_(const LetStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
+  virtual R VisitStmt_(const BindNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const AttrStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const IfThenElseNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const ForNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
@@ -106,7 +106,7 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
   // initialize the vtable.
   static FType InitVTable() {
     FType vtable;
-    IR_STMT_FUNCTOR_DISPATCH(LetStmtNode);
+    IR_STMT_FUNCTOR_DISPATCH(BindNode);
     IR_STMT_FUNCTOR_DISPATCH(AttrStmtNode);
     IR_STMT_FUNCTOR_DISPATCH(IfThenElseNode);
     IR_STMT_FUNCTOR_DISPATCH(ForNode);
@@ -159,9 +159,9 @@ class TVM_DLL StmtVisitor : protected StmtFunctor<void(const Stmt&)> {
    */
   virtual void VisitBufferUse(const Buffer& buffer);
   // statement visitor
+  void VisitStmt_(const BindNode* op) override;
   void VisitStmt_(const AttrStmtNode* op) override;
   void VisitStmt_(const IfThenElseNode* op) override;
-  void VisitStmt_(const LetStmtNode* op) override;
   void VisitStmt_(const ForNode* op) override;
   void VisitStmt_(const WhileNode* op) override;
   void VisitStmt_(const AllocBufferNode* op) override;
@@ -273,9 +273,9 @@ class TVM_DLL StmtMutator : protected StmtFunctor<Stmt(const Stmt&)> {
    */
   virtual Buffer VisitBufferUse(const Buffer& buffer);
   // statement visitor
+  Stmt VisitStmt_(const BindNode* op) override;
   Stmt VisitStmt_(const AttrStmtNode* op) override;
   Stmt VisitStmt_(const IfThenElseNode* op) override;
-  Stmt VisitStmt_(const LetStmtNode* op) override;
   Stmt VisitStmt_(const ForNode* op) override;
   Stmt VisitStmt_(const WhileNode* op) override;
   Stmt VisitStmt_(const AllocBufferNode* op) override;

include/tvm/tir/var.h

@@ -42,7 +42,7 @@ namespace tir {
  * - Allocate
  * - For
  * - Let
- * - LetStmt
+ * - Bind
  */
 class VarNode : public PrimExprNode {
  public:

python/tvm/script/ir_builder/tir/frame.py

@@ -50,13 +50,6 @@ def __enter__(self) -> Var | list[Var]:  # type: ignore[override]
 class AssertFrame(TIRFrame): ...
 
 
-@_register_object("script.ir_builder.tir.LetFrame")
-class LetFrame(TIRFrame):
-    def __enter__(self) -> Var:
-        super().__enter__()
-        return self.var
-
-
 @_register_object("script.ir_builder.tir.AllocateFrame")
 class AllocateFrame(TIRFrame):
     def __enter__(self) -> Buffer:

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

@@ -983,35 +983,37 @@ def Assert(condition: PrimExpr, message, error_kind: str = "RuntimeError") -> fr
     return _ffi_api.Assert(condition, error_kind, message)  # type: ignore[attr-defined] # pylint: disable=no-member
 
 
-def LetStmt(  # pylint: disable=invalid-name
+def Bind(  # pylint: disable=invalid-name
     value: PrimExpr,
     type_annotation: Type | None = None,  # pylint: disable=redefined-outer-name
     *,
     var: Var | None = None,  # pylint: disable=redefined-outer-name
-) -> frame.LetFrame:
-    """Create a LetStmt binding
+) -> Var:
+    """Create a Bind (variable binding).
+
+    Emits a flat Bind statement to the current frame and returns the bound variable.
 
     Parameters
     ----------
     value : PrimExpr
         The value to be bound.
     type_annotation : Optional[Type] = None
-        The type annotation of the let binding. Usually it is used for fine-grained var typing,
+        The type annotation of the binding. Usually it is used for fine-grained var typing,
         particularly, PointerType.
     var : Optional[Var] = None
         The variable to bind. If not specified, a new variable will be created.
 
     Returns
     -------
-    let_frame : frame.LetFrame
-        The result LetFrame.
+    var : Var
+        The bound variable.
     """
     if type_annotation is not None:
         if callable(type_annotation):
             type_annotation = type_annotation()
         if isinstance(type_annotation, Var):
             type_annotation = type_annotation.type_annotation
-    return _ffi_api.LetStmt(value, type_annotation, var)  # type: ignore[attr-defined] # pylint: disable=no-member
+    return _ffi_api.Bind(value, type_annotation, var)  # type: ignore[attr-defined] # pylint: disable=no-member
 
 
 def Let(  # pylint: disable=invalid-name
@@ -1028,7 +1030,7 @@ def let(
     v: Var,
     value: PrimExpr,
     body: PrimExpr = None,
-) -> frame.LetFrame:
+) -> Var:
     """Create a new let binding.
 
     Parameters
@@ -1044,17 +1046,17 @@ def let(
 
     Returns
     -------
-    res : frame.LetFrame
-        The result LetFrame.
+    res : Var
+        The bound variable.
     """
 
     @deprecated("T.let", "T.Let")
     def let_expr(v: Var, value: PrimExpr, body: PrimExpr) -> PrimExpr:
         return tir.Let(v, value, body)
 
-    @deprecated("T.let", "T.LetStmt")
-    def let_stmt(v: Var, value: PrimExpr) -> frame.LetFrame:
-        return _ffi_api.LegacyLetStmt(v, value)  # type: ignore[attr-defined] # pylint: disable=no-member
+    @deprecated("T.let", "T.Bind")
+    def let_stmt(v: Var, value: PrimExpr) -> Var:
+        return Bind(value, var=v)
 
     if body is None:
         return let_stmt(v, value)
@@ -2343,7 +2345,7 @@ def wrapped(*args, **kwargs):
     "Call",
     "CallEffectKind",
     "let",
-    "LetStmt",
+    "Bind",
     "Let",
     "IterVar",
     "CommReducer",

python/tvm/script/parser/core/parser.py

@@ -528,6 +528,9 @@ def _duplicate_lhs_check(self, target: doc.expr) -> bool | set[str]:
             return {target.id}
         elif isinstance(target, doc.Starred):
             return self._duplicate_lhs_check(target.value)
+        elif isinstance(target, doc.Attribute):
+            # Attribute assignment like packedB.data = ..., treated as rebinding.
+            return {target.attr}
         else:
             self.report_error(target, "Invalid type in assign statement")
             raise NotImplementedError

python/tvm/script/parser/tir/parser.py

@@ -145,11 +145,8 @@ def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -
         return value
     else:
         value = tvm.runtime.convert(value)
-        frame = T.LetStmt(value)
-        var = frame.var
+        var = T.Bind(value)
         IRBuilder.name(var_name, var)
-        frame.add_callback(partial(frame.__exit__, None, None, None))
-        frame.__enter__()
         return var
 
 
@@ -352,9 +349,7 @@ def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
     if not isinstance(ann_var, Var):
         self.report_error(node.annotation, "Annotation should be Var")
     self.eval_assign(target=lhs, source=ann_var, bind_value=bind_assign_value)
-    frame = T.LetStmt(rhs, var=ann_var)
-    frame.add_callback(partial(frame.__exit__, None, None, None))
-    frame.__enter__()
+    T.Bind(rhs, var=ann_var)
 
 
 @dispatch.register(token="tir", type_name="With")
@@ -471,6 +466,11 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
     elif isinstance(res, Frame):
         res.add_callback(partial(res.__exit__, None, None, None))
         res.__enter__()
+    elif isinstance(res, Var):
+        # Standalone Var expression (e.g. from T.Bind(value, var=v)) --
+        # the Bind statement was already emitted to the parent frame by the FFI call,
+        # so just discard the returned Var.
+        pass
     elif isinstance(res, PrimExpr):
         T.evaluate(res)
     elif isinstance(res, int | bool):

python/tvm/tir/__init__.py

@@ -29,7 +29,7 @@
 from .expr import Select, BufferLoad, ProducerLoad, Ramp, Broadcast, Shuffle
 from .expr import Call, CallEffectKind, Let, IterVar, CommReducer
 
-from .stmt import Stmt, LetStmt, AssertStmt, ForKind, For, While
+from .stmt import Stmt, Bind, AssertStmt, ForKind, For, While
 from .stmt import (
     BufferStore,
     AllocBuffer,