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/06/06 13:59:19 UTC

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

manupa-arm commented on code in PR #10189:
URL: https://github.com/apache/tvm/pull/10189#discussion_r885793548


##########
src/tir/usmp/transform/assign_pool_info.cc:
##########
@@ -49,18 +49,35 @@ class PoolInfoAssigner : public StmtExprMutator {
     WorkspaceMemoryPools workspace_pools =
         module->GetAttr<WorkspaceMemoryPools>(tvm::attr::kWorkspaceMemoryPools)
             .value_or(WorkspaceMemoryPools({CreateDefaultMemoryPool(module)}));
-    Array<PoolInfo> pool_infos = workspace_pools->pools;
-    for (const PoolInfo& pool_info : pool_infos) {
-      for (const auto& kv : pool_info->target_access) {
-        Target target = kv.first;
-        String target_str = target->str();
-        if (target_pool_infos_.find(target_str) == target_pool_infos_.end()) {
-          target_pool_infos_.Set(target_str, Array<PoolInfo>());
+    // make default ConstantPoolInfo if no constant and no workspace pool infos supplied
+    ConstantMemoryPools constant_pools =
+        module->GetAttr<ConstantMemoryPools>(tvm::attr::kConstantMemoryPools)
+            .value_or(
+                module->GetAttr<WorkspaceMemoryPools>(tvm::attr::kWorkspaceMemoryPools).defined()
+                    ? ConstantMemoryPools()
+                    : ConstantMemoryPools({ConstantPoolInfo(

Review Comment:
   nit : can we factor this out to a seperate function to improve readability. e.g. CreateDefaultConstantMemoryPool



##########
src/runtime/aot_executor/aot_executor.cc:
##########
@@ -62,9 +65,38 @@ AotExecutor::AotExecutor(tvm::runtime::Module module, const std::vector<Device>&
                                       output->dtype(), devices_[0]));
   }
 
-  for (auto pool : metadata_->pools()) {
-    args_.emplace_back(NDArray::Empty(ShapeTuple(pool->shape().begin(), pool->shape().end()),
-                                      pool->dtype(), devices_[0]));
+  auto get_array_byte_size = [](DataType data_type, auto shape) {

Review Comment:
   Any reason why this was chosen to be a lambda ?
   (Feels like a good candidate for common utility)



##########
tests/cpp/aot_metadata_test.cc:
##########
@@ -39,9 +40,19 @@ const struct TVMTensorInfo kNormalOutputs[1] = {
 
 const int64_t kNormalPool1Shape[3] = {3, 8, 8};
 const struct TVMTensorInfo kNormalPools[1] = {{"pool1", kNormalPool1Shape, 3, DLDataType{3, 4, 7}}};
+const struct TVMConstantInfo kNormalConsts[1] = {{"consts1", 0, 0, {}}};
 
 const struct TVMMetadata kNormal = {
-    TVM_METADATA_VERSION, kNormalInputs, 2, kNormalOutputs, 1, kNormalPools, 1, "default",
+    TVM_METADATA_VERSION,
+    kNormalInputs,
+    2,
+    kNormalOutputs,
+    1,
+    kNormalPools,

Review Comment:
   kNormalWorkspacePools might be better here.



##########
tests/cpp/aot_metadata_test.cc:
##########
@@ -324,7 +350,8 @@ TEST(DiscoverArraysVisitor, DiscoverArrays) {
                                    DiscoveredNameEq("kTvmgenMetadata_outputs_0_shape"),
                                    DiscoveredNameEq("kTvmgenMetadata_outputs"),
                                    DiscoveredNameEq("kTvmgenMetadata_pools_0_shape"),
-                                   DiscoveredNameEq("kTvmgenMetadata_pools")}));
+                                   DiscoveredNameEq("kTvmgenMetadata_pools"),

Review Comment:
   lets us use pools --> workspace_pools + constant_pools throughout.



##########
src/relay/backend/executor.cc:
##########
@@ -91,7 +91,8 @@ TVM_REGISTER_EXECUTOR("aot")
     .add_attr_option<Bool>("link-params", Bool(true))
     .add_attr_option<Bool>("unpacked-api")
     .add_attr_option<String>("interface-api")
-    .add_attr_option<Integer>("workspace-byte-alignment");
+    .add_attr_option<Integer>("workspace-alignment")

Review Comment:
   agreed. we should have the unit.



##########
src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc:
##########
@@ -53,14 +54,20 @@ class PoolAllocationToOffsetConverter : public StmtExprMutator {
       : pool_allocations_(pool_allocations), emit_tvmscript_printable_(emit_tvmscript_printable) {
     module_ = module->ShallowCopy();
     for (const auto& kv : pool_allocations) {
-      // TODO(@manupa-arm): add AllocateConstNode when it is available
-      ICHECK(kv.first->IsInstance<AllocateNode>());
-      Allocate allocate_node = Downcast<Allocate>(kv.first);
+      size_t extent_size = -1;
+      if (kv.first->IsInstance<AllocateNode>()) {
+        Allocate allocate_node = Downcast<Allocate>(kv.first);
+        extent_size = CalculateExtentsSize(allocate_node.operator->());
+      } else if (kv.first->IsInstance<AllocateConstNode>()) {
+        AllocateConst allocate_const_node = Downcast<AllocateConst>(kv.first);
+        extent_size = CalculateExtentsSize(allocate_const_node.operator->());
+      } else {
+        LOG(FATAL) << "Not supported";

Review Comment:
   We need this addressed



##########
tests/cpp/aot_metadata_test.cc:
##########
@@ -39,9 +40,19 @@ const struct TVMTensorInfo kNormalOutputs[1] = {
 
 const int64_t kNormalPool1Shape[3] = {3, 8, 8};
 const struct TVMTensorInfo kNormalPools[1] = {{"pool1", kNormalPool1Shape, 3, DLDataType{3, 4, 7}}};
+const struct TVMConstantInfo kNormalConsts[1] = {{"consts1", 0, 0, {}}};
 
 const struct TVMMetadata kNormal = {
-    TVM_METADATA_VERSION, kNormalInputs, 2, kNormalOutputs, 1, kNormalPools, 1, "default",
+    TVM_METADATA_VERSION,
+    kNormalInputs,
+    2,
+    kNormalOutputs,
+    1,
+    kNormalPools,
+    1,
+    kNormalConsts,

Review Comment:
   kNormalConstantPools might be better here.



-- 
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