diff --git a/python/tvm/relay/backend/contrib/ethosu/tir_to_cs_translator.py b/python/tvm/relay/backend/contrib/ethosu/tir_to_cs_translator.py index 33bde5a5ea9b..c53c32801db2 100644 --- a/python/tvm/relay/backend/contrib/ethosu/tir_to_cs_translator.py +++ b/python/tvm/relay/backend/contrib/ethosu/tir_to_cs_translator.py @@ -81,6 +81,12 @@ class RegionOffset(NamedTuple): def analyze_scratch_memory_acesses(mod: tvm.IRModule, candidate_regions_for_scratch: List[int]): """ + This function analyzes the IRModule for intermediary tensors that can be resulting + from a offset of pool variables (via Let nodes) and/or allocate nodes. The allocate + nodes will be folded into a single TVMBackendallocWorkspace call with offsets. Ultimately + this will produce a mapping from each such node to a RegionOffset named tuple that + has the region and the obtained offset, as mentioned above. + Parameters ---------- mod: tvm.IRModule @@ -90,7 +96,7 @@ def analyze_scratch_memory_acesses(mod: tvm.IRModule, candidate_regions_for_scra Returns ------- - scratch_region_map : Dict[tvm.tir.Var, int] + scratch_region_map : Dict[tvm.tir.Var, RegionOffset] A map between buffer vars to scratch regions they are assigned tvm_backend_alloc_workspace_size : int The size of tvm_backend_alloc_workspace call required to service @@ -195,6 +201,11 @@ def translate(tir_module, params): base_addresses : List[util.BaseAddress] base addresses to be used by the driver """ + + # The NPU has 6 usable regions ranging from 0-6 + # The regions 0, 3, and 4 is already used for input, + # output and constant, respectively (See _get_regions()). + # Thus, for scratch we are left with 5, 2 and 1. candidate_regions_for_scratch = [5, 2, 1] ( scratch_region_map, diff --git a/src/tir/usmp/transform/assign_pool_info.cc b/src/tir/usmp/transform/assign_pool_info.cc index 9df4cdda31db..9d8e36137c37 100644 --- a/src/tir/usmp/transform/assign_pool_info.cc +++ b/src/tir/usmp/transform/assign_pool_info.cc @@ -84,11 +84,11 @@ PoolInfo PoolInfoAssigner::CreateDefaultMemoryPool(const tvm::IRModule& module) for (const auto& kv : module->functions) { BaseFunc func = kv.second; Optional target = func->GetAttr(tvm::attr::kTarget); - target_access.Set(target.value_or(target_host), PoolInfo::kTargetPoolReadWriteAccess); + target_access.Set(target.value_or(target_host), kTargetPoolReadWriteAccess); } return PoolInfo("global_workspace", target_access, kUnrestrictedPoolSizeHint, - kUnknownClockFrequency, kUnknownReadBandwidth, - kUnknownWriteBandwidth, 0, 0, {}, Bool(true)); + kUnknownClockFrequency, kUnknownReadBandwidth, kUnknownWriteBandwidth, 0, 0, {}, + Bool(true)); } Stmt PoolInfoAssigner::VisitStmt_(const AllocateNode* op) {