Fused SDDMM example.

include/tvm/tir/sparse.h

@@ -214,6 +214,11 @@ class FusedAxis;
 /*! \brief Derivation axis, constructed by T.fuse(axis1, axis2, ...) */
 class FusedAxisNode : public DenseFixedAxisNode {
  public:
+  /* The group of axes to be fused. */
+  Array<Axis> group;
+  /* The index of current FusedAxis in the group. */
+  int index;
+
   void VisitAttrs(AttrVisitor* v) {
     DenseFixedAxisNode::VisitAttrs(v);
     v->Visit("group", &group);
@@ -231,11 +236,6 @@ class FusedAxisNode : public DenseFixedAxisNode {
     hash_reduce(index);
   }
 
-  /* The group of axes to be fused. */
-  Array<Axis> group;
-  /* The index of current FusedAxis in the group. */
-  int index;
-
   static constexpr const char* _type_key = "tir.sparse.FusedAxis";
   TVM_DECLARE_FINAL_OBJECT_INFO(FusedAxisNode, DenseFixedAxisNode);
 };
@@ -313,8 +313,6 @@ class AttachedAxisNode : public DenseVariableAxisNode {
 
   PrimExpr Aggregate(SparseCtx* ctx, PrimExpr index) const;
 
-  PrimExpr Aggregate(SparseCtx* ctx, PrimExpr index) const;
-
   static constexpr const char* _type_key = "tir.sparse.AttachedAxis";
   TVM_DECLARE_FINAL_OBJECT_INFO(AttachedAxisNode, DenseVariableAxisNode);
 };

src/tir/transforms/lower_sparse_tir.cc

@@ -101,15 +101,12 @@ class SparseBlockCtx : public SparseCtx {
   struct Scope {
     explicit Scope(SparseBlock sp_block) : sp_block(std::move(sp_block)) {
       for (const SpIterVar& sp_iter_var : this->sp_block->sp_iter_vars) {
-        axis2sp_iter.Set(sp_iter_var->axis, sp_iter_var);
         sp_iter_var_map.Set(sp_iter_var->var, sp_iter_var);
       }
     }
 
     /*! \brief The sparse block */
     SparseBlock sp_block;
-    /*! \brief A mapping from axes to the sparse iterators that go over them */
-    Map<Axis, SpIterVar> axis2sp_iter;
     /*! \brief A mapping from the internal variables of sparse iterators to the iterators */
     Map<Var, SpIterVar> sp_iter_var_map;
     /*! \brief The stored offsets of the axis in the sparse block */
@@ -125,14 +122,40 @@ class SparseBlockCtx : public SparseCtx {
     stack_.emplace_back(GetRef<SparseBlock>(sp_block));
     /* Compute offsets and coordinates */
     size_t n_iters = sp_block->sp_iter_vars.size();
-    for (size_t i = 0; i < n_iters; ++i) {
+    for (size_t i = 0; i < n_iters;) {
       SpIterVar sp_iter_var = sp_block->sp_iter_vars[i];
       Axis axis = sp_iter_var->axis;
 
-      PrimExpr offset = AggregateOffset(this, axis, sp_iter_var->var, ana_);
-      SetOffset(axis, offset);
-      PrimExpr coordinate = axis->Decompress(this, offset, sp_iter_var->var);
-      SetCoordinate(axis, coordinate);
+      PrimExpr offset, index;
+      if (auto fused_axis = axis.as<FusedAxisNode>()) {
+        auto group = fused_axis->group;
+        offset = sp_block->sp_iter_vars[i + group.size() - 1]->var;
+        for (int j = group.size() - 1; j >= 0; --j) {
+          Axis orig = group[j];
+          SetOffset(orig, offset);
+          if (j > 0) {
+            // TODO(zihao): support more than sv axis.
+            offset = lower_bound(Downcast<SparseVariableAxis>(orig)->indptr->data, offset,
+                                 Integer(0), orig->GetNNZ());
+          }
+        }
+        for (size_t j = 0; j < group.size(); ++j) {
+          Axis orig = group[j];
+          offset = GetOffset(orig);
+          PrimExpr lb = std::get<0>(orig->GetOffsetExtent(this));
+          index = offset - lb;
+          PrimExpr coordinate = orig->Decompress(this, offset, index);
+          SetCoordinate(orig, coordinate);
+          i++;
+        }
+      } else {
+        offset = AggregateOffset(this, axis, sp_iter_var->var, ana_);
+        index = sp_iter_var->var;
+        PrimExpr coordinate = axis->Decompress(this, offset, index);
+        SetOffset(axis, offset);
+        SetCoordinate(axis, coordinate);
+        i++;
+      }
     }
   }
 
@@ -164,7 +187,7 @@ class SparseBlockCtx : public SparseCtx {
    */
   PrimExpr GetOffset(Axis axis) const {
     Optional<PrimExpr> try_offset = top()->cached_offsets.Get(axis);
-    CHECK(try_offset.defined()) << "The offset of axis not defined yet.";
+    CHECK(try_offset.defined()) << "The offset of axis " << axis->name << " not defined yet.";
     PrimExpr offset = try_offset.value();
     return std::move(offset);
   }
@@ -202,7 +225,7 @@ class SparseBlockCtx : public SparseCtx {
   }
 
   Optional<Axis> MatchAxis(SparseCtx* buf_ctx, Axis axis) {
-    if (!top()->axis2sp_iter.Get(axis).defined()) {
+    if (!top()->cached_offsets.Get(axis).defined()) {
       return NullOpt;
     } else {
       Axis axis_ = axis;
@@ -233,7 +256,11 @@ class SparseBlockCtx : public SparseCtx {
     if (!try_sp_iter_var.defined()) {
       return false;
     }
-    return try_sp_iter_var.value()->axis == matched_axis.value();
+    Axis axis = try_sp_iter_var.value()->axis;
+    if (auto fused_axis = axis.as<FusedAxisNode>()) {
+      axis = fused_axis->group[fused_axis->index];
+    }
+    return axis == matched_axis.value();
   }
 
  private:
@@ -403,7 +430,11 @@ class IndexTransformer : public StmtExprMutator {
     auto try_sp_iter = sp_blk_ctx_.GetSparseIterVar(var);
     if (try_sp_iter.defined()) {
       SpIterVar sp_iter = try_sp_iter.value();
-      return sp_blk_ctx_.GetCoordinate(sp_iter->axis);
+      Axis axis = sp_iter->axis;
+      if (auto fused_axis = axis.as<FusedAxisNode>()) {
+        axis = fused_axis->group[fused_axis->index];
+      }
+      return sp_blk_ctx_.GetCoordinate(axis);
     } else {
       return GetRef<PrimExpr>(var);
     }

tests/python/sparsetir/test_tir_sparse_lower.py

@@ -662,6 +662,7 @@ def test_sddmm():
 def test_fused_sddmm():
     mod = tvm.IRModule.from_expr(fused_sddmm)
     mod = tvm.tir.transform.LowerSparseTIR()(mod)
+    print(mod["main"].script())
     # TODO
 
 
@@ -754,7 +755,7 @@ def test_square_sum_two_K():
     test_ellpack_mm()
     test_csr_element_wise()
     test_sddmm()
-    # test_fused_sddmm()
+    test_fused_sddmm()
     test_bmm()
     test_square_sum()
     test_square_sum_two_K()