code: def conv2d(data, weight, pad="SAME", stride=1):
      """
      Create an extern op that compute data * weight and return result in output
  
      Parameters:
      ----------
      data: Tensor
          The input data, format NHWC
      weight: Tensor
          The conv weight, format output_feature * kH * kW * input_feature
      pad: str
          Padding method, 'SAME' or 'VALID'
      stride: int
          convolution stride
  
      Returns
      -------
      output: Tensor
          The result tensor
      """
      n, hi, wi, ci = data.shape
      co, kh, kw, ciw = weight.shape
      padding = 0 if pad == "SAME" else 1
      ho = hi // stride
      wo = wi // stride
  
      return te.extern(
          (n, ho, wo, co),
          [data, weight],
          lambda ins, outs: tvm.tir.call_packed(
              "tvm.contrib.mps.conv2d",
              ins[0],
              ins[1],
              outs[0],
              padding,
              stride,
          ),
          name="C",
      )
  
imports: import tvm
  from tvm import te
label_desc: Use image node_id 2 for calling a global function with example usage: conv2d(data, weight, pad, stride) and returns: te