[SparseTIR] Hack IsAffineBinding check (#27)

* [TensorIR][Schedule] Inherit block anotation upon creating new blocks

* Fix SDDMM test

* Hack IsAffineBinding for sparse blocks

Author: MasterJH5574

include/tvm/tir/schedule/state.h

@@ -162,6 +162,13 @@ class ScheduleStateNode : public Object {
    * \return A boolean flag indicating if the block has quasi-affine bindings
    */
   bool IsAffineBlockBinding(const StmtSRef& block_sref) const {
+    // (SparseTIR Hack) Always return true for sparse blocks.
+    const auto* block = block_sref->StmtAs<BlockNode>();
+    Optional<ObjectRef> sparse_attr = block != nullptr ? block->annotations.Get("sparse") : NullOpt;
+    if (sparse_attr.defined() && sparse_attr.as<IntImmNode>()->value == 1) {
+      return true;
+    }
+
     return GetBlockInfo(block_sref).affine_binding;
   }
   /*!

src/tir/schedule/analysis/analysis.cc

@@ -168,12 +168,7 @@ Definition of a scope that is a stage pipeline:
         if (it_atomic != block->annotations.end()) {
           is_atomic = ((*it_atomic).second).as<IntImmNode>()->value;
         }
-        auto&& it_sparse = block->annotations.find("sparse");
-        bool is_sparse = false;
-        if (it_sparse != block->annotations.end()) {
-          is_sparse = ((*it_sparse).second).as<IntImmNode>()->value;
-        } 
-        if (!is_sparse && !is_atomic) {
+        if (!is_atomic) {
           throw NotCompactDataFlowError(self->mod, GetRef<Stmt>(scope_root_subtree->stmt),
                                         GetRef<Block>(block));
         }
@@ -445,6 +440,12 @@ void CheckNotOutputBlock(const ScheduleState& self, const StmtSRef& block_sref,
 
 bool IsAffineBinding(const BlockRealize& realize, const Map<Var, Range>& loop_var_ranges,
                      arith::Analyzer* analyzer) {
+  // (SparseTIR Hack) Always return true for sparse blocks.
+  Optional<ObjectRef> sparse_attr = realize->block->annotations.Get("sparse");
+  if (sparse_attr.defined() && sparse_attr.as<IntImmNode>()->value == 1) {
+    return true;
+  }
+
   if (loop_var_ranges.empty()) {
     return true;
   }

src/tir/schedule/primitive/reduction.cc

@@ -212,6 +212,7 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref,
   ObjectPtr<BlockNode> init_block = make_object<BlockNode>();
   ObjectPtr<BlockRealizeNode> init_realize = make_object<BlockRealizeNode>();
   init_block->name_hint = block->name_hint + "_init";
+  init_block->annotations = block->annotations;
   init_realize->iter_values = {};
   init_realize->block = Block(init_block);
   // Step 1. Create new block vars and their bindings
@@ -580,7 +581,10 @@ class BaseBlockCreator {
         /*body=*/new_reduction_update_,
         /*init=*/
         BufferStore(new_reduction_update_->buffer, reducer_->identity_element[0],
-                    new_reduction_update_->indices));
+                    new_reduction_update_->indices),
+        /*alloc_buffers=*/{},
+        /*match_buffers=*/{},
+        /*annotations=*/old_block_realize_->block->annotations);
     new_block_realize_ = BlockRealize(iter_values_, predicate, new_block_);
   }
 

tests/python/sparsetir/test_tir_sparse_correctness.py

@@ -267,10 +267,9 @@ def test_sddmm():
     )
     blk = sch.get_block("sddmm")
     ij, k = sch.get_loops(blk)
-    # TODO(zihao): fix the behavior in the future.
-    # sch.decompose_reduction(blk, ij)
     sch.bind(ij, "blockIdx.x")
     sch.bind(k, "threadIdx.x")
+    sch.decompose_reduction(blk, k)
 
     # convert numpy tensor to tvm ndarray
     C_indices = tvm.nd.array(indices.astype("int32"), device=tvm.cuda(0))

tests/python/unittest/test_tir_schedule_reduction.py

