[TIR][USMP] Added buffer info extraction pass

Adding more documentation for data structures
and the approach

Change-Id: Ide2bfffaeff9add86853b6992017264e5d796299

Author: manupak

include/tvm/tir/usmp/utils.h

@@ -33,7 +33,17 @@ namespace tvm {
 namespace tir {
 namespace usmp {
 
+/*!
+ * \brief The string parameter to indicate read and write access to a pool
+ * This needs to be kept in sync with PoolInfo.READ_WRITE_ACCESS in
+ * python/tvm/tir/usmp/utils.py
+ */
 static constexpr const char* kTargetPoolReadWriteAccess = "rw";
+/*!
+ * \brief The string parameter to indicate read only access to a pool
+ * This needs to be kept in sync with PoolInfo.READ_ONLY_ACCESS in
+ * python/tvm/tir/usmp/utils.py
+ */
 static constexpr const char* kTargetPoolReadOnlyAccess = "ro";
 
 /*!
@@ -43,8 +53,8 @@ struct PoolInfoNode : public Object {
   /*! \brief The name of the memory pool */
   String pool_name;
   /*! \brief The expected size hint to be used by the allocator.
-   * The size_hint is defaulted to -1 to indicate the pool is not
-   * size restricted.
+   * The size_hint_bytes is defaulted to kUnrestrictedPoolSizeHint
+   * to indicate the pool is not size restricted.
    */
   Integer size_hint_bytes;
   /*! \brief The accessibility from each Target*/
@@ -71,10 +81,15 @@ struct PoolInfoNode : public Object {
   TVM_DECLARE_FINAL_OBJECT_INFO(PoolInfoNode, Object);
 };
 
+/*!
+ * \brief The PoolSize is unrestricted for the memory planner
+ */
+static const int kUnrestrictedPoolSizeHint = -1;
+
 class PoolInfo : public ObjectRef {
  public:
   TVM_DLL PoolInfo(String pool_name, Map<Target, String> target_access,
-                   Integer size_hint_bytes = -1);
+                   Integer size_hint_bytes = kUnrestrictedPoolSizeHint);
   TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
 };
 
@@ -172,7 +187,12 @@ class PoolAllocation : public ObjectRef {
  */
 Array<BufferInfo> CreateArrayBufferInfo(const Map<Stmt, BufferInfo>& buffer_info_map);
 
-static constexpr const char* kPoolCandidatesIRModAttr = "candidate_memory_pools";
+/*!
+ * \brief The allocate node attribute to indicate candidate memory pools.
+ * This needs to be kept in sync with CANDIDATE_MEMORY_POOL_ATTR in
+ * python/tvm/tir/usmp/utils.py.
+ */
+static constexpr const char* kPoolCandidatesAllocateAttr = "candidate_memory_pools";
 
 }  // namespace usmp
 }  // namespace tir

python/tvm/tir/usmp/utils.py

@@ -24,6 +24,10 @@
 from tvm.target import Target
 from . import _ffi_api
 
+
+# The allocate node attribute to indicate candidate memory pools.
+# This needs to be kept in sync with CANDIDATE_MEMORY_POOL_ATTR in
+# include/tvm/tir/usmp/utils.h
 CANDIDATE_MEMORY_POOL_ATTR = "candidate_memory_pools"
 
 
@@ -50,11 +54,20 @@ class PoolInfo(Object):
 
     """
 
+    # The string parameter to indicate read and write access to a pool
+    # This needs to be kept in sync with kTargetPoolReadWriteAccess in
+    # include/tvm/tir/usmp/utils.h
     READ_WRITE_ACCESS = "rw"
+    # The string parameter to indicate read only access to a pool
+    # This needs to be kept in sync with kTargetPoolReadOnlyAccess in
+    # include/tvm/tir/usmp/utils.h
     READ_ONLY_ACCESS = "ro"
 
     def __init__(
-        self, pool_name: str, target_access: Dict[Target, str], size_hint_bytes: Optional[int] = -1
+        self,
+        pool_name: str,
+        target_access: Dict[Target, str],
+        size_hint_bytes: Optional[int] = None,
     ):
         self.__init_handle_by_constructor__(
             _ffi_api.PoolInfo,  # type: ignore # pylint: disable=no-member

src/tir/usmp/analysis/extract_buffer_info.cc

@@ -18,8 +18,12 @@
  */
 
 /*!
- * \file tir/analysis/usmp/convert_for_loops_serial.cc
- * \brief Convert all for loops to serial for lesser memory consumption
+ * \file tir/analysis/usmp/extract_buffer_info.cc
+ *
+ * \brief This analysis pass consumes a TIR IRModule with a main function
+ * that defines a ordering in the calles to operators and produces BufferInfo
+ * objects that contains information about tir.allocate nodes and liveness
+ * conflicts between other tir.allocate nodes.
  */
 #include <tvm/arith/analyzer.h>
 #include <tvm/runtime/device_api.h>
@@ -34,6 +38,12 @@ namespace tvm {
 namespace tir {
 namespace usmp {
 
+/*! \brief This class takes a TIR IRModule and a main PrimFunc that contains
+ * that defines a ordering in the calles to operators and produces BufferInfo
+ * objects that contains information about tir.allocate nodes and liveness
+ * conflicts between other tir.allocate nodes.
+ */
+
 class BufferInfoExtractor : public StmtExprVisitor {
  public:
   explicit BufferInfoExtractor(const IRModule& module) : module_(module) {
@@ -63,6 +73,11 @@ class BufferInfoExtractor : public StmtExprVisitor {
 
   std::unordered_set<Stmt, ObjectPtrHash, ObjectPtrEqual> currently_live_allocates;
   int current_stmt_idx_ = 0;
+  // This structure is supposed to contain information
+  // around the scope the visitor is currently in.
+  // We only check whether the current scope belong to
+  // a Serial ForKind. We are not planning for Parallel
+  // ForKind just yet.
   struct ScopeInfo {
     For for_loop;
   };
@@ -77,39 +92,44 @@ void BufferInfoExtractor::VisitStmt(const Stmt& n) {
   StmtExprVisitor::VisitStmt(n);
 }
 
-static size_t CalculateExtentsSize(const AllocateNode* op) {
+static Integer CalculateExtentsSize(const AllocateNode* op) {
   size_t element_size_bytes = op->dtype.bytes();
   size_t num_elements = 1;
   for (const auto& ext : op->extents) {
     if (ext->IsInstance<IntImmNode>()) {
       num_elements *= Downcast<IntImm>(ext)->value;
     } else {
       // We can't statically calculate workspace for dynamic shapes
-      num_elements = 0;
+      return Integer();
     }
   }
-  return (num_elements * element_size_bytes);
+  return Integer(num_elements * element_size_bytes);
 }
 
 void BufferInfoExtractor::VisitStmt_(const AllocateNode* op) {
   const auto& currect_scope_info = scope_stack_.top();
   const auto& type = Downcast<PointerType>(op->buffer_var->type_annotation);
   const auto& storage_scope = type->storage_scope;
 
-  // If the allocate is in a for loop,
-  // USMP currently only looks at serial for loops.
+  // If the allocate is in a for loop, USMP currently only looks at serial for loops.
+  // If its not a serial for loop, then memory planner will omit them in the current memory planning
+  // process leaving them to as tir.allocate nodes for codegen. Additionally, the USMP can only work
+  // with buffers that have global storage_scope
   if ((!currect_scope_info.for_loop.defined()) ||
       (currect_scope_info.for_loop.defined() &&
        currect_scope_info.for_loop->kind == ForKind::kSerial && storage_scope == "global")) {
-    // USMP can only work with buffers that have global storage_scope
     auto size_bytes = CalculateExtentsSize(op);
     // We only statically memory plan only allocates with known
     // compile time sizes.
-    if (size_bytes) {
+    if (size_bytes.defined()) {
       // By default, the core compiler is assumed to attach the a default pool to each allocate.
-      ICHECK(op->annotations.count(kPoolCandidatesIRModAttr))
+      ICHECK(op->annotations.count(kPoolCandidatesAllocateAttr))
           << "Every statically sized allocate node needs an pool candidate attribute";
-      auto pool_candidates = Downcast<Array<PoolInfo>>(op->annotations[kPoolCandidatesIRModAttr]);
+      auto pool_candidates =
+          Downcast<Array<PoolInfo>>(op->annotations[kPoolCandidatesAllocateAttr]);
+
+      // TODO(@manupa-arm): improve the error when the responsible component for attaching a single
+      // pool is added
       ICHECK(pool_candidates.size() > 0)
           << "The core compiler should at least attach a single PoolInfo. If there were no "
              "user-given arguments for memory pools, the default behaviour is a single size "
@@ -203,6 +223,13 @@ void BufferInfoExtractor::VisitExpr_(const CallNode* op) {
 Map<BufferInfo, tir::Stmt> BufferInfoExtractor::operator()(const PrimFunc& main_func) {
   this->VisitStmt(main_func->body);
 
+  // A liveness event is an event that when
+  // traversing the tir.Stmts where tir.allocate node
+  // begins or ceases to be Live. This particular struct
+  // is used to solve interval overlap problem using
+  // a sweep-line algorithm. For that, we need to record
+  // where the liveness event occurred in a chronological
+  // order.
   enum LivenessEventType { START = 0, END = 1 };
   struct LivenessEvent {
     size_t tick;
@@ -216,6 +243,8 @@ Map<BufferInfo, tir::Stmt> BufferInfoExtractor::operator()(const PrimFunc& main_
     }
   };
 
+  // Create a vector of liveness events
+  // associated with each BufferNodes.
   std::vector<LivenessEvent> le_events;
   for (const auto& kv : buffer_info_map_) {
     if (!kv.second->IsInstance<AllocateNode>()) {
@@ -240,6 +269,9 @@ Map<BufferInfo, tir::Stmt> BufferInfoExtractor::operator()(const PrimFunc& main_
     le_events.push_back(le_event_end);
   }
 
+  // Sort the liveness events based on the chronological
+  // ordering. For events that are simultaneous, START event
+  // takes precedence.
   std::sort(le_events.begin(), le_events.end(),
             [](const LivenessEvent& lhs, const LivenessEvent& rhs) {
               if (lhs.tick < rhs.tick) {
@@ -249,6 +281,9 @@ Map<BufferInfo, tir::Stmt> BufferInfoExtractor::operator()(const PrimFunc& main_
               }
               return false;
             });
+
+  // Traverse the liveness events using a open set to track what
+  // is live while updating the conflicts through out the linear traversal
   std::unordered_set<BufferInfo, ObjectPtrHash, ObjectPtrEqual> open_set;
   for (const auto& le_event : le_events) {
     if (le_event.le_type == START) {
@@ -258,7 +293,6 @@ Map<BufferInfo, tir::Stmt> BufferInfoExtractor::operator()(const PrimFunc& main_
       }
       open_set.insert(le_event.buffer_info);
     } else {
-      ICHECK(le_event.le_type == END);
       open_set.erase(le_event.buffer_info);
     }
   }

src/tir/usmp/utils.cc

@@ -78,7 +78,10 @@ TVM_REGISTER_NODE_TYPE(PoolInfoNode);
 TVM_REGISTER_GLOBAL("tir.usmp.PoolInfo")
     .set_body_typed([](String pool_name, Map<Target, String> target_access,
                        Integer size_hint_bytes) {
-      return PoolInfo(pool_name, target_access, size_hint_bytes);
+      if (size_hint_bytes.defined()) {
+        return PoolInfo(pool_name, target_access, size_hint_bytes);
+      }
+      return PoolInfo(pool_name, target_access);
     });
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)

tests/python/unittest/test_tir_usmp_analysis_extract_bufferinfo.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import pytest
+import sys
 
 import tvm
 from tvm import tir, script
@@ -79,6 +80,9 @@ def _assign_poolinfos_to_allocates_in_irmodule(mod, pool_infos):
     return ret
 
 
+# These are test IRModules that contains varied topologies of operator graphs
+# that includes a main TIR function that includes call to such operators.
+
 # fmt: off
 @tvm.script.ir_module
 class LinearStructure:
@@ -846,4 +850,4 @@ def test_inception_structure():
 
 
 if __name__ == "__main__":
-    pytest.main([__file__])
+    pytest.main([__file__] + sys.argv[1:])

tests/python/unittest/test_tir_usmp_utils.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import pytest
+import sys
 
 import tvm
 from tvm.script import tir as T
@@ -188,4 +189,4 @@ def test_create_array_buffer_info():
 
 
 if __name__ == "__main__":
-    pytest.main([__file__])
+    pytest.main([__file__] + sys.argv[1:])