code: def main(
      p0: T.Buffer((128, 128), "float32"),
      p1: T.Buffer((128, 128), "float32"),
      T_add: T.Buffer((128, 128), "float32"),
  ) -> None:
      # function attr dict
      T.func_attr(
          {
              "global_symbol": "main",
              "tir.noalias": True,
              "layout_free_buffers": [1],
          }
      )
      # body
      # with T.block("root")
      T_matmul_NT_local = T.alloc_buffer(
          [128, 128], dtype="float32", scope="local"
      )
      p0_shared = T.alloc_buffer(
          [128, 128], dtype="float32", scope="shared"
      )
      p1_shared = T.alloc_buffer(
          [128, 128], dtype="float32", scope="shared"
      )
      for i0_0_i1_0_fused in T.thread_binding(
          32,
          thread="blockIdx.x",
          annotations={
              "pragma_auto_unroll_max_step": 64,
              "pragma_unroll_explicit": 1,
          },
      ):
          for i0_1_i1_1_fused in T.thread_binding(
              1, thread="vthread.x"
          ):
              for i0_2_i1_2_fused in T.thread_binding(
                  128, thread="threadIdx.x"
              ):
                  for (
                      i0_3_init,
                      i1_3_init,
                      i0_4_init,
                      i1_4_init,
                  ) in T.grid(1, 4, 1, 1):
                      with T.block("T_matmul_NT_init"):
                          i = T.axis.spatial(
                              128,
                              i0_0_i1_0_fused // 4 * 16
                              + i0_2_i1_2_fused // 8
                              + i0_3_init
                              + i0_4_init,
                          )
                          j = T.axis.spatial(
                              128,
                              i0_0_i1_0_fused % 4 * 32
                              + i0_2_i1_2_fused % 8 * 4
                              + i1_3_init
                              + i1_4_init,
                          )
                          T.reads()
                          T.writes(T_matmul_NT_local[i, j])
                          T.block_attr(
                              {
                                  "layout_free_placeholders": [],
                                  "meta_schedule.thread_extent_high_inclusive": 256,
                                  "meta_schedule.thread_extent_low_inclusive": 16,
                                  "meta_schedule.tiling_structure": "SSSRRSRS",
                              }
                          )
                          T_matmul_NT_local[i, j] = T.float32(0)
                  for i2_0 in T.serial(32):
                      for ax0_ax1_fused_0 in T.serial(1):
                          for ax0_ax1_fused_1 in T.thread_binding(
                              128, thread="threadIdx.x"
                          ):
                              for ax0_ax1_fused_2 in T.vectorized(2):
                                  with T.block("p0_shared"):
                                      T.where(
                                          (
                                              ax0_ax1_fused_0 * 128
                                              + ax0_ax1_fused_1
                                          )
                                          * 2
                                          + ax0_ax1_fused_2
                                          < 64
                                      )
                                      v0 = T.axis.spatial(
                                          128,
                                          i0_0_i1_0_fused // 4 * 16
                                          + (
                                              ax0_ax1_fused_0 * 256
                                              + ax0_ax1_fused_1 * 2
                                              + ax0_ax1_fused_2
                                          )
                                          // 4,
                                      )
                                      v1 = T.axis.spatial(
                                          128,
                                          i2_0 * 4
                                          + (
                                              ax0_ax1_fused_0 * 256
                                              + ax0_ax1_fused_1 * 2
                                              + ax0_ax1_fused_2
                                          )
                                          % 4,
                                      )
                                      T.reads(p0[v0, v1])
                                      T.writes(p0_shared[v0, v1])
                                      p0_shared[v0, v1] = p0[v0, v1]
                      for ax0_ax1_fused_0 in T.serial(1):
                          for ax0_ax1_fused_1 in T.thread_binding(
                              128, thread="threadIdx.x"
                          ):
                              for ax0_ax1_fused_2 in T.vectorized(4):
                                  with T.block("p1_shared"):
                                      T.where(
                                          (
                                              ax0_ax1_fused_0 * 128
                                              + ax0_ax1_fused_1
                                          )
                                          * 4
                                          + ax0_ax1_fused_2
                                          < 128
                                      )
                                      v0 = T.axis.spatial(
                                          128,
                                          i0_0_i1_0_fused % 4 * 32
                                          + (
                                              ax0_ax1_fused_0 * 512
                                              + ax0_ax1_fused_1 * 4
                                              + ax0_ax1_fused_2
                                          )
                                          // 4,
                                      )
                                      v1 = T.axis.spatial(
                                          128,
                                          i2_0 * 4
                                          + (
                                              ax0_ax1_fused_0 * 512
                                              + ax0_ax1_fused_1 * 4
                                              + ax0_ax1_fused_2
                                          )
                                          % 4,
                                      )
                                      T.reads(p1[v0, v1])
                                      T.writes(p1_shared[v0, v1])
                                      p1_shared[v0, v1] = p1[v0, v1]
                      for i2_1, i0_3, i1_3, i2_2, i0_4, i1_4 in T.grid(
                          1, 1, 4, 4, 1, 1
                      ):
                          with T.block("T_matmul_NT_update"):
                              i = T.axis.spatial(
                                  128,
                                  i0_0_i1_0_fused // 4 * 16
                                  + i0_2_i1_2_fused // 8
                                  + i0_3
                                  + i0_4,
                              )
                              j = T.axis.spatial(
                                  128,
                                  i0_0_i1_0_fused % 4 * 32
                                  + i0_2_i1_2_fused % 8 * 4
                                  + i1_3
                                  + i1_4,
                              )
                              k = T.axis.reduce(
                                  128, i2_0 * 4 + i2_1 * 4 + i2_2
                              )
                              T.reads(
                                  T_matmul_NT_local[i, j],
                                  p0_shared[i, k],
                                  p1_shared[j, k],
                              )
                              T.writes(T_matmul_NT_local[i, j])
                              T.block_attr(
                                  {
                                      "layout_free_placeholders": [],
                                      "meta_schedule.thread_extent_high_inclusive": 256,
                                      "meta_schedule.thread_extent_low_inclusive": 16,
                                      "meta_schedule.tiling_structure": "SSSRRSRS",
                                  }
                              )
                              T_matmul_NT_local[i, j] = (
                                  T_matmul_NT_local[i, j]
                                  + p0_shared[i, k] * p1_shared[j, k]
                              )
                  for ax0, ax1 in T.grid(1, 4):
                      with T.block("T_matmul_NT_local"):
                          v0 = T.axis.spatial(
                              128,
                              i0_0_i1_0_fused // 4 * 16
                              + i0_2_i1_2_fused // 8
                              + ax0,
                          )
                          v1 = T.axis.spatial(
                              128,
                              i0_0_i1_0_fused % 4 * 32
                              + i0_2_i1_2_fused % 8 * 4
                              + ax1,
                          )
                          T.reads(T_matmul_NT_local[v0, v1])
                          T.writes(T_add[v0, v1])
                          T_add[v0, v1] = T_matmul_NT_local[
                              v0, v1
                          ] + T.float32(1)
  
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 the DenseAdd_scheduled_gpu obj's underlying member method code with example usage: obj.main(p0, p1, T_add) without return types