[SparseTIR] Index Lowering (#8)

* Add StmtFunctor/ExprFunctor for SparseBufferStore/Load * Add basic index lowering * Finish index lowering (maybe) * Address comments * Convert CRLF to LF
yzh119 · Feb 15, 2022 · affb1af · affb1af
1 parent 65ce747
commit affb1af
Show file tree

Hide file tree

Showing 8 changed files with 379 additions and 2 deletions.
diff --git a/include/tvm/tir/expr_functor.h b/include/tvm/tir/expr_functor.h
@@ -119,6 +119,7 @@ class ExprFunctor<R(const PrimExpr& n, Args...)> {
     return VisitExpr_(static_cast<const VarNode*>(op), std::forward<Args>(args)...);
   }
   virtual R VisitExpr_(const BufferLoadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+  virtual R VisitExpr_(const SparseBufferLoadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const ProducerLoadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const LoadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const LetNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
@@ -165,6 +166,7 @@ class ExprFunctor<R(const PrimExpr& n, Args...)> {
     IR_EXPR_FUNCTOR_DISPATCH(SizeVarNode);
     IR_EXPR_FUNCTOR_DISPATCH(LoadNode);
     IR_EXPR_FUNCTOR_DISPATCH(BufferLoadNode);
+    IR_EXPR_FUNCTOR_DISPATCH(SparseBufferLoadNode);
     IR_EXPR_FUNCTOR_DISPATCH(ProducerLoadNode);
     IR_EXPR_FUNCTOR_DISPATCH(LetNode);
     IR_EXPR_FUNCTOR_DISPATCH(CallNode);
@@ -217,6 +219,7 @@ class TVM_DLL ExprVisitor : public ExprFunctor<void(const PrimExpr&)> {
   void VisitExpr_(const SizeVarNode* op) override;
   void VisitExpr_(const LoadNode* op) override;
   void VisitExpr_(const BufferLoadNode* op) override;
+  void VisitExpr_(const SparseBufferLoadNode* op) override;
   void VisitExpr_(const ProducerLoadNode* op) override;
   void VisitExpr_(const LetNode* op) override;
   void VisitExpr_(const CallNode* op) override;
@@ -264,6 +267,7 @@ class TVM_DLL ExprMutator : protected ExprFunctor<PrimExpr(const PrimExpr&)> {
   PrimExpr VisitExpr_(const SizeVarNode* op) override;
   PrimExpr VisitExpr_(const LoadNode* op) override;
   PrimExpr VisitExpr_(const BufferLoadNode* op) override;
+  PrimExpr VisitExpr_(const SparseBufferLoadNode* op) override;
   PrimExpr VisitExpr_(const ProducerLoadNode* op) override;
   PrimExpr VisitExpr_(const LetNode* op) override;
   PrimExpr VisitExpr_(const CallNode* op) override;

diff --git a/include/tvm/tir/stmt_functor.h b/include/tvm/tir/stmt_functor.h
@@ -89,6 +89,7 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
   virtual R VisitStmt_(const AllocateNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const StoreNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const BufferStoreNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
+  virtual R VisitStmt_(const SparseBufferStoreNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const BufferRealizeNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const AssertStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const ProducerStoreNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
@@ -121,6 +122,7 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
     IR_STMT_FUNCTOR_DISPATCH(SeqStmtNode);
     IR_STMT_FUNCTOR_DISPATCH(EvaluateNode);
     IR_STMT_FUNCTOR_DISPATCH(BufferStoreNode);
+    IR_STMT_FUNCTOR_DISPATCH(SparseBufferStoreNode);
     IR_STMT_FUNCTOR_DISPATCH(BufferRealizeNode);
     IR_STMT_FUNCTOR_DISPATCH(BlockNode);
     IR_STMT_FUNCTOR_DISPATCH(BlockRealizeNode);
@@ -157,6 +159,7 @@ class TVM_DLL StmtVisitor : protected StmtFunctor<void(const Stmt&)> {
   void VisitStmt_(const AllocateNode* op) override;
   void VisitStmt_(const StoreNode* op) override;
   void VisitStmt_(const BufferStoreNode* op) override;
+  void VisitStmt_(const SparseBufferStoreNode* op) override;
   void VisitStmt_(const BufferRealizeNode* op) override;
   void VisitStmt_(const AssertStmtNode* op) override;
   void VisitStmt_(const ProducerStoreNode* op) override;
@@ -257,6 +260,7 @@ class TVM_DLL StmtMutator : protected StmtFunctor<Stmt(const Stmt&)> {
   Stmt VisitStmt_(const AllocateNode* op) override;
   Stmt VisitStmt_(const StoreNode* op) override;
   Stmt VisitStmt_(const BufferStoreNode* op) override;
+  Stmt VisitStmt_(const SparseBufferStoreNode* op) override;
   Stmt VisitStmt_(const BufferRealizeNode* op) override;
   Stmt VisitStmt_(const AssertStmtNode* op) override;
   Stmt VisitStmt_(const ProducerStoreNode* op) override;

diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h
@@ -500,6 +500,12 @@ TVM_DLL Pass ConvertForLoopsToSerial();
  */
 TVM_DLL Pass UnifiedStaticMemoryPlanner();
 
+/*!
+ * \brief Lower SparseTIR to TIR.
+ * \return The pass.
+ */
+TVM_DLL Pass LowerSparseTIR();
+
 }  // namespace transform
 }  // namespace tir
 }  // namespace tvm

diff --git a/python/tvm/tir/transform/transform.py b/python/tvm/tir/transform/transform.py
@@ -760,3 +760,14 @@ def ConvertForLoopsToSerial():
         The result pass
     """
     return _ffi_api.ConvertForLoopsToSerial()  # type: ignore
+
+
+def LowerSparseTIR():
+    """Lower SparseTIR to TIR
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.LowerSparseTIR()  # type: ignore
diff --git a/src/tir/ir/expr_functor.cc b/src/tir/ir/expr_functor.cc
@@ -43,6 +43,10 @@ void ExprVisitor::VisitExpr_(const BufferLoadNode* op) {
   VisitArray(op->indices, [this](const PrimExpr& e) { this->VisitExpr(e); });
 }
 
+void ExprVisitor::VisitExpr_(const SparseBufferLoadNode* op) {
+  VisitArray(op->indices, [this](const PrimExpr& e) { this->VisitExpr(e); });
+}
+
 void ExprVisitor::VisitExpr_(const ProducerLoadNode* op) {
   VisitArray(op->indices, [this](const PrimExpr& e) { this->VisitExpr(e); });
 }
@@ -146,6 +150,16 @@ PrimExpr ExprMutator::VisitExpr_(const BufferLoadNode* op) {
   }
 }
 
+PrimExpr ExprMutator::VisitExpr_(const SparseBufferLoadNode* op) {
+  auto fmutate = [this](const PrimExpr& e) { return this->VisitExpr(e); };
+  Array<PrimExpr> indices = MutateArray(op->indices, fmutate);
+  if (indices.same_as(op->indices)) {
+    return GetRef<PrimExpr>(op);
+  } else {
+    return SparseBufferLoad(op->buffer, indices);
+  }
+};
+
 PrimExpr ExprMutator::VisitExpr_(const ProducerLoadNode* op) {
   auto fmutate = [this](const PrimExpr& e) { return this->VisitExpr(e); };
   Array<PrimExpr> indices = MutateArray(op->indices, fmutate);

diff --git a/src/tir/ir/sparse.cc b/src/tir/ir/sparse.cc
@@ -21,6 +21,7 @@
  * \file sparse.cc
  * \brief buffers and formats in sparse tir.
  */
+#include <tvm/arith/analyzer.h>
 #include <tvm/runtime/registry.h>
 #include <tvm/tir/buffer.h>
 #include <tvm/tir/sparse.h>
@@ -158,9 +159,21 @@ SpIterVar::SpIterVar(String name, PrimExpr max_extent, SpIterKind kind, bool is_
                      Optional<Axis> axis) {
   ObjectPtr<SpIterVarNode> node = make_object<SpIterVarNode>();
 
+  arith::Analyzer ana;
+  if (axis.defined()) {
+    CHECK(ana.CanProveEqual(axis.value()->length, max_extent));
+  }
   if (kind != SpIterKind::kDenseFixed) {
     CHECK(axis.defined()) << "ValueError: To create a SpIterVar that is not fixed-dense, one must "
                              "specify the axis over which the SpIterVar iterates";
+    const char* err_str = "ValueError: The given kind doesn't match the type of the given axis";
+    if (kind == SpIterKind::kDenseVariable) {
+      CHECK(axis.value()->IsInstance<DenseFixedAxisNode>()) << err_str;
+    } else if (kind == SpIterKind::kSparseFixed) {
+      CHECK(axis.value()->IsInstance<SparseFixedAxisNode>()) << err_str;
+    } else if (kind == SpIterKind::kSparseVariable) {
+      CHECK(axis.value()->IsInstance<SparseVariableAxisNode>()) << err_str;
+    }
   }
 
   node->var = Var(std::move(name));
@@ -174,9 +187,9 @@ SpIterVar::SpIterVar(String name, PrimExpr max_extent, SpIterKind kind, bool is_
 TVM_REGISTER_NODE_TYPE(SpIterVarNode);
 
 TVM_REGISTER_GLOBAL("tir.sparse.SpIterVar")
-    .set_body_typed([](String name, PrimExpr max_extent, SpIterKind kind, bool is_reduction,
+    .set_body_typed([](String name, PrimExpr max_extent, int kind, bool is_reduction,
                        Optional<Axis> axis) {
-      return SpIterVar(name, max_extent, kind, is_reduction, axis);
+      return SpIterVar(name, max_extent, SpIterKind(kind), is_reduction, axis);
     });
 
 }  // namespace tir

diff --git a/src/tir/ir/stmt_functor.cc b/src/tir/ir/stmt_functor.cc
@@ -69,6 +69,11 @@ void StmtVisitor::VisitStmt_(const BufferStoreNode* op) {
   VisitArray(op->indices, [this](const PrimExpr& e) { this->VisitExpr(e); });
 }
 
+void StmtVisitor::VisitStmt_(const SparseBufferStoreNode* op) {
+  this->VisitExpr(op->value);
+  VisitArray(op->indices, [this](const PrimExpr& e) { this->VisitExpr(e); });
+}
+
 void StmtVisitor::VisitStmt_(const BufferRealizeNode* op) {
   VisitArray(op->bounds, [this](const Range& r) {
     this->VisitExpr(r->min);
@@ -367,6 +372,20 @@ Stmt StmtMutator::VisitStmt_(const BufferStoreNode* op) {
   }
 }
 
+Stmt StmtMutator::VisitStmt_(const SparseBufferStoreNode* op) {
+  PrimExpr value = this->VisitExpr(op->value);
+  Array<PrimExpr> indices = Internal::Mutate(this, op->indices);
+
+  if (value.same_as(op->value) && indices.same_as(op->indices)) {
+    return GetRef<Stmt>(op);
+  } else {
+    auto n = CopyOnWrite(op);
+    n->value = std::move(value);
+    n->indices = std::move(indices);
+    return Stmt(n);
+  }
+}
+
 Stmt StmtMutator::VisitStmt_(const BufferRealizeNode* op) {
   Region bounds = Internal::Mutate(this, op->bounds);
   PrimExpr condition = this->VisitExpr(op->condition);