[TIR.Constant] U1 usecase

Constants are now aggregated into one struct and initialized in default_lib0.c
file

Change-Id: I34d61f8139c8a92c06944fe990ba892a660476fd

Author: d-smirnov

include/tvm/tir/stmt.h

@@ -673,7 +673,9 @@ class AllocateConst : public Stmt {
    * create AllocateConstNode with irmod_storage_idx or data
    */
   TVM_DLL AllocateConst(Var buffer_var, DataType dtype, Array<PrimExpr> extents,
-                        ObjectRef data_or_idx, Stmt body, Span span = Span());
+                        ObjectRef data_or_idx, Stmt body,
+                        Map<String, ObjectRef> annotations = Map<String, ObjectRef>(),
+                        Span span = Span());
   TVM_DEFINE_OBJECT_REF_METHODS(AllocateConst, Stmt, AllocateConstNode);
 };
 

include/tvm/tir/usmp/algorithms.h

@@ -53,6 +53,17 @@ Map<BufferInfo, PoolAllocation> GreedyBySize(const Array<BufferInfo>& buffer_inf
  */
 Map<BufferInfo, PoolAllocation> GreedyByConflicts(const Array<BufferInfo>& buffer_info_arr,
                                                   const Integer& memory_pressure);
+/*!
+ *\brief The Hill-Climb algoritm to plan memory
+ *
+ * This will perform an attempt to utilize probabalistic approach to memory
+ * allocation. Typically better than greedy family, but quite slow due to large
+ * number of iterations.
+ *
+ * \return A Map of BufferInfo objects and their associated PoolAllocation
+ */
+Map<BufferInfo, PoolAllocation> HillClimb(const Array<BufferInfo>& buffer_info_arr,
+                                          const Integer& memory_pressure);
 
 /*!
  * \brief The Hill-Climb algorithm to plan memory

include/tvm/tir/usmp/utils.h

@@ -44,6 +44,126 @@ constexpr const char* kUSMPAlgorithmOption = "tir.usmp.algorithm";
 
 namespace tir {
 namespace usmp {
+/*
+ * \brief The ConstantInfoNode contains numeric literal in RO pool
+ */
+struct ConstantInfoNode : public Object {
+  String name_hint;
+  Integer byte_alignment;
+  Integer byte_offset;
+  runtime::NDArray data;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("constant_names", &name_hint);
+    v->Visit("constant_alignment", &byte_alignment);
+    v->Visit("constant_offsets", &byte_offset);
+    v->Visit("constant_data", &data);
+  }
+
+  bool SEqualReduce(const ConstantInfoNode* other, SEqualReducer equal) const {
+    return equal(name_hint, other->name_hint) && equal(byte_alignment, other->byte_alignment) &&
+           equal(byte_offset, other->byte_offset) && equal(data, other->data);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce(name_hint);
+    hash_reduce(byte_alignment);
+    hash_reduce(byte_offset);
+    hash_reduce(data);
+  }
+
+  static constexpr const char* _type_key = "tir.usmp.ConstantInfo";
+  static constexpr bool _type_has_method_sequal_reduce = true;
+  static constexpr bool _type_has_method_shash_reduce = true;
+  TVM_DECLARE_FINAL_OBJECT_INFO(ConstantInfoNode, Object);
+};
+
+class ConstantInfo : public ObjectRef {
+ public:
+  TVM_DLL ConstantInfo(String name, Integer byte_alignment, Integer byte_offset,
+                       runtime::NDArray data);
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ConstantInfo, ObjectRef, ConstantInfoNode);
+};
+
+#if 0
+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_bytes is set to kUnrestrictedPoolSizeHint
+   * to indicate the pool is not size restricted.
+   */
+  Integer size_hint_bytes;
+  /*! \brief The accessibility from each Target */
+  Map<Target, String> target_access;  // 'rw' or 'ro'
+  /*! \brief The clock frequency of the memory in Hz */
+  Integer clock_frequency_hz;
+  /*! \brief The read bandwidth in bytes/cycle */
+  Integer read_bandwidth_bytes_per_cycle;
+  /*! \brief The write bandwidth in bytes/cycle */
+  Integer write_bandwidth_bytes_per_cycle;
+  /*! \brief The read latency in cycles */
+  Integer read_latency_cycles;
+  /*! \brief The write latency in cycles */
+  Integer write_latency_cycles;
+  /*! \brief The burst length in bytes for each Target */
+  Map<Target, Integer> target_burst_bytes;
+  /*! \brief Whether pool is internally generated.
+   * The internal pools will be generated as part of
+   * the entry point code generation of the executor
+   */
+  bool is_internal = false;
+
+  Array<ConstantInfo> constant_info_arr;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("pool_name", &pool_name);
+    v->Visit("size_hint_bytes", &size_hint_bytes);
+    v->Visit("target_access", &target_access);
+    v->Visit("clock_frequency_hz", &clock_frequency_hz);
+    v->Visit("read_bandwidth_bytes_per_cycle", &read_bandwidth_bytes_per_cycle);
+    v->Visit("write_bandwidth_bytes_per_cycle", &write_bandwidth_bytes_per_cycle);
+    v->Visit("read_latency_cycles", &read_latency_cycles);
+    v->Visit("write_latency_cycles", &write_latency_cycles);
+    v->Visit("target_burst_bytes", &target_burst_bytes);
+    v->Visit("is_internal", &is_internal);
+    v->Visit("constant_info_arr", &constant_info_arr);
+  }
+
+  bool SEqualReduce(const PoolInfoNode* other, SEqualReducer equal) const {
+    return equal(pool_name, other->pool_name) && equal(size_hint_bytes, other->size_hint_bytes) &&
+           equal(target_access, other->target_access) &&
+           equal(target_access, other->target_access) &&
+           equal(clock_frequency_hz, other->clock_frequency_hz) &&
+           equal(read_bandwidth_bytes_per_cycle, other->read_bandwidth_bytes_per_cycle) &&
+           equal(write_bandwidth_bytes_per_cycle, other->write_bandwidth_bytes_per_cycle) &&
+           equal(read_latency_cycles, other->read_latency_cycles) &&
+           equal(write_latency_cycles, other->write_latency_cycles) &&
+           equal(target_burst_bytes, other->target_burst_bytes) &&
+           equal(is_internal, other->is_internal) &&
+           equal(constant_info_arr, other->constant_info_arr);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce(pool_name);
+    hash_reduce(size_hint_bytes);
+    hash_reduce(target_access);
+    hash_reduce(clock_frequency_hz);
+    hash_reduce(read_bandwidth_bytes_per_cycle);
+    hash_reduce(write_bandwidth_bytes_per_cycle);
+    hash_reduce(read_latency_cycles);
+    hash_reduce(write_latency_cycles);
+    hash_reduce(target_burst_bytes);
+    hash_reduce(is_internal);
+    hash_reduce(constant_info_arr);
+  }
+
+  static constexpr const char* _type_key = "tir.usmp.PoolInfo";
+  static constexpr bool _type_has_method_sequal_reduce = true;
+  static constexpr bool _type_has_method_shash_reduce = true;
+  TVM_DECLARE_FINAL_OBJECT_INFO(PoolInfoNode, Object);
+};
+#endif
 
 /*!
  * \brief Describes an abstract memory buffer that will get allocated inside a pool.
@@ -150,20 +270,25 @@ class BufferInfoAnalysis : public ObjectRef {
 struct PoolAllocationNode : public Object {
   /*! \brief The assigned PoolInfo object */
   PoolInfo pool_info;
