code: def depthwise_conv2d_compute(
      ifm: te.Tensor,
      weight: te.Tensor,
      scale_bias: te.Tensor,
      lut: te.Tensor,
      ifm_scale: float,
      ifm_zero_point: int,
      weight_zero_point: int,
      ofm_scale: float,
      ofm_zero_point: int,
      strides: Tuple[int, int],
      padding: Tuple[int, int, int, int],
      dilation: Union[Tuple[int, int], List[int]],
      activation: str,
      clip_min: int,
      clip_max: int,
      rounding_mode: str,
      upscale: str,
      ifm_layout: str,
      ofm_layout: str,
      ofm_dtype: str,
  ) -> te.Tensor:
      """A compute operator representing the capabilities of 2D convolution for the NPU.
  
      Parameters
      ----------
      ifm : te.Tensor
          The Input Feature Map tensor (IFM).
      weight : te.Tensor
          The weight tensor.
      scale_bias : te.Tensor
          The packed per-channel weight scale and bias tensor.
      lut : te.Tensor
          The look-up table of values to use if activation = "LUT".
      ifm_scale : float
          The quantization scale for the Input Feature Map tensor.
      ifm_zero_point : int
          The quantization zero point for the Input Feature Map tensor.
      weight_zero_point : int
          The quantization zero point for the weight tensor.
      ofm_scale : float
          The quantization scale for the Output Feature Map tensor.
      ofm_zero_point : int
          The quantization zero point for the Output Feature Map tensor.
      strides : tuple
          The 2 dimensional strides as (stride_height, stride_width).
      padding : tuple
          The 4 dimensional padding as (pad_top, pad_left, pad_bottom, pad_right).
      dilation : Union[int, tuple, list]
          The 2 dimensional dilation as (dilation_height, dilation_width).
      activation : str
          The activation function to use.
              "NONE" - no activation function.
              "CLIP" - clip the output between clip_min and clip_max.
              "TANH" - tanh activation function.
              "SIGMOID" - sigmoid activation function.
              "LUT" - use a look-up table to perform the activation function.
      clip_min : int
          The minimum clipping value if activation = "CLIP".
      clip_max : int
          The maximum clipping value if activation = "CLIP".
      rounding_mode : str
          The rounding mode to apply to the Output Feature Map tensor.
              "TFL" - Tensorflow Lite rounding scheme.
              "TRUNCATE" - Truncate towards zero.
              "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.
      upscale : str
          The 2x2 upscaling mode to apply to the Input Feature Map tensor.
              "NONE" - no upscaling.
              "NEAREST" - upscale using nearest neighbour.
              "ZEROS" - upscale using zeros.
      ifm_layout : str
          The layout of the Input Feature Map tensor. Can be "NHWC" or "NHCWB16".
      ofm_layout : str
          The layout of the Output Feature Map tensor. Can be "NHWC" or "NHCWB16".
      ofm_dtype : str, optional
          The Output Feature Map tensor data type. Can be 'int8', 'uint8' or 'int16'.
  
      Returns
      -------
      te.Tensor
          The OFM tensor.
  
      """
      assert ifm.shape[0] == 1, "Only batch size 1 is supported"
      assert ifm_layout in {"NHWC", "NHCWB16"}
      assert ofm_layout in {"NHWC", "NHCWB16"}
  
      padding = [int(v) for v in padding]
      stride_h, stride_w = [int(v) for v in strides]
      dilation_h, dilation_w = [int(v) for v in dilation]
      channels, kernel_h, kernel_w, _ = [int(v) for v in weight.shape]
  
      # Compute operation for the IFM DMA pipeline
      dmaed_ifm = dma_ifm_compute(
          ifm, ifm_layout, ifm_zero_point, ifm_scale, channels, padding
      )
  
      # 2D Depthwise Convolution compute operation
      dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
      dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
      ofm_height = (
          dmaed_ifm.shape[1] - dilated_kernel_h
      ) // stride_h + 1
      ofm_width = (
          dmaed_ifm.shape[2] - dilated_kernel_w
      ) // stride_w + 1
      rh = te.reduce_axis((0, kernel_h), name="ry")
      rw = te.reduce_axis((0, kernel_w), name="rx")
  
      depthwise_conv2d_attrs = {
          "op": "ethosu_depthwise_conv2d",
          "weight_zero_point": weight_zero_point,
          "activation": activation,
          "clip_min": clip_min,
          "clip_max": clip_max,
          "rounding_mode": rounding_mode,
          "upscale": upscale,
          "stride_h": stride_h,
          "stride_w": stride_w,
          "dilation_h": dilation_h,
          "dilation_w": dilation_w,
      }
  
      has_lut = activation in ("TANH", "LUT", "SIGMOID")
  
      # This is a trick to insert the LUT tensor into the TE graph if LUT is present
      lut_expr = (lut[0] + lut[255]).astype(ifm.dtype) if has_lut else 0
  
      # Add the LUT tensor to the attributes to be able to later tell which tensor is the LUT
      if has_lut:
          depthwise_conv2d_attrs["lut"] = lut
  
      depthwise = te.compute(
          (1, ofm_height, ofm_width, channels),
          lambda nn, hh, ww, cc: te.sum(
              (
                  dmaed_ifm(
                      nn,
                      hh * stride_h + rh * dilation_h,
                      ww * stride_w + rw * dilation_w,
                      cc,
                  ).astype(ifm.dtype)
                  * weight[cc, rh, rw, 0].astype(ifm.dtype)
                  # This is a trick to load 10 elements of the scale_bias at once, not accurate maths
                  + (
                      scale_bias[cc, 0] * scale_bias[cc, 9] + lut_expr
                  ).astype(ifm.dtype)
              ).astype(ofm_dtype),
              axis=[rh, rw],
          ),
          name="ethosu_depthwise_conv2d",
          attrs=depthwise_conv2d_attrs,
      )
  
      nhwc_to_nhcwb16, nhcwb16_to_nhwc = get_layout_transform_matrices(
          channels
      )
  
      ifm_matrix = [
          [1, 0, 0, 0, 0],
          [0, stride_h, 0, 0, (dilated_kernel_h - stride_h)],
          [0, 0, stride_w, 0, (dilated_kernel_w - stride_w)],
          [0, 0, 0, 1, 0],
          [0, 0, 0, 0, 1],
      ]
      weights_matrix = [
          [0, 0, 0, 1, 0],
          [0, 0, 0, 0, kernel_h],
          [0, 0, 0, 0, kernel_w],
          [0, 0, 0, 0, 1],
          [0, 0, 0, 0, 1],
      ]
      bias_matrix = [
          [0, 0, 0, 1, 0],
          [0, 0, 0, 0, 10],
          [0, 0, 0, 0, 1],
      ]
      if ofm_layout == "NHCWB16":
          ifm_matrix = np.matmul(ifm_matrix, nhcwb16_to_nhwc).tolist()
          weights_matrix = np.matmul(
              weights_matrix, nhcwb16_to_nhwc
          ).tolist()
          bias_matrix = np.matmul(bias_matrix, nhcwb16_to_nhwc).tolist()
      if ifm_layout == "NHCWB16":
          ifm_matrix = np.matmul(nhwc_to_nhcwb16, ifm_matrix).tolist()
      ifm_propagator = Propagator(
          ifm_matrix,
          [0, -padding[0], -padding[1], 0]
          if ifm_layout == "NHWC"
          else [0, -padding[0], 0, -padding[1], 0],
      )
      weights_propagator = Propagator(
          weights_matrix,
          [0, 0, 0, 0],
      )
      bias_propagator = Propagator(
          bias_matrix,
          [0, 0],
      )
      propagator_attrs = {
          "ifm_propagator": ifm_propagator,
          "weights_propagator": weights_propagator,
          "bias_propagator": bias_propagator,
      }
  
      # Compute operation for the OFM DMA pipeline
      return dma_ofm_compute(
          depthwise,
          ofm_layout,
          ofm_zero_point,
          ofm_scale,
          channels,
          attrs=propagator_attrs,
      )
  
file_path: tvm/python/tvm/relay/backend/contrib/ethosu/te/depthwise.py
imports: from typing import Tuple, Union, List
  import numpy
  from tvm import te
  from tvm.contrib.ethosu.cascader import TESubgraph, EthosuPart, Propagator, register_matcher
  from .dma import dma_ofm_compute, dma_ifm_compute
  from .common import get_layout_transform_matrices
label_desc: Use image node_id 1 for calling a global function with example usage: depthwise_conv2d_compute(ifm, weight, scale_bias, lut, ifm_scale, ifm_zero_point, weight_zero_point, ofm_scale, ofm_zero_point, strides, padding, dilation, activation, clip_min, clip_max, rounding_mode, upscale, ifm_layout, ofm_layout, ofm_dtype) and returns: dma_ofm_compute