You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2022/04/01 13:30:21 UTC

[GitHub] [tvm] manupa-arm commented on a change in pull request #10189: [USMP] Adding support for U1 usecase for constant pools

manupa-arm commented on a change in pull request #10189:
URL: https://github.com/apache/tvm/pull/10189#discussion_r840570103



##########
File path: include/tvm/ir/memory_pools.h
##########
@@ -40,8 +40,12 @@ struct PoolInfoNode : public Object {
    * to indicate the pool is not size restricted.
    */
   Integer size_hint_bytes;
+
+ private:
   /*! \brief The accessibility from each Target */
   Map<Target, String> target_access;  // 'rw' or 'ro'

Review comment:
       Lets remove this in the favor of array of targets

##########
File path: python/tvm/ir/memory_pools.py
##########
@@ -112,6 +112,137 @@ def __init__(
         )
 
 
+@register_object("ir.PoolInfoProperties")
+class PoolInfoProperties(Object):
+    """PoolInfo object holds information related to memory pools
+    where the statically sized allocate nodes will pooled into.
+
+    Parameters
+    ----------
+    size_hint_bytes : Optional[int]
+        The expected size hint to be used by the allocator.
+        The default value would be -1 which means the pool
+        is not size restricted.
+
+    clock_frequency_hz : Optional[int]
+        The clock frequency that the memory pool runs at in Hz.
+        If not specified/known, this will default to -1 indicating
+        it hasn't been defined.
+
+    read_bandwidth_bytes_per_cycle : Optional[int]
+        The read bandwidth of the memory pool in bytes/cycle.
+        If not specified/known, this will default to -1 indicating
+        it hasn't been defined.
+
+    write_bandwidth_bytes_per_cycle : Optional[int]
+        The write bandwidth of the memory pool in bytes/cycle.
+        If not specified/known, this will default to -1 indicating
+        it hasn't been defined.
+
+    read_latency_cycles : Optional[int]
+        The read latency of the memory pool in cycles.
+        If not specified/known, this will default to 0.
+
+    write_latency_cycles : Optional[int]
+        The write latency of the memory pool in cycles.
+        If not specified/known, this will default to 0.
+
+    target_burst_bytes : Optional[Union[Dict[Target, int], None]]
+        The burst length of the memory pool in bytes per target.
+        If not specified/known for a given target, a burst length
+        of 1 byte will be assumed.
+
+    """
+
+    def __init__(
+        self,
+        size_hint_bytes: Optional[int] = -1,
+        clock_frequency_hz: Optional[int] = -1,
+        read_bandwidth_bytes_per_cycle: Optional[int] = -1,
+        write_bandwidth_bytes_per_cycle: Optional[int] = -1,
+        read_latency_cycles: Optional[int] = 0,
+        write_latency_cycles: Optional[int] = 0,
+        target_burst_bytes=None,  # Optional[Union[Dict[target.Target, int], None]]
+    ):
+        if not target_burst_bytes:
+            target_burst_bytes = dict()
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.PoolInfoProperties,  # type: ignore # pylint: disable=no-member
+            size_hint_bytes,
+            clock_frequency_hz,
+            read_bandwidth_bytes_per_cycle,
+            write_bandwidth_bytes_per_cycle,
+            read_latency_cycles,
+            write_latency_cycles,
+            target_burst_bytes,
+        )
+
+
+@register_object("ir.WorkspacePoolInfo")
+class WorkspacePoolInfo(PoolInfo):
+    """WorkspacePoolInfo object holds information related to RW memory pools
+    where the statically sized allocate nodes will pooled into.
+
+    Parameters
+    ----------
+    pool_name : str
+        The name of the memory pool
+
+    targets : list[Target]
+        A list of targets which could access the pool
+
+    pool_info_properties : PoolInfoProperties
+        The properties of the pool.
+    """
+
+    # pylint: disable=W0231

Review comment:
       Why are disabling this linter warning ? (i.e. could we not call base class init here ? )

##########
File path: python/tvm/ir/memory_pools.py
##########
@@ -112,6 +112,137 @@ def __init__(
         )
 
 
