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 6a2fcd3678ef8..6e8efa579367e 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 @@ -424,14 +424,16 @@ def replace_npu_address_range_with_address(npu_addr_range): assert isinstance(npu_addr_range.address, tvm.tir.Load) buffer = npu_addr_range.address.buffer_var if buffer in scratch_region_map.keys(): + # Check whether the Load itself has an offset + tensor_idx = int(npu_addr_range.address.index.value) return vapi.NpuAddressRange( scratch_region_map[buffer].region, - scratch_region_map[buffer].offset, + scratch_region_map[buffer].offset + tensor_idx, npu_addr_range.length, ) assert buffer in buffer_addresses.keys(), f"searching for buffer : {buffer}, but not found" address, buffer_type = buffer_addresses[buffer] - address = address + npu_addr_range.address.index.value + address = address + int(npu_addr_range.address.index.value) return vapi.NpuAddressRange(_get_region(buffer_type), address, npu_addr_range.length) def replace_tir_loads(npu_object):