Add docstring for sparse tir lowering (#21)

* add docstring * upd
yzh119 · Jan 24, 2022 · a4213c9 · a4213c9
1 parent 9a58063
commit a4213c9
Show file tree

Hide file tree

Showing 4 changed files with 234 additions and 102 deletions.
diff --git a/include/tvm/tir/stmt.h b/include/tvm/tir/stmt.h
@@ -1288,7 +1288,7 @@ class SparseBlockNode : public StmtNode {
   /*! \brief The sparse data structures */
   Array<ObjectRef> sp_structs;
   /*! \brief The mapping from sparse data structures to the PrimFunc parameters */
-  Map<ObjectRef, Array<Var>> sp_struct2param_map;
+  Map<ObjectRef, Array<Var>> sp_struct_param_map;
   /*! \brief The name of the block */
   String name;
   /*! \brief The body of the block */
@@ -1299,7 +1299,7 @@ class SparseBlockNode : public StmtNode {
   void VisitAttrs(AttrVisitor* v) {
     v->Visit("sp_iter_vars", &sp_iter_vars);
     v->Visit("sp_structs", &sp_structs);
-    v->Visit("sp_struct2param_map", &sp_struct2param_map);
+    v->Visit("sp_struct_param_map", &sp_struct_param_map);
     v->Visit("name", &name);
     v->Visit("body", &body);
     v->Visit("init", &init);

diff --git a/src/printer/tvmscript_printer.cc b/src/printer/tvmscript_printer.cc
@@ -1372,7 +1372,7 @@ Doc TVMScriptPrinter::PrintSparseStructDefinitions(const SparseBlockNode* sp_blo
   std::vector<Doc> sp_buf_docs;
 
   for (const ObjectRef& obj : sp_block->sp_structs) {
-    Array<Var> params = sp_block->sp_struct2param_map.Get(obj).value();
+    Array<Var> params = sp_block->sp_struct_param_map.Get(obj).value();
 
     Doc doc;
     doc << Print(obj) << " = " << tir_prefix_ << ".";

diff --git a/src/tir/ir/stmt.cc b/src/tir/ir/stmt.cc
@@ -981,7 +981,7 @@ SparseBlock::SparseBlock(Array<SpIterVar> sp_iter_vars, Array<ObjectRef> sp_stru
   CHECK_EQ(sp_structs.size(), sp_struct_params.size())
       << "ValueError: The length of `sp_struct_params` is expected to be equal to the length "
          "`sp_structs`, which is the number of sparse data structures";
-  Map<ObjectRef, Array<Var>> sp_struct2param_map;
+  Map<ObjectRef, Array<Var>> sp_struct_param_map;
   for (int i = 0; i < static_cast<int>(sp_structs.size()); ++i) {
     ObjectRef obj = sp_structs[i];
     Array<Var> params = sp_struct_params[i];
@@ -1005,13 +1005,13 @@ SparseBlock::SparseBlock(Array<SpIterVar> sp_iter_vars, Array<ObjectRef> sp_stru
       LOG(FATAL) << "ValueError: " << obj->_type_key << " is not a sparse data structure";
     }
 
-    sp_struct2param_map.Set(obj, params);
+    sp_struct_param_map.Set(obj, params);
   }
 
   ObjectPtr<SparseBlockNode> node = make_object<SparseBlockNode>();
   node->sp_iter_vars = std::move(sp_iter_vars);
   node->sp_structs = std::move(sp_structs);
-  node->sp_struct2param_map = std::move(sp_struct2param_map);
+  node->sp_struct_param_map = std::move(sp_struct_param_map);
   node->name = std::move(name);
   node->body = std::move(body);
   node->init = std::move(init);