+@register_object("ir.PoolInfoProperties")
+class PoolInfoProperties(Object):
+    """PoolInfo object holds information related to memory pools
+    where the statically sized allocate nodes will pooled into.
+
+    Parameters
+    ----------
+    size_hint_bytes : Optional[int]
+        The expected size hint to be used by the allocator.
+        The default value would be -1 which means the pool
+        is not size restricted.
+
+    clock_frequency_hz : Optional[int]
+        The clock frequency that the memory pool runs at in Hz.
+        If not specified/known, this will default to -1 indicating
+        it hasn't been defined.
+
+    read_bandwidth_bytes_per_cycle : Optional[int]
+        The read bandwidth of the memory pool in bytes/cycle.
+        If not specified/known, this will default to -1 indicating
+        it hasn't been defined.
+
+    write_bandwidth_bytes_per_cycle : Optional[int]
+        The write bandwidth of the memory pool in bytes/cycle.
+        If not specified/known, this will default to -1 indicating
+        it hasn't been defined.
+
+    read_latency_cycles : Optional[int]
+        The read latency of the memory pool in cycles.
+        If not specified/known, this will default to 0.
+
+    write_latency_cycles : Optional[int]
+        The write latency of the memory pool in cycles.
+        If not specified/known, this will default to 0.
+
+    target_burst_bytes : Optional[Union[Dict[Target, int], None]]
+        The burst length of the memory pool in bytes per target.
+        If not specified/known for a given target, a burst length
+        of 1 byte will be assumed.
+
+    """
+
+    def __init__(
+        self,
+        size_hint_bytes: Optional[int] = -1,
+        clock_frequency_hz: Optional[int] = -1,
+        read_bandwidth_bytes_per_cycle: Optional[int] = -1,
+        write_bandwidth_bytes_per_cycle: Optional[int] = -1,
+        read_latency_cycles: Optional[int] = 0,
+        write_latency_cycles: Optional[int] = 0,
+        target_burst_bytes=None,  # Optional[Union[Dict[target.Target, int], None]]
+    ):
+        if not target_burst_bytes:
+            target_burst_bytes = dict()
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.PoolInfoProperties,  # type: ignore # pylint: disable=no-member
+            size_hint_bytes,
+            clock_frequency_hz,
+            read_bandwidth_bytes_per_cycle,
+            write_bandwidth_bytes_per_cycle,
+            read_latency_cycles,
+            write_latency_cycles,
+            target_burst_bytes,
+        )
+
+
+@register_object("ir.WorkspacePoolInfo")
+class WorkspacePoolInfo(PoolInfo):
+    """WorkspacePoolInfo object holds information related to RW memory pools
+    where the statically sized allocate nodes will pooled into.
+
+    Parameters
+    ----------
+    pool_name : str
+        The name of the memory pool
+
+    targets : list[Target]
+        A list of targets which could access the pool
+
+    pool_info_properties : PoolInfoProperties
+        The properties of the pool.
+    """
+
+    # pylint: disable=W0231
+    def __init__(
+        self,
+        pool_name: str,
+        targets,  # Dict[Target, str]
+        pool_info_properties=PoolInfoProperties(),
+    ):
+        self.__init_handle_by_constructor__(
+            _ffi_api.WorkspacePoolInfo,  # type: ignore # pylint: disable=no-member
+            pool_name,
+            targets,
+            pool_info_properties,
+        )
+
+
+@register_object("ir.ConstantPoolInfo")
+class ConstantPoolInfo(PoolInfo):
+    """ConstantPoolInfo object holds information related to RO memory pools
+    where the statically sized allocate nodes will pooled into.
+
+    Parameters
+    ----------
+    pool_name : str
+        The name of the memory pool
+
+    targets : list[Target]
+        describes which targets could access the pool
+
+    pool_info_properties : PoolInfoProperties
+        The properties of the pool.
+    """
+
+    # pylint: disable=W0231

Review comment:
       Why are disabling this linter warning ? (i.e. could we not call base class init here ? )

##########
File path: 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

Review comment:
       We need more description here to explain why the byte offset is not sufficient, because naturally if the byte offset is calculated respecting the byte alignment, byte_alignment feels redundant.

##########
File path: src/tir/usmp/unified_static_memory_planner.cc
##########
@@ -54,14 +55,62 @@ IRModule PlanMemory(const IRModule& mod, String algo) {
   PrimFunc main_func = Downcast<PrimFunc>(mod->Lookup(::tvm::runtime::symbol::tvm_module_main));
   BufferInfoAnalysis buffer_info_analysis = ExtractBufferInfo(main_func, mod);
   Array<BufferInfo> buffer_info_arr =
-      CreateArrayBufferInfo(buffer_info_analysis->buffer_info_stmts);
+      ConvertToArrayOfBufferInfo(buffer_info_analysis->buffer_info_stmts);
+
+  // All items in RO pool should have conflicts with each other in this RO pool

Review comment:
       This part should not be here.
   
   This needs to part of ExtractBufferInfo pass.
   I would prefer if we can make it part of ExtractBufferInfo pass because that analysis pass is reponsible in producing the BufferInfo nodes with correct conflicts rather than doing a correction later like this.
   
   Please keep this part of code for integration of sub functionality of USMP.

##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -802,16 +797,15 @@ class AOTExecutorCodegen : public MixedModeVisitor {
       }
     }
     function_metadata_.Set(runtime::symbol::tvm_module_main, main_func_info);
