code: def test_assign_addresses():
      test_cases = [
          {
              # Stimulus
              "tir_module": WeightStreamOnly,
              "param_dict": {
                  WeightStreamOnly["main"].attrs["constants"][
                      "buffer"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [128],
                      "uint8",
                  ),
                  WeightStreamOnly["main"].attrs["constants"][
                      "buffer_1"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [32],
                      "uint8",
                  ),
                  WeightStreamOnly["main"].attrs["constants"][
                      "buffer_2"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [112],
                      "uint8",
                  ),
                  WeightStreamOnly["main"].attrs["constants"][
                      "buffer_3"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [32],
                      "uint8",
                  ),
                  WeightStreamOnly["main"].attrs["constants"][
                      "buffer_4"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [112],
                      "uint8",
                  ),
                  WeightStreamOnly["main"].attrs["constants"][
                      "buffer_5"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [32],
                      "uint8",
                  ),
                  WeightStreamOnly["main"].attrs["constants"][
                      "buffer_6"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [112],
                      "uint8",
                  ),
                  WeightStreamOnly["main"].attrs["constants"][
                      "buffer_7"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [32],
                      "uint8",
                  ),
              },
          },
          {
              # Stimulus
              "tir_module": MixedRead,
              "param_dict": {
                  MixedRead["main"].attrs["constants"][
                      "buffer"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [592],
                      "uint8",
                  ),
                  MixedRead["main"].attrs["constants"][
                      "buffer_1"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [160],
                      "uint8",
                  ),
                  MixedRead["main"].attrs["constants"][
                      "buffer_2"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [80],
                      "uint8",
                  ),
                  MixedRead["main"].attrs["constants"][
                      "buffer_3"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [32],
                      "uint8",
                  ),
                  MixedRead["main"].attrs["constants"][
                      "buffer_4"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [80],
                      "uint8",
                  ),
                  MixedRead["main"].attrs["constants"][
                      "buffer_5"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [32],
                      "uint8",
                  ),
                  MixedRead["main"].attrs["constants"][
                      "buffer_6"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [80],
                      "uint8",
                  ),
                  MixedRead["main"].attrs["constants"][
                      "buffer_7"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [32],
                      "uint8",
                  ),
                  MixedRead["main"].attrs["constants"][
                      "buffer_8"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [80],
                      "uint8",
                  ),
                  MixedRead["main"].attrs["constants"][
                      "buffer_9"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [32],
                      "uint8",
                  ),
              },
          },
          {
              # Stimulus
              "tir_module": MixedConstantDatatypes,
              "param_dict": {
                  MixedConstantDatatypes["main"].attrs["constants"][
                      "buffer"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [160],
                      "uint8",
                  ),
                  MixedConstantDatatypes["main"].attrs["constants"][
                      "buffer_2"
                  ]: np.random.randint(
                      np.iinfo("int16").min,
                      np.iinfo("int16").max,
                      [1],
                      "int16",
                  ),
                  MixedConstantDatatypes["main"].attrs["constants"][
                      "buffer_1"
                  ]: np.random.randint(
                      np.iinfo("uint8").min,
                      np.iinfo("uint8").max,
                      [272],
                      "uint8",
                  ),
              },
          },
      ]
  
      def extract_call_extern_list(mod):
          """This function will obtain all ethosu_conv2d
          calls from a NPU TIR module
          Parameters
          ----------
          mod : tvm.IRModule
              This is a NPU TIR Module
  
          Returns
          -------
          list
              of tvm.tir.Call objects
              that are tir extern calls
              for ethosu_conv2d
          """
          # There should only be a single function
          assert len(mod.functions.items()) == 1
          primfunc = mod.functions.items()[0][1]
  
          extern_calls = list()
  
          def populate_extern_calls(stmt):
              if (
                  isinstance(stmt, tvm.tir.Call)
                  and stmt.op.name == "tir.call_extern"
              ):
                  extern_calls.append(stmt)
  
          stmt_functor.post_order_visit(
              primfunc.body, populate_extern_calls
          )
          return extern_calls
  
      def collect_tir_buffer_info(npu_ops):
          """This is run prior to address assigning to collect tir buffer information
          for verification later on"""
          _npu_op_tir_buffers = dict()
          for npu_op in npu_ops:
              if isinstance(npu_op, vapi.NpuDmaOperation):
                  _npu_op_tir_buffers[npu_op] = (
                      npu_op.src.address,
                      npu_op.dest.address,
                  )
              elif issubclass(type(npu_op), vapi.NpuBlockOperation):
                  _npu_op_tir_buffers[npu_op] = (
                      npu_op.ifm.tiles.addresses[0],
                      npu_op.ofm.tiles.addresses[0],
                      npu_op.weights,
                      npu_op.biases,
                  )
          return _npu_op_tir_buffers
  
      def _check_buffer(address, region, length, buffer_var):
          """Checks whether the buffer information is valid with
          original tir buffers.
          - If its constant, this will check
            the slice in the constant tensor has the values.
          - If its scratch, this will check
            the slice is within scratch and does not have conflicts
            with other scratch tensors.
          - If its input/output, this will check the
            address is zero
          """
          inverse_region_map = {
              0: tir_to_cs_translator.BufferType.constant,
              1: tir_to_cs_translator.BufferType.scratch,
              3: tir_to_cs_translator.BufferType.input,
              4: tir_to_cs_translator.BufferType.output,
          }
          buffer_type = inverse_region_map[region]
          buffer_dtype = buffer_var.type_annotation.element_type.dtype
          dtype_bytes = np.iinfo(np.dtype(buffer_dtype)).bits // 8
          if buffer_type == tir_to_cs_translator.BufferType.constant:
              ref = buffer_info[buffer_var].values
              hex_from = address * dtype_bytes * 2
              hex_to = hex_from + length * dtype_bytes * 2
              constant_hex = constant_hex_string[hex_from:hex_to]
              constant_tensor = np.frombuffer(
                  bytearray.fromhex(constant_hex), dtype=buffer_dtype
              )
              np.array_equal(constant_tensor, ref)
              # Every buffer is adjusted to align to 16 bytes
              length = util.round_up(length, 16)
              # Mark these constants are read at least once
              constant_tensor_read_mask[
                  address : address + length
              ] = np.ones(length, dtype=buffer_dtype)
          elif buffer_type == tir_to_cs_translator.BufferType.scratch:
              assert address < tvmbaw_workspace_size
  
              size_in_bytes = allocate_node_sizes[buffer_var]
              # Every buffer is adjusted to align to 16 bytes
              size_in_bytes = util.round_up(size_in_bytes, 16)
              assert address + size_in_bytes <= tvmbaw_workspace_size
              # The scratch area should not be used by any other buffer
              assert not tvmbaw_workspace_mask[
                  address : address + size_in_bytes
              ].any()
              # The scratch area is marked as used
              tvmbaw_workspace_mask[
                  address : address + size_in_bytes
              ] = np.ones(size_in_bytes, dtype="uint8")
          elif buffer_type == tir_to_cs_translator.BufferType.input:
              assert address == 0
          else:
              assert (
                  buffer_type == tir_to_cs_translator.BufferType.output
              )
              assert address == 0
  
      def _get_allocate_node_sizes(mod):
          # There should only be a single function
          assert len(mod.functions.items()) == 1
          primfunc = mod.functions.items()[0][1]
          _allocate_node_sizes = dict()
  
          def analyze_remaining_allocates(stmt):
              if isinstance(stmt, tvm.tir.stmt.Allocate):
                  allocate = stmt
                  pointer_type = allocate.buffer_var.type_annotation
                  storage_scope = pointer_type.storage_scope
                  if storage_scope == "global":
                      dtype_bytes = (
                          np.iinfo(np.dtype(allocate.dtype)).bits // 8
                      )
                      size_in_bytes = int(
                          dtype_bytes * np.prod(list(allocate.extents))
                      )
                      # Every memory address the NPU access have to be 16 byte aligned
                      size_in_bytes = util.round_up(size_in_bytes, 16)
                      _allocate_node_sizes[
                          allocate.buffer_var
                      ] = size_in_bytes
  
          tvm.tir.stmt_functor.post_order_visit(
              primfunc.body, analyze_remaining_allocates
          )
          return _allocate_node_sizes
  
      def verify(npu_ops):
          """This wrapper verifies the allocated addresses matches with original tir buffers"""
          checked_buffers = set()
  
          def check_buffer(address, region, length, buffer_var):
              if buffer_var not in checked_buffers:
                  _check_buffer(address, region, length, buffer_var)
                  checked_buffers.add(buffer_var)
  
          for npu_op in npu_ops:
              if isinstance(npu_op, vapi.NpuDmaOperation):
                  src_tir_buffer_var = npu_op_tir_buffers[npu_op][
                      0
                  ].buffer.data
                  check_buffer(
                      npu_op.src.address,
                      npu_op.src.region,
                      npu_op.src.length,
                      src_tir_buffer_var,
                  )
                  dest_tir_load = npu_op_tir_buffers[npu_op][
                      1
                  ].buffer.data
                  check_buffer(
                      npu_op.dest.address,
                      npu_op.dest.region,
                      npu_op.dest.length,
                      dest_tir_load,
                  )
              elif issubclass(type(npu_op), vapi.NpuBlockOperation):
                  ifm_tir_buffer_var = npu_op_tir_buffers[npu_op][
                      0
                  ].buffer.data
                  ifm_length = (
                      npu_op.ifm.shape.height
                      * npu_op.ifm.shape.width
                      * npu_op.ifm.shape.depth
                  )
                  check_buffer(
                      npu_op.ifm.tiles.addresses[0],
                      npu_op.ifm.region,
                      ifm_length,
                      ifm_tir_buffer_var,
                  )
                  ofm_tir_buffer_var = npu_op_tir_buffers[npu_op][
                      1
                  ].buffer.data
                  ofm_length = (
                      npu_op.ofm.shape.height
                      * npu_op.ofm.shape.width
                      * npu_op.ofm.shape.depth
                  )
                  check_buffer(
                      npu_op.ofm.tiles.addresses[0],
                      npu_op.ofm.region,
                      ofm_length,
                      ofm_tir_buffer_var,
                  )
                  for idx, weight in enumerate(
                      npu_op_tir_buffers[npu_op][2]
                  ):
                      assert isinstance(weight, vapi.NpuAddressRange)
                      check_buffer(
                          npu_op.weights[idx].address,
                          npu_op.weights[idx].region,
                          npu_op.weights[idx].length,
                          weight.address.buffer.data,
                      )
                  for idx, bias in enumerate(
                      npu_op_tir_buffers[npu_op][3]
                  ):
                      assert isinstance(bias, vapi.NpuAddressRange)
                      check_buffer(
                          npu_op.biases[idx].address,
                          npu_op.biases[idx].region,
                          npu_op.biases[idx].length,
                          bias.address.buffer.data,
                      )
  
      for test_case in test_cases:
          tir_mod = test_case["tir_module"]
          tir_mod["main"] = tir_mod["main"].with_attr(
              "target", tvm.target.Target("ethos-u", host="ethos-u")
          )
          tir_mod = tvm.tir.transform.MakeUnpackedAPI()(tir_mod)
          candidate_regions_for_scratch = [5, 2, 1]
          (
              scratch_region_map,
              tvmbaw_workspace_size,
              _,
          ) = tir_to_cs_translator.analyze_scratch_memory_acesses(
              tir_mod, candidate_regions_for_scratch
          )
          allocate_node_sizes = _get_allocate_node_sizes(tir_mod)
          buffer_info = tir_to_cs_translator.extract_buffer_info(
              tir_mod, test_case["param_dict"]
          )
          extern_calls = extract_call_extern_list(tir_mod)
          _npu_ops = list()
          for extern_call in extern_calls:
              _npu_ops.append(
                  tir_to_cs_translator.translate_ethosu_tir_call_extern(
                      extern_call
                  )
              )
          npu_op_tir_buffers = collect_tir_buffer_info(_npu_ops)
          (
              _npu_ops,
              constant_hex_string,
          ) = tir_to_cs_translator.assign_addresses(
              buffer_info, _npu_ops, scratch_region_map
          )
          tvmbaw_workspace_mask = np.zeros(
              tvmbaw_workspace_size, dtype="uint8"
          )
          constant_tensor_read_mask = np.zeros(
              len(constant_hex_string) // 2, dtype="uint8"
          )
          verify(_npu_ops)
          # This will be only 1 if all allocated scratch is used.
          assert np.prod(tvmbaw_workspace_mask) == 1
          # This will be only 1 if all constant tensors is read at least once.
          assert np.prod(constant_tensor_read_mask) == 1
  
imports: import pytest
  import numpy
  import tvm
  from tvm.tir import stmt_functor
  from tvm.script import tir
  from tvm.relay.backend.contrib.ethosu import tir_to_cs_translator
  from tvm.relay.backend.contrib.ethosu import util
  import ethosu.vela.api
label_desc: Use image node_id 5 for calling a global function with example usage: test_assign_addresses() and returns: extern_calls, _npu_op_tir_buffers, _allocate_node_sizes