[SparseTIR][Schedule] GetSpIters (#24)

Author: MasterJH5574

include/tvm/tir/schedule/schedule.h

@@ -547,6 +547,12 @@ class ScheduleNode : public runtime::Object {
    * \note Indexing error is raised if 0 or multiple blocks exist with the specific name
    */
   virtual SparseBlockRV GetSparseBlock(const String& name, const String& func_name = "main") = 0;
+  /*!
+   * \brief Retrieve the sparse iterators of a given sparse block
+   * \param block_rv The block to be queried
+   * \return The sparse iterators of the input sparse block
+   */
+  virtual Array<SpIterVar> GetSpIters(const SparseBlockRV& block_rv) = 0;
   /*!
    * \brief Reorder a list of sparse iterators. It requires the new order to not break the iterator
    * dependency.

python/tvm/tir/schedule/schedule.py

@@ -1930,6 +1930,24 @@ def get_sparse_block(
             func_name,
         )
 
+    def get_sp_iters(self, block: SparseBlockRV) -> List[SpIterVar]:
+        """Retrieve the sparse iterators of a given sparse block
+
+        Parameters
+        ----------
+        block : SparseBlockRV
+            The block to be queried
+
+        Returns
+        -------
+        sp_iters : List[SpIterVar]
+            The sparse iterators of the input sparse block
+        """
+        return _ffi_api.ScheduleGetSpIters(  # type: ignore # pylint: disable=no-member
+            self,
+            block,
+        )
+
     def sparse_reorder(self, block: SparseBlockRV, new_order: List[SpIterVar]) -> None:
         """Reorder a list of sparse iterators. It requires the new order to not break the iterator
         dependency.

src/tir/schedule/concrete_schedule.cc

@@ -700,6 +700,10 @@ SparseBlockRV ConcreteScheduleNode::GetSparseBlock(const String& name, const Str
   return CreateRV(GetRef<SparseBlock>(block));
 }
 
+Array<SpIterVar> ConcreteScheduleNode::GetSpIters(const SparseBlockRV& block_rv) {
+  return this->Get(block_rv)->sp_iter_vars;
+}
+
 void ConcreteScheduleNode::SparseReorder(const SparseBlockRV& block_rv,
                                          const Array<SpIterVar>& new_order) {
   SparseBlock old_block = this->Get(block_rv);

src/tir/schedule/concrete_schedule.h

@@ -135,6 +135,7 @@ class ConcreteScheduleNode : public ScheduleNode {
   void EnterPostproc() override {}
   /******** Schedule: SparseTIR schedules ********/
   SparseBlockRV GetSparseBlock(const String& name, const String& func_name = "main") override;
+  Array<SpIterVar> GetSpIters(const SparseBlockRV& block_rv) override;
   void SparseReorder(const SparseBlockRV& block_rv, const Array<SpIterVar>& new_order) override;
 
  protected:

src/tir/schedule/schedule.cc

@@ -228,6 +228,8 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleEnterPostproc")
 /******** (FFI) SparseTIR schedules ********/
 TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetSparseBlock")
     .set_body_method<Schedule>(&ScheduleNode::GetSparseBlock);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetSpIters")
+    .set_body_method<Schedule>(&ScheduleNode::GetSpIters);
 TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSparseReorder")
     .set_body_method<Schedule>(&ScheduleNode::SparseReorder);
 

src/tir/schedule/traced_schedule.cc

@@ -416,6 +416,14 @@ SparseBlockRV TracedScheduleNode::GetSparseBlock(const String& name, const Strin
   return result;
 }
 
+Array<SpIterVar> TracedScheduleNode::GetSpIters(const SparseBlockRV& block_rv) {
+  Array<SpIterVar> result = ConcreteScheduleNode::GetSpIters(block_rv);
+
+  // Do not support traced schedule so far.
+
+  return result;
+}
+
 void TracedScheduleNode::SparseReorder(const SparseBlockRV& block_rv,
                                        const Array<SpIterVar>& new_order) {
   ConcreteScheduleNode::SparseReorder(block_rv, new_order);

src/tir/schedule/traced_schedule.h

@@ -96,6 +96,7 @@ class TracedScheduleNode : public ConcreteScheduleNode {
   void EnterPostproc() final;
   /******** Schedule: SparseTIR schedules ********/
   SparseBlockRV GetSparseBlock(const String& name, const String& func_name = "main") final;
+  Array<SpIterVar> GetSpIters(const SparseBlockRV& block_rv) final;
   void SparseReorder(const SparseBlockRV& block, const Array<SpIterVar>& new_order) final;
 };
 

tests/python/sparsetir/test_tir_sparse_schedule.py

@@ -408,35 +408,43 @@ def test_get_sparse_block():
     assert block.same_as(csrmm.body)
 
 
+def test_get_sp_iters():
+    sch = tir.Schedule(csrmm, debug_mask="all")
+    block = sch.get_sparse_block("csrmm")
+    vi, vj, vk = sch.get_sp_iters(block)
+    assert vi.same_as(csrmm.body.sp_iter_vars[0])
+    assert vj.same_as(csrmm.body.sp_iter_vars[1])
+    assert vk.same_as(csrmm.body.sp_iter_vars[2])
+
+
 def test_reorder():
     sch = tir.Schedule(bsrmm, debug_mask="all")
-    block_rv = sch.get_sparse_block("bsrmm")
-    block = sch.get(block_rv)
-    i, j, bi, bj, f = block.sp_iter_vars
-    sch.sparse_reorder(block_rv, [bi, bj, i, j, f])
+    block = sch.get_sparse_block("bsrmm")
+    i, j, bi, bj, f = sch.get_sp_iters(block)
+    sch.sparse_reorder(block, [bi, bj, i, j, f])
     tvm.ir.assert_structural_equal(sch.mod["main"], reordered_bsrmm, True)
+    assert sch.get(block).name == "bsrmm"
 
 
 def test_reorder_fail_on_dependency():
     sch = tir.Schedule(bsrmm, debug_mask="all")
-    block_rv = sch.get_sparse_block("bsrmm")
-    block = sch.get(block_rv)
-    i, j, bi, bj, f = block.sp_iter_vars
+    block = sch.get_sparse_block("bsrmm")
+    i, j, bi, bj, f = sch.get_sp_iters(block)
     with pytest.raises(tvm.tir.ScheduleError):
-        sch.sparse_reorder(block_rv, [bi, bj, j, i, f])
+        sch.sparse_reorder(block, [bi, bj, j, i, f])
 
 
 def test_reorder_fail_on_new_order_length():
     sch = tir.Schedule(bsrmm, debug_mask="all")
-    block_rv = sch.get_sparse_block("bsrmm")
-    block = sch.get(block_rv)
-    i, j, bi, bj, f = block.sp_iter_vars
+    block = sch.get_sparse_block("bsrmm")
+    i, j, bi, bj, f = sch.get_sp_iters(block)
     with pytest.raises(tvm.tir.ScheduleError):
-        sch.sparse_reorder(block_rv, [bi, bj, i, j])
+        sch.sparse_reorder(block, [bi, bj, i, j])
 
 
 if __name__ == "__main__":
     test_get_sparse_block()
+    test_get_sp_iters()
     test_reorder()
     test_reorder_fail_on_dependency()
     test_reorder_fail_on_new_order_length()