code: def EncodeConstants(const_dict):
      """the NPU requires that weights are compressed and bias/scales are 'encoded', both
      of which are performed by this pass.
  
      This pass modifies both the constant dict to contain the post-encoding values of the
      constants and the IR to adjust buffer types/sizes/accesses so they align with the
      encoded constants. Calls to the Vela API are made to perform the actual compression/
      encoding.
  
      """
      new_const_dict = {}
  
      def collect_encoding_definitions(stmt, old_buffer_var_to_const):
          # Map from copy destination to copy source.
          copy_map = {}
          # List of buffer copies that occurred
          copied_buffers = []
          # List of encoded buffer information
          constant_buffer_replacements = []
  
          def _align_scale_bias(tir_extern_call, bias):
              """Align the scale_bias to 16 bytes."""
              value_bytes = bytearray()
              value_bytes.extend(bias.tobytes())
              # Align to 16
              remainder = (len(value_bytes)) % 16
              if remainder > 0:
                  value_bytes.extend(bytearray(16 - remainder))
              value = np.frombuffer(value_bytes, dtype="uint8")
              return value
  
          accel_config = vela_api.get_accelerator_config()
  
          def _encode_weights(tir_extern_call, weights):
              """Encode the weights for a TIR extern call."""
              value_bytes = vela_api.encode_weights(
                  tir_extern_call, weights, accel_config
              )
              value = np.frombuffer(value_bytes, dtype="uint8")
              return value
  
          def _declare_constant_buffer(
              old_buffer, encoded_constants, split_idx
          ):
              """Create a new buffer and add the old buffer and its pointer to the
              rewriting maps."""
              new_buffer = tvm.tir.decl_buffer(
                  shape=[len(encoded_constants)],
                  dtype=str(encoded_constants.dtype),
                  name=old_buffer.name + "_encoded",
                  scope=old_buffer.scope(),
              )
  
              constant_buffer_replacements.append(
                  {
                      "old_buffer": old_buffer,
                      "new_buffer": new_buffer,
                      "encoded_constants": encoded_constants,
                      "split_idx": split_idx,
                  }
              )
  
          def _encode_weights_or_bias(
              buffer1, buffer2, stmt, encode_func
          ):
              """Encode the weights or align the bias either for one or two cores,
              depending on the variant."""
              constant = old_buffer_var_to_const[buffer1.data]
  
              # If we have just one core, encode the whole constant
              if buffer2 is None:
                  new_const = encode_func(stmt, constant)
                  return new_const, None
  
              # Assume that the constant tensor has not been flattened yet
              assert len(constant.shape) != 1
              channels = constant.shape[0]
              split_const = np.split(constant, channels, axis=0)
  
              const_list = [
                  split_const[i] for i in range(channels) if i % 2 == 0
              ]
              const_to_encode = np.concatenate(const_list, axis=0)
  
              new_const = encode_func(stmt, const_to_encode)
              split_idx = len(new_const)
  
              # Encode half of the constant separately for the other core if it exists
              assert buffer1.same_as(buffer2)
              const2_list = [
                  split_const[i] for i in range(channels) if i % 2 == 1
              ]
              const2_to_encode = np.concatenate(const2_list, axis=0)
  
              new_const2 = encode_func(stmt, const2_to_encode)
              new_const = np.append(new_const, new_const2).astype(
                  "uint8"
              )
  
              return new_const, split_idx
  
          def _visit(stmt):
              if isinstance(stmt, tvm.tir.Call):
                  op = str(stmt.args[0].value)
                  # Handle copies as a special-case by propagating the buffer information
                  # from the read to the write pointer.
                  if op == "ethosu_copy":
                      read_buffer = stmt.args[1].buffer
                      write_buffer = stmt.args[3].buffer
                      # Assert writing to the base of the write_var (pre-StorageRewrite)
                      assert list(stmt.args[3].indices) == [0]
                      assert list(stmt.args[1].indices) == [0]
                      copied_buffers.append(
                          {"source": read_buffer, "dest": write_buffer}
                      )
                      copy_map[write_buffer] = read_buffer
  
                  ops_with_weights = {
                      "ethosu_conv2d": tirtocs.translate_ethosu_conv2d,
                      "ethosu_depthwise_conv2d": tirtocs.translate_ethosu_depthwise_conv2d,
                  }
                  if op in ops_with_weights:
                      npu_op, _ = ops_with_weights[op](stmt)
  
                      # Encode the weights
                      weights_buffer = npu_op.weights[0].address.buffer
                      if weights_buffer in copy_map:
                          weights_buffer = copy_map[weights_buffer]
  
                      # In case of U65 512 mac variant the weights are split across two cores
                      # and need to be encoded separately
                      weights2_buffer = (
                          npu_op.weights[1].address.buffer
                          if accel_config
                          == vapi.NpuAccelerator.Ethos_U65_512
                          else None
                      )
                      if weights2_buffer in copy_map:
                          weights2_buffer = copy_map[weights2_buffer]
  
                      new_weights, split_idx = _encode_weights_or_bias(
                          weights_buffer,
                          weights2_buffer,
                          stmt,
                          _encode_weights,
                      )
                      _declare_constant_buffer(
                          weights_buffer, new_weights, split_idx
                      )
  
                      # Align the scale_bias to 16 bytes
                      scale_bias_buffer = npu_op.biases[
                          0
                      ].address.buffer
                      if scale_bias_buffer in copy_map:
                          scale_bias_buffer = copy_map[
                              scale_bias_buffer
                          ]
                      scale_bias2_buffer = (
                          npu_op.biases[1].address.buffer
                          if accel_config
                          == vapi.NpuAccelerator.Ethos_U65_512
                          else None
                      )
                      if scale_bias2_buffer in copy_map:
                          scale_bias2_buffer = copy_map[
                              scale_bias2_buffer
                          ]
  
                      (
                          new_scale_bias,
                          split_idx,
                      ) = _encode_weights_or_bias(
                          scale_bias_buffer,
                          scale_bias2_buffer,
                          stmt,
                          _align_scale_bias,
                      )
  
                      _declare_constant_buffer(
                          scale_bias_buffer, new_scale_bias, split_idx
                      )
  
          tvm.tir.stmt_functor.post_order_visit(stmt, _visit)
  
          return {
              "copied_buffers": copied_buffers,
              "constant_buffer_replacements": constant_buffer_replacements,
          }
  
      def transform_stmt(
          stmt,
          buf_remap,
          var_remap,
          pointer_to_buffer,
          new_buffer_var_to_const,
          new_buffer_to_split_idx,
      ):
          def _visit_rewrite(stmt):
              if isinstance(stmt, tvm.tir.Call):
                  # For extern calls, we need to rewrite pairs of arguments corresponding to
                  # base address load and the length of the load.
                  old_args = list(stmt.args)
  
                  new_args = [stmt.args[0]]
                  for prev_arg, arg in zip(old_args[:-1], old_args[1:]):
                      # If the previous argument was a load from an
                      # encoded buffer, the current should be a length.
                      if (
                          isinstance(prev_arg, tvm.tir.BufferLoad)
                          and prev_arg.buffer.data
                          in new_buffer_var_to_const
                      ):
                          buffer_size = np.prod(
                              list(prev_arg.buffer.shape)
                          )
                          arg = buffer_size
                          # We have to check for split weights/bias for conv2d and depthwise_conv2d
                          if old_args[0] in (
                              "ethosu_conv2d",
                              "depthwise_conv2d",
                          ):
                              # We have split weights/bias
                              if (
                                  prev_arg.buffer
                                  in new_buffer_to_split_idx
                              ):
                                  split_idx = new_buffer_to_split_idx[
                                      prev_arg.buffer
                                  ]
                                  # The first half of the split buffer
                                  if prev_arg.indices[0] == 0:
                                      arg = split_idx
                                  # the second half of the split buffer
                                  else:
                                      arg = buffer_size - split_idx
  
                      new_args.append(arg)
  
                  return tvm.tir.Call(
                      stmt.dtype, stmt.op, new_args, stmt.span
                  )
  
              if isinstance(stmt, tvm.tir.Allocate):
                  # Where a pointer needs rewriting, the allocate for it must be rewritten
                  allocate_pointer = stmt.buffer_var
                  if allocate_pointer in var_remap:
                      new_allocate_pointer = var_remap[allocate_pointer]
                      new_buffer = pointer_to_buffer[
                          new_allocate_pointer
                      ]
  
                      return tvm.tir.Allocate(
                          new_buffer.data,
                          new_buffer.dtype,
                          new_buffer.shape,
                          stmt.condition,
                          stmt.body,
                          stmt.span,
                      )
  
              # The following rewrites would be better expressed by just
              # rewriting the Buffers. However ir_transform doesn't
              # visit Buffers, so instead we do the next best thing and
              # rewrite the nodes which contain the Buffers.
              if isinstance(stmt, tvm.tir.BufferLoad):
                  if stmt.buffer in buf_remap:
                      new_buffer = buf_remap[stmt.buffer]
                      new_indices = stmt.indices
                      offset = new_indices[0]
                      if (
                          offset != 0
                          and new_buffer in new_buffer_to_split_idx
                      ):
                          offset = new_buffer_to_split_idx[new_buffer]
                      return tvm.tir.BufferLoad(
                          buf_remap[stmt.buffer], [offset], stmt.span
                      )
  
              if isinstance(stmt, tvm.tir.AttrStmt):
                  node_pointer = stmt.node
                  if node_pointer in var_remap:
                      return tvm.tir.AttrStmt(
                          var_remap[node_pointer],
                          stmt.attr_key,
                          stmt.value,
                          stmt.body,
                          stmt.span,
                      )
  
              return None
  
          return tvm.tir.stmt_functor.ir_transform(
              stmt,
              None,
              _visit_rewrite,
              [
                  "tir.Call",
                  "tir.Allocate",
                  "tir.BufferLoad",
                  "tir.AttrStmt",
              ],
          )
  
      def _collect_parameter_buffer_aliases(prim_func):
          buffer_vars = {}
          for param in prim_func.params:
              if param in prim_func.buffer_map:
                  buf = prim_func.buffer_map[param]
                  buffer_vars[buf.data] = {buf}
  
          def visit(node):
              if isinstance(
                  node,
                  (
                      tvm.tir.BufferStore,
                      tvm.tir.BufferLoad,
                      tvm.tir.DeclBuffer,
                  ),
              ):
                  buf = node.buffer
                  if buf.data in buffer_vars:
                      buffer_vars[buf.data].add(buf)
  
          tvm.tir.stmt_functor.post_order_visit(prim_func.body, visit)
          return buffer_vars
  
      def _ftransform(f, mod, ctx):
          param_buffer_var_usage = _collect_parameter_buffer_aliases(f)
  
          # Step 0: Unpack the constant dictionary in terms of the
          # functions buffers.
          old_buffer_var_to_const = {}
          for i, param in enumerate(f.params):
              if i in const_dict:
                  old_buffer_var_to_const[
                      f.buffer_map[param].data
                  ] = const_dict[i]
  
          # Step 1: Collect information on the buffers that will be
          # replaced by encodings.
          buffer_information = collect_encoding_definitions(
              f.body, old_buffer_var_to_const
          )
  
          # Step 2: Generate variable/buffer remaps, based on the
          # collected information.
          buf_remap = {}
          new_buffer_var_to_const = {}
          new_buffer_to_split_idx = {}
  
          def define_remap(old_buf, new_buf):
              try:
                  old_buffers = param_buffer_var_usage[old_buf.data]
              except KeyError:
                  old_buffers = [old_buf]
  
              for old_buffer in old_buffers:
                  buf_remap[old_buffer] = new_buf
  
          # Any encoded buffers must be replaced
          for info in buffer_information[
              "constant_buffer_replacements"
          ]:
              define_remap(info["old_buffer"], info["new_buffer"])
  
              new_buffer_var_to_const[info["new_buffer"].data] = info[
                  "encoded_constants"
              ]
  
              if info["split_idx"]:
                  new_buffer_to_split_idx[info["new_buffer"]] = info[
                      "split_idx"
                  ]
  
          # Any buffers that are copied into from an encoded buffer must
          # be replaced.
          for info in buffer_information["copied_buffers"]:
              copy_source = info["source"]
              while copy_source in buf_remap:
                  copy_source = buf_remap[copy_source]
  
              copy_dest = info["dest"]
  
              if (
                  copy_source.shape != copy_dest.shape
                  or copy_source.dtype != copy_dest.dtype
              ):
                  new_dest = tvm.tir.decl_buffer(
                      shape=copy_source.shape,
                      dtype=copy_source.dtype,
                      name=copy_dest.name,
                      scope=copy_dest.scope(),
                  )
                  define_remap(copy_dest, new_dest)
                  if copy_source.data in new_buffer_var_to_const:
                      new_buffer_var_to_const[
                          new_dest.data
                      ] = new_buffer_var_to_const[copy_source.data]
  
                  if copy_source in new_buffer_to_split_idx:
                      new_buffer_to_split_idx[
                          new_dest
                      ] = new_buffer_to_split_idx[copy_source]
  
          # Define additional dependent lookup tables.
          var_remap = {
              old.data: new.data for (old, new) in buf_remap.items()
          }
          pointer_to_buffer = {
              buf.data: buf
              for (old, new) in buf_remap.items()
              for buf in [old, new]
          }
  
          # Step 3: Then perform the rewrites
          new_body = transform_stmt(
              f.body,
              buf_remap,
              var_remap,
              pointer_to_buffer,
              new_buffer_var_to_const,
              new_buffer_to_split_idx,
          )
  
          # Step 4: Rewrite the buffer map and const dict to instead use the encoded versions
          new_buffer_map = {}
          for i, param in enumerate(f.params):
              buffer = f.buffer_map[param]
              if buffer in buf_remap:
                  buffer = buf_remap[buffer]
  
              if buffer.data in new_buffer_var_to_const:
                  new_const_dict[i] = new_buffer_var_to_const[
                      buffer.data
                  ].flatten()
              elif buffer.data in old_buffer_var_to_const:
                  new_const_dict[i] = old_buffer_var_to_const[
                      buffer.data
                  ].flatten()
  
              new_buffer_map[param] = buffer
  
          new_f = tvm.tir.PrimFunc(
              f.params,
              new_body,
              f.ret_type,
              new_buffer_map,
              f.attrs,
              f.span,
          )
          return new_f
  
      def _encode_constants(mod):
          mod, divided_const_dict = DivideConstants(const_dict)(mod)
          const_dict.clear()
          for key, value in divided_const_dict.items():
              const_dict[key] = value
          transform_func = tvm.tir.transform.prim_func_pass(
              _ftransform,
              opt_level=0,
              name="tir.contrib.ethos-u.encode_constants",
          )
          new_func = transform_func(mod)
          return new_func, new_const_dict
  
      return _encode_constants
  
imports: from collections import namedtuple
  from typing import Optional
  import numpy
  from ethosu.vela import api
  import tvm
  from tvm.relay.backend.contrib.ethosu import vela_api
  from tvm.relay.backend.contrib.ethosu import tir_to_cs_translator
  from .convolution import get_conv2d_params
  from .depthwise import get_depthwise_conv2d_params
  from .pooling import get_pooling_params
  from .binary_elementwise import get_binary_elementwise_params
  from .identity import get_identity_params
  from .unary_elementwise import get_unary_elementwise_params
  from .transform import get_copy_params
  from .producers_consumers import ProducersConsumers
  from ..None import _ffi_api
label_desc: Use image node_id 4 for calling a global function with example usage: EncodeConstants(const_dict) and returns: _encode_constants, dict, tvm, buffer_vars, new_f, new_func,  new_const_dict, value, value, new_const,  split_idx, None, new_const,  None, tvm, tvm, tvm, tvm