+    LOG(INFO) << "LOWERED MOD ATTRS " << PrettyPrint(lowered_mod->attrs) << "\n";

Review comment:
       remove

##########
File path: src/target/source/source_module.cc
##########
@@ -236,20 +241,72 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
     return reference_arg + "_tvm_value";
   }
 
-  void GenerateInternalWorkspaceBuffers() {
+  void GenerateInternalBuffers() {
     if (metadata_->pool_inputs.defined()) {
       for (const auto& kv : metadata_->pool_inputs.value()) {
         tir::usmp::AllocatedPoolInfo allocated_pool_info = kv.second;
         if (allocated_pool_info->pool_info->is_internal) {
-          code_ << "__attribute__((section(\".data.tvm\"), ";
-          code_ << "aligned(" << 16 << ")))\n";
-          code_ << "static uint8_t " << allocated_pool_info->pool_info->pool_name << "["
-                << allocated_pool_info->allocated_size->value << "];\n";
+          if (const auto* pool_info = allocated_pool_info->pool_info.as<ConstantPoolInfoNode>()) {
+            GenerateConstantBuffer(pool_info, allocated_pool_info->allocated_size->value);
+          } else {
+            GenerateWorkspaceBuffer(allocated_pool_info->pool_info.as<WorkspacePoolInfoNode>(),
+                                    allocated_pool_info->allocated_size->value);
+          }
         }
       }
     }
   }
 
+  void GenerateConstantBuffer(const ConstantPoolInfoNode* pool_info, size_t allocated_size) {
+    size_t offset = 0;
+    if (pool_info->constant_info_array.size() > 0) {
+      // Pool is RO, form an initialized struct
+      code_ << "__attribute__((section(\".rodata.tvm\"), ";
+      code_ << "))\n";
+      code_ << "static struct " << pool_info->pool_name << " {\n";
+      // emit struct field names
+      std::vector<ConstantInfo> const_info_vec(pool_info->constant_info_array.begin(),
+                                               pool_info->constant_info_array.end());
+      std::sort(const_info_vec.begin(), const_info_vec.end(),
+                [](const ConstantInfo& a, const ConstantInfo& b) {
+                  return a->byte_offset->value < b->byte_offset->value;
+                });
+      for (const auto& const_info : const_info_vec) {
+        const auto& data = const_info->data;
+        const auto& offs = const_info->byte_offset;
+        int64_t num_elements = std::accumulate(data.Shape().begin(), data.Shape().end(), 1,
+                                               std::multiplies<int64_t>());
+        code_ << "  ";
+        codegen_c_base_.PrintType(data.DataType(), code_);
+        code_ << " " << const_info->name_hint << "[" << num_elements
+              << "] __attribute__((packed, aligned(" << metadata_->workspace_byte_alignment

Review comment:
       Should not this be constant byte alignment ?

##########
File path: 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

Review comment:
       On having read the PR, I dont see where this is used. Can you point me if and where this is used ?

##########
File path: src/tir/usmp/transform/assign_pool_info.cc
##########
@@ -48,18 +48,26 @@ class PoolInfoAssigner : public StmtExprMutator {
     ICHECK(target_host) << "main function does not have a target attr";
     WorkspaceMemoryPools workspace_pools =
         module->GetAttr<WorkspaceMemoryPools>(tvm::attr::kWorkspaceMemoryPools)
-            .value_or(WorkspaceMemoryPools({CreateDefaultMemoryPool(module)}));
+            .value_or(WorkspaceMemoryPools(

Review comment:
       ConstantPoolInfo should not be a part of WorskpaceMemoryPools

##########
File path: include/tvm/ir/memory_pools.h
##########
@@ -130,15 +137,159 @@ static const int kUnknownReadBandwidth = -1;
 static const int kUnknownWriteBandwidth = -1;
 
 class PoolInfo : public ObjectRef {
- public:
+ protected:
   TVM_DLL PoolInfo(String pool_name, Map<Target, String> target_access,
                    Integer size_hint_bytes = kUnrestrictedPoolSizeHint,
                    Integer clock_frequency_hz = kUnknownClockFrequency,
                    Integer read_bandwidth_bytes_per_cycle = kUnknownReadBandwidth,
                    Integer write_bandwidth_bytes_per_cycle = kUnknownWriteBandwidth,
                    Integer read_latency_cycles = 0, Integer write_latency_cycles = 0,
                    Map<Target, Integer> target_burst_bytes = {}, Bool is_internal = Bool(false));
-  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
+
+ public:
+  // TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
+  TVM_DEFINE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
+};
+
+/*!
+ * \brief Describes a pool of memory properties
+ */
+struct PoolInfoPropertiesNode : public Object {
+  /*! \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 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;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("size_hint_bytes", &size_hint_bytes);
+    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);
+  }
+
+  bool SEqualReduce(const PoolInfoPropertiesNode* other, SEqualReducer equal) const {
+    return equal(size_hint_bytes, other->size_hint_bytes) &&
+           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);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce(size_hint_bytes);
+    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);
+  }
+
+  static constexpr const char* _type_key = "ir.PoolInfoProperties";
+  TVM_DECLARE_FINAL_OBJECT_INFO(PoolInfoPropertiesNode, Object);
+};
+
+class PoolInfoProperties : public ObjectRef {
+ public:
+  TVM_DLL PoolInfoProperties(Integer size_hint_bytes = kUnrestrictedPoolSizeHint,
+                             Integer clock_frequency_hz = kUnknownClockFrequency,
+                             Integer read_bandwidth_bytes_per_cycle = kUnknownReadBandwidth,
+                             Integer write_bandwidth_bytes_per_cycle = kUnknownWriteBandwidth,
+                             Integer read_latency_cycles = 0, Integer write_latency_cycles = 0,
+                             Map<Target, Integer> target_burst_bytes = {},
+                             Bool is_internal = Bool(false));
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfoProperties, ObjectRef, PoolInfoPropertiesNode);
+};
+
+struct WorkspacePoolInfoNode : public PoolInfoNode {
+  static constexpr const char* _type_key = "ir.WorkspacePoolInfo";
+  TVM_DECLARE_FINAL_OBJECT_INFO(WorkspacePoolInfoNode, PoolInfoNode);
+};
+
+class WorkspacePoolInfo : public PoolInfo {
+ public:
+  TVM_DLL WorkspacePoolInfo(
+      String pool_name, Array<Target> targets,
+      PoolInfoProperties properties = PoolInfoProperties(kUnrestrictedPoolSizeHint));
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(WorkspacePoolInfo, PoolInfo, WorkspacePoolInfoNode);
+};
+
+/*
+ * \brief The ConstantInfoNode contains numeric literal in RO pool
+ * Used in initialisation of RO memory
+ */
+struct ConstantInfoNode : public Object {
+  String name_hint;
+  Integer byte_offset;
+  runtime::NDArray data;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("name_hint", &name_hint);
+    v->Visit("byte_offset", &byte_offset);
+    v->Visit("data", &data);
+  }
+
+  bool SEqualReduce(const ConstantInfoNode* other, SEqualReducer equal) const {
+    return equal(name_hint, other->name_hint) && equal(byte_offset, other->byte_offset) &&
+           equal(data, other->data);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce(name_hint);
+    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_offset, runtime::NDArray data);
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ConstantInfo, ObjectRef, ConstantInfoNode);
+};
+
+struct ConstantPoolInfoNode : public PoolInfoNode {
+  Array<ConstantInfo> constant_info_array;
+  static constexpr const char* _type_key = "ir.ConstantPoolInfo";
+  TVM_DECLARE_FINAL_OBJECT_INFO(ConstantPoolInfoNode, PoolInfoNode);
+};
+
+class ConstantPoolInfo : public PoolInfo {
+ public:
+  TVM_DLL ConstantPoolInfo(
+      String pool_name, Array<Target> targets, Array<ConstantInfo> constant_info_array,
+      PoolInfoProperties properties = PoolInfoProperties(kUnrestrictedPoolSizeHint));
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ConstantPoolInfo, PoolInfo, ConstantPoolInfoNode);
 };
 
 struct WorkspaceMemoryPoolsNode : public Object {

Review comment:
       We need something similiar for ConstantMemoryPools.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org