code: def get_conv2d_vnni_mod(intrin_id):
      @tvm.script.ir_module
      class Conv2dInt8_NCHWc_scheduled:
          @T.prim_func
          def main(
              p0: T.Buffer((1, 32, 7, 7, 16), "uint8"),
              p1: T.Buffer((128, 32, 1, 1, 4, 16, 4), "int8"),
              p2: T.Buffer((1, 128, 1, 1, 16), "int32"),
              p3: T.Buffer((1, 128, 1, 1, 16), "float32"),
              p4: T.Buffer(1, "float32"),
              p5: T.Buffer((1, 128, 7, 7, 16), "uint8"),
              T_cast: T.Buffer((1, 128, 7, 7, 16), "int32"),
          ) -> None:
              # function attr dict
              T.func_attr(
                  {"global_symbol": "main", "tir.noalias": True}
              )
              # body
              # with T.block("root")
              conv2d_NCHWc_int8 = T.alloc_buffer(
                  [1, 128, 7, 7, 16], dtype="int32"
              )
              for (
                  i0_0_i1_0_i2_0_i3_0_i4_0_0_i0_1_i1_1_fused
              ) in T.parallel(
                  128,
                  annotations={
                      "pragma_auto_unroll_max_step": 64,
                      "pragma_unroll_explicit": 1,
                  },
              ):
                  for i2_1, i3_1, i4_0_1 in T.grid(7, 1, 1):
                      for (
                          i0_2_init,
                          i1_2_init,
                          i2_2_init,
                          i3_2_init,
                          i4_0_2_init,
                          i0_3_init,
                          i1_3_init,
                          i2_3_init,
                          i3_3_init,
                          i4_0_3_init,
                      ) in T.grid(1, 1, 1, 1, 1, 1, 1, 1, 7, 1):
                          with T.block("conv2d_NCHWc_int8_o_init"):
                              n = T.axis.spatial(
                                  1, i0_2_init + i0_3_init
                              )
                              oc_chunk = T.axis.spatial(
                                  128,
                                  i0_0_i1_0_i2_0_i3_0_i4_0_0_i0_1_i1_1_fused
                                  + i1_2_init
                                  + i1_3_init,
                              )
                              oh = T.axis.spatial(
                                  7, i2_1 + i2_2_init + i2_3_init
                              )
                              ow = T.axis.spatial(
                                  7,
                                  i3_1 * 7 + i3_2_init * 7 + i3_3_init,
                              )
                              oc_block_o = T.axis.spatial(
                                  1, i4_0_1 + i4_0_2_init + i4_0_3_init
                              )
                              T.reads()
                              T.writes(
                                  conv2d_NCHWc_int8[
                                      n, oc_chunk, oh, ow, 0:16
                                  ]
                              )
                              for i4_1 in T.vectorized(16):
                                  with T.block(
                                      "conv2d_NCHWc_int8_init"
                                  ):
                                      oc_block_i_init = T.axis.spatial(
                                          16, i4_1
                                      )
                                      T.reads()
                                      T.writes(
                                          conv2d_NCHWc_int8[
                                              n,
                                              oc_chunk,
                                              oh,
                                              ow,
                                              oc_block_i_init,
                                          ]
                                      )
                                      conv2d_NCHWc_int8[
                                          n,
                                          oc_chunk,
                                          oh,
                                          ow,
                                          oc_block_i_init,
                                      ] = 0
                      for (
                          i5_0,
                          i6_0,
                          i7_0,
                          i8_0,
                          i9_0_0,
                          i0_2,
                          i1_2,
                          i2_2,
                          i3_2,
                          i4_0_2,
                          i5_1,
                          i6_1,
                          i7_1,
                          i8_1,
                          i9_0_1,
                          i0_3,
                          i1_3,
                          i2_3,
                          i3_3,
                          i4_0_3,
                      ) in T.grid(
                          1,
                          1,
                          4,
                          4,
                          1,
                          1,
                          1,
                          1,
                          1,
                          1,
                          1,
                          1,
                          8,
                          1,
                          1,
                          1,
                          1,
                          1,
                          7,
                          1,
                      ):
                          with T.block("conv2d_NCHWc_int8_o_update"):
                              n = T.axis.spatial(1, i0_2 + i0_3)
                              oc_chunk = T.axis.spatial(
                                  128,
                                  i0_0_i1_0_i2_0_i3_0_i4_0_0_i0_1_i1_1_fused
                                  + i1_2
                                  + i1_3,
                              )
                              oh = T.axis.spatial(7, i2_1 + i2_2 + i2_3)
                              ow = T.axis.spatial(
                                  7, i3_1 * 7 + i3_2 * 7 + i3_3
                              )
                              oc_block_o = T.axis.spatial(
                                  1, i4_0_1 + i4_0_2 + i4_0_3
                              )
                              kh = T.axis.reduce(1, i5_0 + i5_1)
                              kw = T.axis.reduce(1, i6_0 + i6_1)
                              ic_outer = T.axis.reduce(
                                  32, i7_0 * 8 + i7_1
                              )
                              ic_f_inner = T.axis.reduce(4, i8_0 + i8_1)
                              ic_s_inner_o = T.axis.reduce(
                                  1, i9_0_0 + i9_0_1
                              )
                              T.reads(
                                  conv2d_NCHWc_int8[
                                      n, oc_chunk, oh, ow, 0:16
                                  ],
                                  p0[
                                      n,
                                      ic_outer,
                                      oh + kh,
                                      ow + kw,
                                      ic_f_inner * 4 : ic_f_inner * 4
                                      + 4,
                                  ],
                                  p1[
                                      oc_chunk,
                                      ic_outer,
                                      kh,
                                      kw,
                                      ic_f_inner,
                                      0:16,
                                      0:4,
                                  ],
                              )
                              T.writes(
                                  conv2d_NCHWc_int8[
                                      n, oc_chunk, oh, ow, 0:16
                                  ]
                              )
                              A = T.match_buffer(
                                  p0[
                                      n,
                                      ic_outer,
                                      oh + kh,
                                      ow + kw,
                                      ic_f_inner * 4 : ic_f_inner * 4
                                      + 4,
                                  ],
                                  [4],
                                  dtype="uint8",
                                  offset_factor=1,
                              )
                              B = T.match_buffer(
                                  p1[
                                      oc_chunk,
                                      ic_outer,
                                      kh,
                                      kw,
                                      ic_f_inner,
                                      0:16,
                                      0:4,
                                  ],
                                  [16, 4],
                                  dtype="int8",
                                  offset_factor=1,
                              )
                              C = T.match_buffer(
                                  conv2d_NCHWc_int8[
                                      n, oc_chunk, oh, ow, 0:16
                                  ],
                                  [16],
                                  dtype="int32",
                                  offset_factor=1,
                              )
                              A_u8x4: T.uint8x4 = A[0:4]
                              A_i32: T.int32 = T.reinterpret(
                                  A_u8x4, dtype="int32"
                              )
                              B_i8x64: T.int8x64 = B[0, 0:64]
                              B_i32x16: T.int32x16 = T.reinterpret(
                                  B_i8x64, dtype="int32x16"
                              )
                              C_i32x16: T.int32x16 = C[0:16]
                              C[0:16] = T.call_llvm_pure_intrin(
                                  T.uint32(intrin_id),
                                  T.uint32(3),
                                  C_i32x16,
                                  T.broadcast(A_i32, 16),
                                  B_i32x16,
                                  dtype="int32x16",
                              )
                      for ax0, ax1, ax2, ax3 in T.grid(1, 1, 1, 7):
                          for ax4_fused in T.vectorized(16):
                              with T.block("T_cast_8"):
                                  ax0_1 = T.axis.spatial(1, ax0)
                                  ax1_1 = T.axis.spatial(
                                      128,
                                      i0_0_i1_0_i2_0_i3_0_i4_0_0_i0_1_i1_1_fused
                                      + ax1,
                                  )
                                  ax2_1 = T.axis.spatial(7, i2_1 + ax2)
                                  ax3_1, ax4 = T.axis.remap(
                                      "SS", [ax3, ax4_fused]
                                  )
                                  T.reads(
                                      conv2d_NCHWc_int8[
                                          ax0_1,
                                          ax1_1,
                                          ax2_1,
                                          ax3_1,
                                          ax4,
                                      ],
                                      p2[ax0_1, ax1_1, 0, 0, ax4],
                                      p3[ax0_1, ax1_1, 0, 0, ax4],
                                      p4[0],
                                      p5[
                                          ax0_1,
                                          ax1_1,
                                          ax2_1,
                                          ax3_1,
                                          ax4,
                                      ],
                                  )
                                  T.writes(
                                      T_cast[
                                          ax0_1,
                                          ax1_1,
                                          ax2_1,
                                          ax3_1,
                                          ax4,
                                      ]
                                  )
                                  T_cast[
                                      ax0_1, ax1_1, ax2_1, ax3_1, ax4
                                  ] = T.cast(
                                      T.max(
                                          T.min(
                                              T.cast(
                                                  T.max(
                                                      T.min(
                                                          T.cast(
                                                              T.floor(
                                                                  T.float32(
                                                                      0.95489668846130371
                                                                  )
                                                                  * (
                                                                      T.cast(
                                                                          T.cast(
                                                                              T.max(
                                                                                  T.min(
                                                                                      T.cast(
                                                                                          T.floor(
                                                                                              T.cast(
                                                                                                  conv2d_NCHWc_int8[
                                                                                                      ax0_1,
                                                                                                      ax1_1,
                                                                                                      ax2_1,
                                                                                                      ax3_1,
                                                                                                      ax4,
                                                                                                  ]
                                                                                                  + p2[
                                                                                                      ax0_1,
                                                                                                      ax1_1,
                                                                                                      0,
                                                                                                      0,
                                                                                                      ax4,
                                                                                                  ],
                                                                                                  "float32",
                                                                                              )
                                                                                              * p3[
                                                                                                  ax0_1,
                                                                                                  ax1_1,
                                                                                                  0,
                                                                                                  0,
                                                                                                  ax4,
                                                                                              ]
                                                                                              + T.float32(
                                                                                                  65.5
                                                                                              ),
                                                                                              dtype="float32",
                                                                                          ),
                                                                                          "int32",
                                                                                      ),
                                                                                      255,
                                                                                  ),
                                                                                  0,
                                                                              ),
                                                                              "uint8",
                                                                          ),
                                                                          "float32",
                                                                      )
                                                                      - p4[
                                                                          0
                                                                      ]
                                                                  )
                                                                  + T.float32(
                                                                      0.5
                                                                  ),
                                                                  dtype="float32",
                                                              ),
                                                              "int32",
                                                          )
                                                          + T.cast(
                                                              T.floor(
                                                                  T.float32(
                                                                      0.71245479583740234
                                                                  )
                                                                  * T.cast(
                                                                      p5[
                                                                          ax0_1,
                                                                          ax1_1,
                                                                          ax2_1,
                                                                          ax3_1,
                                                                          ax4,
                                                                      ],
                                                                      "float32",
                                                                  )
                                                                  + T.float32(
                                                                      0.5
                                                                  ),
                                                                  dtype="float32",
                                                              ),
                                                              "int32",
                                                          ),
                                                          255,
                                                      ),
                                                      0,
                                                  ),
                                                  "uint8",
                                              ),
                                              T.uint8(255),
                                          ),
                                          T.uint8(0),
                                      ),
                                      "int32",
                                  )
  
      return Conv2dInt8_NCHWc_scheduled
  
imports: import pytest
  import tvm
  import tvm.meta_schedule
  import tvm.testing
  from tvm.script import tir
  from tvm.target import Target
  from tvm.target.codegen import llvm_lookup_intrinsic_id
  from tvm.tir import Schedule, floordiv, floormod
  from tvm.tir.tensor_intrin.cuda import *
  from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN
label_desc: Use image node_id 1 for calling a global function with example usage: get_conv2d_vnni_mod(intrin_id) and returns: Conv2dInt8_NCHWc_scheduled