[SparseTIR] Parser, Printer, Roundtrip (#14)

* SparseBlock scope handler (part 1)

* SparseBlock scope handler (part 2)

* SparseBlock scope handler (part 3)

* SparseBlock scope handler (fix 1)

* Add SparseBufferLoad/Store on Python side

* Parser for SparseBufferLoad/Store

* Add SparseBlock to Python __init__

* StmtFunctor for SparseBlock

* Ensure at least one dimension for SparseBuffer

* Make `axis` field of SpIterVar mandatory

* SparseBlock scope handler (fix 2)

* Update Axis syntax by removing `name` parameter

* Move to intrin.py

* Add filed `from_sparse` to DenseFixedAxis

* SparseTIR script printer

* Roundtrip test

* `update_symbol` bug fix

* Fix attr visit in SparseBuffer

* Define then compare in SparseBlock

* Fix printer bug for SparseBuffer

* Enable graph match for Axis and SparseBuffer

* Complete HashReduce and EqualReduce for AxisTree and SparseBuffer

* Fix typo

* Rename test

* Bug fix 1

* Bug fix 2

* Add more tests

Author: MasterJH5574

include/tvm/tir/expr.h

@@ -656,7 +656,7 @@ class BufferLoad : public PrimExpr {
  */
 class SparseBufferLoadNode : public PrimExprNode {
  public:
-  /*! \brief The buffer variable. */
+  /*! \brief The buffer to be loaded. */
   SparseBuffer buffer;
   /*! \brief The indices location to be loaded. */
   Array<PrimExpr> indices;

include/tvm/tir/sparse.h

@@ -99,23 +99,48 @@ class DenseAxis : public Axis {
   TVM_DEFINE_OBJECT_REF_METHODS(DenseAxis, Axis, DenseAxisNode);
 };
 
+/*!
+ * \brief Sparse axis whose column indices is not consecutive.
+ */
+class SparseAxisNode : public AxisNode {
+ public:
+  static constexpr const char* _type_key = "tir.sparse.SparseAxis";
+  TVM_DECLARE_BASE_OBJECT_INFO(SparseAxisNode, AxisNode);
+};
+
+/*!
+ * \brief Managed reference to SparseAxisNode.
+ * \sa SparseAxisNode
+ */
+class SparseAxis : public Axis {
+ public:
+  TVM_DEFINE_OBJECT_REF_METHODS(SparseAxis, Axis, SparseAxisNode);
+};
+
 /*!
  * \brief Dense axis with fixed length per row.
  */
 class DenseFixedAxisNode : public DenseAxisNode {
  public:
+  Optional<SparseAxis> from_sparse;
+
   void VisitAttrs(AttrVisitor* v) {
     v->Visit("name", &name);
     v->Visit("length", &length);
+    v->Visit("from_sparse", &from_sparse);
   }
 
-  bool SEqualReduce(const DenseAxisNode* other, SEqualReducer equal) const {
-    return equal(name, other->name) && equal(length, other->length);
+  bool SEqualReduce(const DenseFixedAxisNode* other, SEqualReducer equal) const {
+    equal->MarkGraphNode();
+    return equal(name, other->name) && equal(length, other->length) &&
+           equal(from_sparse, other->from_sparse);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce->MarkGraphNode();
     hash_reduce(name);
     hash_reduce(length);
+    hash_reduce(from_sparse);
   }
 
   static constexpr const char* _type_key = "tir.sparse.DenseFixedAxis";
@@ -128,7 +153,8 @@ class DenseFixedAxisNode : public DenseAxisNode {
  */
 class DenseFixedAxis : public DenseAxis {
  public:
-  TVM_DLL explicit DenseFixedAxis(String name, PrimExpr length);
+  TVM_DLL explicit DenseFixedAxis(String name, PrimExpr length,
+                                  Optional<SparseAxis> from_sparse = NullOpt);
 
   TVM_DEFINE_OBJECT_REF_METHODS(DenseFixedAxis, DenseAxis, DenseFixedAxisNode);
 };
@@ -144,10 +170,12 @@ class DenseVariableAxisNode : public DenseAxisNode {
   }
 
   bool SEqualReduce(const DenseVariableAxisNode* other, SEqualReducer equal) const {
+    equal->MarkGraphNode();
     return equal(name, other->name) && equal(length, other->length) && equal(indptr, other->indptr);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce->MarkGraphNode();
     hash_reduce(name);
     hash_reduce(length);
     hash_reduce(indptr);
@@ -168,24 +196,6 @@ class DenseVariableAxis : public DenseAxis {
   TVM_DEFINE_OBJECT_REF_METHODS(DenseVariableAxis, DenseAxis, DenseVariableAxisNode);
 };
 
-/*!
- * \brief Sparse axis whose column indices is not consecutive.
- */
-class SparseAxisNode : public AxisNode {
- public:
-  static constexpr const char* _type_key = "tir.sparse.SparseAxis";
-  TVM_DECLARE_BASE_OBJECT_INFO(SparseAxisNode, AxisNode);
-};
-
-/*!
- * \brief Managed reference to SparseAxisNode.
- * \sa SparseAxisNode
- */
-class SparseAxis : public Axis {
- public:
-  TVM_DEFINE_OBJECT_REF_METHODS(SparseAxis, Axis, SparseAxisNode);
-};
-
 /*!
  * \brief Sparse axis with fixed number of non-zero columns per row.
  */
@@ -203,11 +213,13 @@ class SparseFixedAxisNode : public SparseAxisNode {
   }
 
   bool SEqualReduce(const SparseFixedAxisNode* other, SEqualReducer equal) const {
+    equal->MarkGraphNode();
     return equal(name, other->name) && equal(length, other->length) &&
            equal(indices, other->indices) && equal(num_cols, other->num_cols);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce->MarkGraphNode();
     hash_reduce(name);
     hash_reduce(length);
     hash_reduce(indices);
@@ -245,11 +257,13 @@ class SparseVariableAxisNode : public SparseAxisNode {
   }
 
   bool SEqualReduce(const SparseVariableAxisNode* other, SEqualReducer equal) const {
+    equal->MarkGraphNode();
     return equal(name, other->name) && equal(length, other->length) &&
            equal(indptr, other->indptr) && equal(indices, other->indices);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce->MarkGraphNode();
     hash_reduce(name);
     hash_reduce(length);
     hash_reduce(indptr);
@@ -277,13 +291,27 @@ class SparseVariableAxis : public SparseAxis {
 class AxisTreeNode : public Object {
  public:
   // unordered map that stores the parent relationship between axes.
-  std::unordered_map<String, Optional<String>, ObjectPtrHash, ObjectPtrEqual> parent;
+  Map<String, Optional<String>> parent;
   // unordered map that stores the children relationship between axes.
-  std::unordered_map<Optional<String>, Array<String>, ObjectPtrHash, ObjectPtrEqual> children;
+  Map<Optional<String>, Array<String>> children;
+
+  void VisitAttrs(AttrVisitor* v) {
+    v->Visit("parent", &parent);
+    v->Visit("children", &children);
+  }
+
+  bool SEqualReduce(const AxisTreeNode* other, SEqualReducer equal) const {
+    return equal(parent, other->parent) && equal(children, other->children);
+  }
 
-  void VisitAttrs(AttrVisitor* v) {}
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce(parent);
+    hash_reduce(children);
+  }
 
   static constexpr const char* _type_key = "tir.sparse.AxisTree";
+  static constexpr const bool _type_has_method_sequal_reduce = true;
+  static constexpr const bool _type_has_method_shash_reduce = true;
   TVM_DECLARE_FINAL_OBJECT_INFO(AxisTreeNode, Object);
 };
 
@@ -313,22 +341,26 @@ class SparseBufferNode : public Object {
   inline int ndim() const { return static_cast<int>(axes.size()); }
 
   void VisitAttrs(AttrVisitor* v) {
-    v->Visit("length", &axes);
-    v->Visit("num_cols", &data);
+    v->Visit("axes", &axes);
+    v->Visit("data", &data);
     v->Visit("name", &name);
   }
 
   bool SEqualReduce(const SparseBufferNode* other, SEqualReducer equal) const {
+    equal->MarkGraphNode();
     return equal(axes, other->axes) && equal(data, other->data) && equal(name, other->name);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce->MarkGraphNode();
     hash_reduce(axes);
     hash_reduce(data);
     hash_reduce(name);
   }
 
   static constexpr const char* _type_key = "tir.sparse.SparseBuffer";
+  static constexpr const bool _type_has_method_sequal_reduce = true;
+  static constexpr const bool _type_has_method_shash_reduce = true;
   TVM_DECLARE_FINAL_OBJECT_INFO(SparseBufferNode, Object);
 };
 
@@ -359,7 +391,7 @@ class SpIterVarNode : public Object {
   PrimExpr max_extent;
   SpIterKind kind;
   bool is_reduction;
-  Optional<Axis> axis;
+  Axis axis;
 
   void VisitAttrs(AttrVisitor* v) {
     v->Visit("var", &var);
@@ -392,7 +424,7 @@ class SpIterVarNode : public Object {
 class SpIterVar : public ObjectRef {
  public:
   TVM_DLL explicit SpIterVar(Var var, PrimExpr max_extent, SpIterKind kind, bool is_reduction,
-                             Optional<Axis> axis = NullOpt);
+                             Axis axis);
 
   /*!
    * \return the corresponding var in the IterVar.

include/tvm/tir/stmt.h

@@ -335,11 +335,11 @@ class BufferStore : public Stmt {
  *  buffer[i, j] = value;
  *
  * \endcode
- * \sa SparseBufferLoad
+ * \sa SparseBufferStore
  */
 class SparseBufferStoreNode : public StmtNode {
  public:
-  /*! \brief The buffer variable. */
+  /*! \brief The sparse buffer to be accessed. */
   SparseBuffer buffer;
   /*! \brief The value to be stored. */
   PrimExpr value;
@@ -1303,17 +1303,17 @@ class SparseBlockNode : public StmtNode {
   }
 
   bool SEqualReduce(const SparseBlockNode* other, SEqualReducer equal) const {
-    return equal(sp_iter_vars, other->sp_iter_vars) &&
-           equal(sp_struct2param_map, other->sp_struct2param_map) && equal(name, other->name) &&
-           equal(body, other->body) && equal(init, other->init);
+    return equal(sp_iter_vars, other->sp_iter_vars) && equal(name, other->name) &&
+           equal(body, other->body) && equal(init, other->init) &&
+           equal(sp_struct2param_map, other->sp_struct2param_map);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
     hash_reduce(sp_iter_vars);
-    hash_reduce(sp_struct2param_map);
     hash_reduce(name);
     hash_reduce(body);
     hash_reduce(init);
+    hash_reduce(sp_struct2param_map);
   }
 
   static constexpr const char* _type_key = "tir.SparseBlock";

include/tvm/tir/stmt_functor.h

@@ -99,6 +99,7 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
   virtual R VisitStmt_(const EvaluateNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const BlockNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const BlockRealizeNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
+  virtual R VisitStmt_(const SparseBlockNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmtDefault_(const Object* op, Args...) {
     LOG(FATAL) << "Do not have a default for " << op->GetTypeKey();
     return R();
@@ -126,6 +127,7 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
     IR_STMT_FUNCTOR_DISPATCH(BufferRealizeNode);
     IR_STMT_FUNCTOR_DISPATCH(BlockNode);
     IR_STMT_FUNCTOR_DISPATCH(BlockRealizeNode);
+    IR_STMT_FUNCTOR_DISPATCH(SparseBlockNode);
     return vtable;
   }
 };
@@ -169,6 +171,7 @@ class TVM_DLL StmtVisitor : protected StmtFunctor<void(const Stmt&)> {
   void VisitStmt_(const EvaluateNode* op) override;
   void VisitStmt_(const BlockNode* op) override;
   void VisitStmt_(const BlockRealizeNode* op) override;
+  void VisitStmt_(const SparseBlockNode* op) override;
 };
 
 /*!
@@ -270,6 +273,7 @@ class TVM_DLL StmtMutator : protected StmtFunctor<Stmt(const Stmt&)> {
   Stmt VisitStmt_(const EvaluateNode* op) override;
   Stmt VisitStmt_(const BlockNode* op) override;
   Stmt VisitStmt_(const BlockRealizeNode* op) override;
+  Stmt VisitStmt_(const SparseBlockNode* op) override;
   /*!
    * \brief Alternative advance method for SeqStmtNode.
    *

python/tvm/script/context_maintainer.py

@@ -219,7 +219,7 @@ def update_symbol(
         self, name: str, symbol: Union[Buffer, Var, SparseBuffer, Axis], node: synr.ast.Node
     ):
         """Append a symbol into current scope"""
-        if isinstance(symbol, (Buffer, Var, SparseBuffer, Axis)):
+        if isinstance(symbol, (Buffer, SparseBuffer, Axis)):
             if name in self.symbols[0]:
                 self.report_error("Duplicate Buffer name: " + symbol.name, node.span)
             self.symbols[0][name] = symbol

python/tvm/script/parser.py

@@ -582,6 +582,14 @@ def transform_SubscriptAssign(self, node):
                 indexes,
                 span=tvm_span_from_synr(node.span),
             )
+        elif isinstance(symbol, tvm.tir.sparse.SparseBuffer):
+            # SparseBufferStore
+            return tvm.tir.SparseBufferStore(
+                symbol,
+                tvm.runtime.convert(rhs, span=rhs_span),
+                indexes,
+                span=tvm_span_from_synr(node.span),
+            )
         else:
             if len(indexes) != 1:
                 self.report_error(
@@ -876,6 +884,8 @@ def transform_Subscript(self, node):
             return BufferSlice(
                 symbol, indexes, self.report_error, span=tvm_span_from_synr(node.span)
             )
+        elif isinstance(symbol, tvm.tir.sparse.SparseBuffer):
+            return tvm.tir.SparseBufferLoad(symbol, indexes, span=tvm_span_from_synr(node.span))
         elif isinstance(symbol, tvm.container.Array):
             if len(indexes) > 1:
                 self.report_error(

python/tvm/script/tir/intrin.py

@@ -17,9 +17,18 @@
 """TVM Script Parser Intrinsic Classes"""
 # pylint: disable=redefined-builtin, relative-beyond-top-level
 import builtins
-from typing import List, Any
+from typing import List, Optional, Any
 
 import tvm.tir
+from tvm.ir import Span
+from tvm.tir.sparse import (
+    Axis,
+    DenseFixedAxis,
+    DenseVariableAxis,
+    SpIterVar,
+    SparseFixedAxis,
+    SparseVariableAxis,
+)
 from ..registry import register
 from ..utils import get_param_list, tvm_span_from_synr
 
@@ -244,3 +253,35 @@ def comm_reducer(lambda_io, identities, span):
         lambda_output = (lambda_output,)
 
     return tvm.tir.CommReducer(x, y, lambda_output, identities, span)
+
+
+@register
+def to_dense(axis: Axis, span: Optional[Span] = None):
+    if isinstance(axis, (SparseFixedAxis, SparseVariableAxis)):
+        return DenseFixedAxis(axis.name + "_dense", axis.length, axis)
+    else:
+        return axis
+
+
+@register
+def cord(axis: Axis, span: Optional[Span] = None):
+    # The field `var` and `is_reduction` will be updated in SparseBlock scope handler
+    var_temp = tvm.te.var()
+    if isinstance(axis, DenseVariableAxis):
+        return SpIterVar(var_temp, axis.length, SpIterVar.DenseVariable, False, axis)
+    else:
+        return SpIterVar(var_temp, axis.length, SpIterVar.DenseFixed, False, axis)
+
+
+@register
+def pos(axis: Axis, span: Optional[Span] = None):
+    # The field `var` and `is_reduction` will be updated in SparseBlock scope handler
+    var_temp = tvm.te.var()
+    if isinstance(axis, DenseFixedAxis):
+        return SpIterVar(var_temp, axis.length, SpIterVar.DenseFixed, False, axis)
+    elif isinstance(axis, DenseVariableAxis):
+        return SpIterVar(var_temp, axis.length, SpIterVar.DenseVariable, False, axis)
+    elif isinstance(axis, SparseFixedAxis):
+        return SpIterVar(var_temp, axis.length, SpIterVar.SparseFixed, False, axis)
+    else:
+        return SpIterVar(var_temp, axis.length, SpIterVar.SparseVariable, False, axis)