[TIRx] Restore BufferRegion as typed expression

include/tvm/ir/base_expr.h

@@ -435,21 +435,24 @@ class PrimExpr : public TypedExpr<PrimType> {
  * This is useful for the FFI to convert the expressions to PrimExpr.
  * \sa PrimExpr
  */
-class PrimExprConvertibleNode : public ffi::Object {
+class PrimExprConvertibleNode : public ExprNode {
  public:
   virtual ~PrimExprConvertibleNode() {}
   virtual PrimExpr ToPrimExpr() const = 0;
-  TVM_FFI_DECLARE_OBJECT_INFO("ir.PrimExprConvertible", PrimExprConvertibleNode, ffi::Object);
+  static constexpr const uint32_t _type_child_slots = 2;
+  TVM_FFI_DECLARE_OBJECT_INFO("ir.PrimExprConvertible", PrimExprConvertibleNode, ExprNode);
 };
 
 /*!
  * \brief Managed reference to PrimExprConvertibleNode.
  * \sa PrimExprConvertibleNode
  */
-class PrimExprConvertible : public ffi::ObjectRef {
+class PrimExprConvertible : public Expr {
  public:
-  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimExprConvertible, ffi::ObjectRef,
-                                             PrimExprConvertibleNode);
+  bool operator==(const PrimExprConvertible& other) const { return this->same_as(other); }
+  bool operator!=(const PrimExprConvertible& other) const { return !(*this == other); }
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimExprConvertible, Expr, PrimExprConvertibleNode);
 };
 
 namespace ffi {

include/tvm/tirx/expr_functor.h

@@ -27,6 +27,7 @@
 
 #include <tvm/ir/node_functor.h>
 #include <tvm/tirx/expr.h>
+#include <tvm/tirx/stmt.h>
 
 #include <utility>
 
@@ -117,6 +118,7 @@ class ExprFunctor<R(const Expr& n, Args...)> {
   virtual R VisitExpr_(const VarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const BufferLoadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const OpaqueExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+  virtual R VisitExpr_(const BufferRegionNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const TupleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const TupleGetItemNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const LetNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
@@ -161,6 +163,7 @@ class ExprFunctor<R(const Expr& n, Args...)> {
     IR_EXPR_FUNCTOR_DISPATCH(VarNode);
     IR_EXPR_FUNCTOR_DISPATCH(BufferLoadNode);
     IR_EXPR_FUNCTOR_DISPATCH(OpaqueExprNode);
+    IR_EXPR_FUNCTOR_DISPATCH(BufferRegionNode);
     IR_EXPR_FUNCTOR_DISPATCH(TupleNode);
     IR_EXPR_FUNCTOR_DISPATCH(TupleGetItemNode);
     IR_EXPR_FUNCTOR_DISPATCH(LetNode);
@@ -213,6 +216,7 @@ class TVM_DLL ExprVisitor : public ExprFunctor<void(const Expr&)> {
   void VisitExpr_(const VarNode* op) override;
   void VisitExpr_(const BufferLoadNode* op) override;
   void VisitExpr_(const OpaqueExprNode* op) override;
+  void VisitExpr_(const BufferRegionNode* op) override;
   void VisitExpr_(const TupleNode* op) override;
   void VisitExpr_(const TupleGetItemNode* op) override;
   void VisitExpr_(const LetNode* op) override;
@@ -261,6 +265,7 @@ class TVM_DLL ExprMutator : protected ExprFunctor<Expr(const Expr&)> {
   Expr VisitExpr_(const VarNode* op) override;
   Expr VisitExpr_(const BufferLoadNode* op) override;
   Expr VisitExpr_(const OpaqueExprNode* op) override;
+  Expr VisitExpr_(const BufferRegionNode* op) override;
   Expr VisitExpr_(const TupleNode* op) override;
   Expr VisitExpr_(const TupleGetItemNode* op) override;
   Expr VisitExpr_(const LetNode* op) override;

include/tvm/tirx/stmt.h

@@ -770,6 +770,29 @@ class Continue : public Stmt {
   TVM_DEFINE_OBJECT_REF_COW_METHOD(ContinueNode);
 };
 
+/*!
+ * \brief The type of a multi-dimensional buffer region expression.
+ */
+class BufferRegionTypeNode : public TypeNode {
+ public:
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<BufferRegionTypeNode>();
+  }
+
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegionType", BufferRegionTypeNode, TypeNode);
+};
+
+/*!
+ * \brief Managed reference to BufferRegionTypeNode.
+ */
+class BufferRegionType : public Type {
+ public:
+  TVM_DLL BufferRegionType(Span span = Span());
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BufferRegionType, Type, BufferRegionTypeNode);
+};
+
 /*!
  * \brief Representing the region of multi-dimensional buffer access.
  */

include/tvm/tirx/stmt_functor.h

@@ -343,6 +343,7 @@ class TVM_DLL StmtExprVisitor : public ExprVisitor, public StmtVisitor {
 
   void VisitExpr(const Expr& e) override { return ExprVisitor::VisitExpr(e); }
   void VisitExpr_(const BufferLoadNode* op) override;
+  void VisitExpr_(const BufferRegionNode* op) override;
 };
 
 /*!
@@ -362,6 +363,7 @@ class TVM_DLL StmtExprMutator : public ExprMutator, public StmtMutator {
   Expr VisitExpr(const Expr& e) override { return ExprMutator::VisitExpr(e); }
   Expr VisitExpr_(const VarNode* op) override;
   Expr VisitExpr_(const BufferLoadNode* op) override;
+  Expr VisitExpr_(const BufferRegionNode* op) override;
 };
 
 /*!

include/tvm/tirx/var.h

@@ -162,12 +162,6 @@ class IterVarNode : public PrimExprConvertibleNode {
    *  set this if this is bound already to a known thread tag.
    */
   ffi::String thread_tag;
-  /*!
-   * \brief Span that points to the original source code.
-   *        Reserved debug information.
-   */
-  mutable Span span;
-
   PrimExpr ToPrimExpr() const final { return var; }
 
   static void RegisterReflection() {
@@ -176,9 +170,7 @@ class IterVarNode : public PrimExprConvertibleNode {
         .def_ro("dom", &IterVarNode::dom)
         .def_ro("var", &IterVarNode::var, refl::AttachFieldFlag::SEqHashDefRecursive())
         .def_ro("iter_type", &IterVarNode::iter_type)
-        .def_ro("thread_tag", &IterVarNode::thread_tag)
-        .def_ro("span", &IterVarNode::span, refl::DefaultValue(Span()),
-                refl::AttachFieldFlag::SEqHashIgnore());
+        .def_ro("thread_tag", &IterVarNode::thread_tag);
   }
 
   static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;

python/tvm/ir/__init__.py

@@ -39,6 +39,7 @@
     Expr,
     GlobalVar,
     OpaqueExpr,
+    PrimExprConvertible,
     Range,
     Tuple,
     TupleGetItem,