+  /*! \brief The byte alignment where the tensor is supposed to be placed within the pool*/
+  Integer byte_alignment;
   /*! \brief The byte offset where the tensor is supposed to be placed within the pool*/
   Integer byte_offset;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("pool_info", &pool_info);
+    v->Visit("byte_alignment", &byte_alignment);
     v->Visit("byte_offset", &byte_offset);
   }
 
   bool SEqualReduce(const PoolAllocationNode* other, SEqualReducer equal) const {
-    return equal(pool_info, other->pool_info) && equal(byte_offset, other->byte_offset);
+    return equal(pool_info, other->pool_info) && equal(byte_alignment, other->byte_alignment) &&
+           equal(byte_offset, other->byte_offset);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
     hash_reduce(pool_info);
+    hash_reduce(byte_alignment);
     hash_reduce(byte_offset);
   }
 
@@ -173,7 +298,7 @@ struct PoolAllocationNode : public Object {
 
 class PoolAllocation : public ObjectRef {
  public:
-  TVM_DLL PoolAllocation(PoolInfo pool_info, Integer byte_offset);
+  TVM_DLL PoolAllocation(PoolInfo pool_info, Integer byte_alignment, Integer byte_offset);
   TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolAllocation, ObjectRef, PoolAllocationNode);
 };
 