@@ -152,6 +152,39 @@ def matmul_decompose4(a: T.handle, b: T.handle, c: T.handle) -> None:
                     C[vi, vj] = C[vi, vj] + (A[vi, vk] * B[vj, vk])
 
 
+@T.prim_func
+def matmul_with_annotation(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j, k in T.grid(128, 128, 128):
+        with T.block("update"):
+            T.block_attr({"test_annotation": 1})
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@T.prim_func
+def matmul_decompose_with_annotation(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+
+    for i, j in T.grid(128, 128):
+        with T.block("init"):
+            T.block_attr({"test_annotation": 1})
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = 0.0
+
+    for i, j, k in T.grid(128, 128, 128):
+        with T.block("update"):
+            T.block_attr({"test_annotation": 1})
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
 # pylint: enable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
 
 
@@ -201,5 +234,14 @@ def test_reduction_decompose4():
     verify_trace_roundtrip(s, mod=matmul)
 
 
+def test_reduction_decompose_with_annotation():
+    s = tir.Schedule(matmul_with_annotation, debug_mask="all")
+    C = s.get_block("update")
+    i, j, k = s.get_loops(C)
+    s.decompose_reduction(C, i)
+    tvm.ir.assert_structural_equal(matmul_decompose_with_annotation, s.mod["main"])
+    verify_trace_roundtrip(s, mod=matmul_with_annotation)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__] + sys.argv[1:]))

tests/python/unittest/test_tir_schedule_rfactor.py

@@ -214,6 +214,43 @@ def square_sum_square_root_rfactor(a: T.handle, d: T.handle) -> None:
             D[b_2] = T.sqrt(C[b_2], dtype="float32")
 
 
+@T.prim_func
+def square_sum_with_annotation(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [16, 256, 256])
+    C = T.match_buffer(c, [16])
+
+    for b0, i0, j0 in T.grid(16, 256, 256):
+        with T.block("C"):
+            T.block_attr({"test_annotation": 1})
+            b, i, j = T.axis.remap("SRR", [b0, i0, j0])
+            with T.init():
+                C[b] = 0.0
+            C[b] = C[b] + A[b, i, j] * A[b, i, j]
+
+
+@T.prim_func
+def square_sum_with_annotation_rfactor(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [16, 256, 256])
+    C = T.match_buffer(c, [16])
+    C_rf = T.alloc_buffer([16, 256])
+
+    for i0, i1, i2 in T.grid(16, 256, 256):
+        with T.block("C_rf"):
+            T.block_attr({"test_annotation": 1})
+            vi2, b, i = T.axis.remap("SSR", [i2, i0, i1])
+            with T.init():
+                C_rf[b, vi2] = 0.0
+            C_rf[b, vi2] = C_rf[b, vi2] + (A[b, i, vi2] * A[b, i, vi2])
+
+    for i0_1, i2_1 in T.grid(16, 256):
+        with T.block("C"):
+            T.block_attr({"test_annotation": 1})
+            vi2_1, b_1 = T.axis.remap("RS", [i2_1, i0_1])
+            with T.init():
+                C[b_1] = 0.0
+            C[b_1] = C[b_1] + C_rf[b_1, vi2_1]
+
+
 @T.prim_func
 def element_wise(a: T.handle, b: T.handle) -> None:
     A = T.match_buffer(a, (128, 128))
@@ -660,5 +697,16 @@ def test_reduction_rfactor_predicate():  # pylint: disable=invalid-name
     verify_trace_roundtrip(s, mod=rowsum_predicate)
 
 
+def test_reduction_rfactor_with_annotation():
+    s = tir.Schedule(square_sum_with_annotation, debug_mask="all")
+    C = s.get_block("C")
+    _, _, j = s.get_loops(C)
+    rf_block = s.rfactor(j, 1)
+    tvm.ir.assert_structural_equal(s.mod["main"], square_sum_with_annotation_rfactor)
+    assert s.get(rf_block).same_as(s.get(s.get_block("C_rf")))
+    assert s.get(C).same_as(s.get(s.get_block("C")))
+    verify_trace_roundtrip(s, mod=square_sum_with_annotation)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__] + sys.argv[1:]))