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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