@@ -187,22 +312,26 @@ struct AllocatedPoolInfoNode : public Object {
   Integer allocated_size;
   /*! \brief An optional associated pool Var*/
   Optional<Var> pool_var;
+  /*! \brief pool initialization data */
+  Array<ConstantInfo> constant_info_arr;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("pool_info", &pool_info);
     v->Visit("allocated_size", &allocated_size);
     v->Visit("pool_var", &pool_var);
+    v->Visit("constant_info_arr", &constant_info_arr);
   }
 
   bool SEqualReduce(const AllocatedPoolInfoNode* other, SEqualReducer equal) const {
     return equal(pool_info, other->pool_info) && equal(allocated_size, other->allocated_size) &&
-           equal(pool_var, other->pool_var);
+           equal(pool_var, other->pool_var) && equal(constant_info_arr, other->constant_info_arr);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
     hash_reduce(pool_info);
     hash_reduce(allocated_size);
     hash_reduce(pool_var);
+    hash_reduce(constant_info_arr);
   }
 
   static constexpr const char* _type_key = "tir.usmp.AllocatedPoolInfo";
@@ -211,7 +340,8 @@ struct AllocatedPoolInfoNode : public Object {
 
 class AllocatedPoolInfo : public ObjectRef {
  public:
-  TVM_DLL AllocatedPoolInfo(PoolInfo pool_info, Integer allocated_size, Var pool_var = Var());
+  TVM_DLL AllocatedPoolInfo(PoolInfo pool_info, Integer allocated_size, Var pool_var = Var(),
+                            Array<ConstantInfo> = {});
   TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(AllocatedPoolInfo, ObjectRef, AllocatedPoolInfoNode);
 };
 
@@ -243,6 +373,13 @@ static constexpr const char* kPoolCandidatesAllocateAttr = "candidate_memory_poo
  */
 Integer CalculateExtentsSize(const AllocateNode* op);
 
+/*!
+ * \brief Calculate the size of the extents in bytes
+ *
+ * \param op the allocate const node
+ */
+Integer CalculateExtentsSize(const AllocateConstNode* op);
+
 /*!
  * \brief Joins the Stmt nodes with PoolAllocation objects
  *
@@ -268,7 +405,6 @@ static constexpr const char* kPoolArgs = "pool_args";
  * as an Array.
  */
 static constexpr const char* kPoolInfoIRModuleAttr = "pool_infos";
-
 }  // namespace attr
 
 }  // namespace tvm

python/tvm/script/tir/scope_handler.py

@@ -166,12 +166,20 @@ class AllocateConst(WithScopeHandler):
     """
 
     def __init__(self):
-        def allocate_const(raw_data, dtype, shape, span=None):
+        def allocate_const(raw_data, dtype, shape, annotations=None, span=None):
             list_data = []
             for i in raw_data:
                 list_data.append(i.value)
             nd_data = tvm.nd.array(np.asarray(list_data, dtype=dtype))
-            n = tvm.tir.AllocateConst(self.buffer_var, dtype, shape, nd_data, self.body, span=span)
+            n = tvm.tir.AllocateConst(
+                self.buffer_var,
+                dtype,
+                shape,
+                nd_data,
+                self.body,
+                annotations=annotations,
+                span=span,
+            )
             return n
 
         super().__init__(allocate_const, concise_scope=True, def_symbol=True)
@@ -199,7 +207,7 @@ def enter_scope(
         else:
             raise Exception("Internal Bug")
 
-        def setup_buffer_var(data, dtype, shape, span: Span = None):
+        def setup_buffer_var(data, dtype, shape, annotations: dict = None, span: Span = None):
             """Setup buffer var for a given type."""
             buffer_ptr_type = tvm.ir.PointerType(tvm.ir.PrimType(dtype))
             self.buffer_var = tvm.tir.Var(name, buffer_ptr_type, span)

python/tvm/tir/stmt.py

@@ -364,13 +364,16 @@ class AllocateConst(Stmt):
     body : Stmt
         The body statement.
 
+    annotations : Optional[Map]
+        Additional annotations about the allocation.
+
     span : Optional[Span]
         The location of this itervar in the source code.
     """
 
