@@ -129,15 +129,15 @@ def analyze_pool_access(stmt):
129129
130130 tvm .tir .stmt_functor .post_order_visit (primfunc .body , analyze_pool_access )
131131
132- tvmbaw_region = None
132+ dynamic_allocation_region = None
133133 if len (candidate_regions_for_scratch ) > 0 :
134- tvmbaw_region = candidate_regions_for_scratch .pop ()
135- tvmbaw_size = 0
134+ dynamic_allocation_region = candidate_regions_for_scratch .pop ()
135+ dynamic_allocation_size = 0
136136
137137 # If there are tir.Allocate remaining by now, they need to be serviced via
138- # TVMBAW calls.
138+ # dynamic_allocation calls.
139139 def analyze_remaining_allocates (stmt ):
140- nonlocal tvmbaw_size
140+ nonlocal dynamic_allocation_size
141141 if isinstance (stmt , tvm .tir .stmt .Allocate ):
142142 allocate = stmt
143143 pointer_type = allocate .buffer_var .type_annotation
@@ -147,18 +147,18 @@ def analyze_remaining_allocates(stmt):
147147 size_in_bytes = int (dtype_bytes * np .prod (list (allocate .extents )))
148148 # Every memory address the NPU access have to be 16 byte aligned
149149 size_in_bytes = util .round_up (size_in_bytes , 16 )
150- address = tvmbaw_size
151- tvmbaw_size += size_in_bytes
150+ address = dynamic_allocation_size
151+ dynamic_allocation_size += size_in_bytes
152152 scratch_region_map [allocate .buffer_var ] = RegionOffset (
153- region = tvmbaw_region , offset = address
153+ region = dynamic_allocation_region , offset = address
154154 )
155155
156156 tvm .tir .stmt_functor .post_order_visit (primfunc .body , analyze_remaining_allocates )
157157
158158 return (
159159 scratch_region_map ,
160- tvmbaw_size ,
161- tvmbaw_region ,
160+ dynamic_allocation_size ,
161+ dynamic_allocation_region ,
162162 )
163163
164164
@@ -209,8 +209,8 @@ def translate(tir_module, params):
209209 candidate_regions_for_scratch = [5 , 2 , 1 ]
210210 (
211211 scratch_region_map ,
212- tvmbaw_workspace_size ,
213- tvmbaw_region ,
212+ dynamic_allocation_size ,
213+ dynamic_allocation_region ,
214214 ) = analyze_scratch_memory_acesses (tir_module , candidate_regions_for_scratch )
215215 buffer_info = extract_buffer_info (tir_module , params )
216216 call_extern_list = extract_call_extern_list (tir_module )
@@ -219,13 +219,13 @@ def translate(tir_module, params):
219219 _npu_ops .append (translate_ethosu_tir_call_extern (call_extern ))
220220 _npu_ops , constant_data = assign_addresses (buffer_info , _npu_ops , scratch_region_map )
221221 base_addresses = extract_param_base_addresses (tir_module , buffer_info , scratch_region_map )
222- if tvmbaw_workspace_size :
222+ if dynamic_allocation_size :
223223 base_addresses .append (
224224 util .BaseAddress (
225- name = "tvmbaw " ,
225+ name = "dynamic_allocation " ,
226226 primfunc_param_idx = None ,
227- region = tvmbaw_region ,
228- size = tvmbaw_workspace_size ,
227+ region = dynamic_allocation_region ,
228+ size = dynamic_allocation_size ,
229229 is_runtime_allocation = True ,
230230 )
231231 )
0 commit comments