-    def __init__(self, buffer_var, dtype, extents, condition, body, span=None):
+    def __init__(self, buffer_var, dtype, extents, condition, body, annotations=None, span=None):
         self.__init_handle_by_constructor__(
-            _ffi_api.AllocateConst, buffer_var, dtype, extents, condition, body, span
+            _ffi_api.AllocateConst, buffer_var, dtype, extents, condition, body, annotations, span
         )
 
 

python/tvm/tir/usmp/utils.py

@@ -83,14 +83,18 @@ class PoolAllocation(Object):
     pool_info : PoolInfo
         The PoolInfo to which this allocation corresponds to
 
+    byte_alignment : int
+        The alignment in the pool where the allocate node should be placed
+
     byte_offset : int
         The offset in the pool where the allocate node should be placed
 
     """
 
-    def __init__(self, pool_info: PoolInfo, byte_offset: int):
+    def __init__(self, pool_info: PoolInfo, byte_alignment: int, byte_offset: int):
         self.__init_handle_by_constructor__(
             _ffi_api.PoolAllocation,  # type: ignore # pylint: disable=no-member
             pool_info,
+            byte_alignment,
             byte_offset,
         )

src/relay/backend/aot_executor_codegen.cc

@@ -706,9 +706,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
    * brief Run USMP to plan memory for lowered IRModule
    */
   IRModule PlanMemoryWithUSMP(const IRModule& mod) {
-    Executor executor_config = mod->GetAttr<Executor>(tvm::attr::kExecutor).value();
-    Integer workspace_byte_alignment =
-        executor_config->GetAttr<Integer>("workspace-byte-alignment").value_or(16);
+    Integer workspace_byte_alignment = getModuleAlignment(mod);
     IRModule lowered_mod = mod->ShallowCopy();
     lowered_mod = tir::transform::UnifiedStaticMemoryPlanner()(lowered_mod);
     // Update workspace size based on the pool allocations.
@@ -748,9 +746,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
    * brief Run StorageRewrite to plan memory for lowered IRModule
    */
   IRModule PlanMemoryWithStorageRewrite(const IRModule& mod) {
-    Executor executor_config = mod->GetAttr<Executor>(tvm::attr::kExecutor).value();
-    Integer workspace_byte_alignment =
-        executor_config->GetAttr<Integer>("workspace-byte-alignment").value_or(16);
+    Integer workspace_byte_alignment = getModuleAlignment(mod);
     IRModule lowered_mod = mod->ShallowCopy();
     // Running StorageRewrite just on the main function
     tir::PrimFunc tir_main_func =
@@ -773,6 +769,11 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     return lowered_mod;
   }
 
+  Integer getModuleAlignment(const IRModule& mod) {
+    Executor executor_config = mod->GetAttr<Executor>(tvm::attr::kExecutor).value();
+    return executor_config->GetAttr<Integer>("workspace-byte-alignment").value_or(16);
+  }
+
  protected:
   /*! \brief mod */
   runtime::Module* mod_;
@@ -837,10 +838,10 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     ICHECK(target_host_.defined()) << "require a target_host to be given for AOT codegen";
     VLOG(1) << "target host: " << target_host_->ToDebugString();
 
+    Integer workspace_byte_alignment = getModuleAlignment(mod);
+
     Executor executor_config = mod->GetAttr<Executor>(tvm::attr::kExecutor).value();
     String interface_api = executor_config->GetAttr<String>("interface-api").value_or("packed");
-    Integer workspace_byte_alignment =
-        executor_config->GetAttr<Integer>("workspace-byte-alignment").value_or(16);
     use_unpacked_api_ = executor_config->GetAttr<Bool>("unpacked-api").value_or(Bool(false));
 
     // TODO(mbs): Plumb from compiler config

src/target/source/codegen_params.cc

@@ -238,7 +238,7 @@ void NDArrayDataToC(::tvm::runtime::NDArray arr, int indent_chars, std::ostream&
     }
 
     default:
-      CHECK(false) << "Data type not supported";
+      CHECK(false) << "Data type '" << arr_type << "' not supported";
   }
 
   os.flags(old_fmtflags);