code: def GmmCuda1(
      X: T.Buffer((1, 128, 128), "float32"),
      Y: T.Buffer((1, 128, 128), "float32"),
      Z: T.Buffer((1, 128, 128), "float32"),
  ) -> None:
      Z_local = T.alloc_buffer(
          [1, 128, 128], dtype="float32", scope="local"
      )
      X_shared = T.alloc_buffer(
          [1, 128, 128], dtype="float32", scope="shared"
      )
      Y_shared = T.alloc_buffer(
          [1, 128, 128], dtype="float32", scope="shared"
      )
      for i0_0_i1_0_i2_0_fused in T.thread_binding(
          16, thread="blockIdx.x"
      ):
          for i0_1_i1_1_i2_1_fused in T.thread_binding(
              1, thread="vthread.x"
          ):
              for i0_2_i1_2_i2_2_fused in T.thread_binding(
                  128, thread="threadIdx.x"
              ):
                  for i1_3_init, i2_4_init in T.grid(4, 2):
                      with T.block("Z_init"):
                          b = T.axis.spatial(1, 0)
                          i = T.axis.spatial(
                              128,
                              i0_0_i1_0_i2_0_fused // 4 * 32
                              + i0_2_i1_2_i2_2_fused // 16 * 4
                              + i1_3_init,
                          )
                          j = T.axis.spatial(
                              128,
                              i0_0_i1_0_i2_0_fused % 4 * 32
                              + i0_2_i1_2_i2_2_fused % 16 * 2
                              + i2_4_init,
                          )
                          T.reads()
                          T.writes(Z_local[b, i, j])
                          Z_local[b, i, j] = T.float32(0)
                  for i3_0 in T.serial(4):
                      for ax0_ax1_ax2_fused_0 in T.serial(4):
                          for ax0_ax1_ax2_fused_1 in T.thread_binding(
                              128, thread="threadIdx.x"
                          ):
                              for ax0_ax1_ax2_fused_2 in T.vectorized(
                                  2
                              ):
                                  with T.block("X_shared"):
                                      v0 = T.axis.spatial(1, 0)
                                      v1 = T.axis.spatial(
                                          128,
                                          i0_0_i1_0_i2_0_fused // 4 * 32
                                          + (
                                              ax0_ax1_ax2_fused_0 * 256
                                              + ax0_ax1_ax2_fused_1 * 2
                                              + ax0_ax1_ax2_fused_2
                                          )
                                          // 32,
                                      )
                                      v2 = T.axis.spatial(
                                          128,
                                          i3_0 * 32
                                          + (
                                              ax0_ax1_ax2_fused_0 * 256
                                              + ax0_ax1_ax2_fused_1 * 2
                                              + ax0_ax1_ax2_fused_2
                                          )
                                          % 32,
                                      )
                                      T.reads(X[v0, v1, v2])
                                      T.writes(X_shared[v0, v1, v2])
                                      X_shared[v0, v1, v2] = X[
                                          v0, v1, v2
                                      ]
                      for ax0_ax1_ax2_fused_0 in T.serial(8):
                          for ax0_ax1_ax2_fused_1 in T.thread_binding(
                              128, thread="threadIdx.x"
                          ):
                              with T.block("Y_shared"):
                                  v0 = T.axis.spatial(1, 0)
                                  v1 = T.axis.spatial(
                                      128,
                                      i3_0 * 32
                                      + (
                                          ax0_ax1_ax2_fused_0 * 128
                                          + ax0_ax1_ax2_fused_1
                                      )
                                      // 32,
                                  )
                                  v2 = T.axis.spatial(
                                      128,
                                      i0_0_i1_0_i2_0_fused % 4 * 32
                                      + (
                                          ax0_ax1_ax2_fused_0 * 128
                                          + ax0_ax1_ax2_fused_1
                                      )
                                      % 32,
                                  )
                                  T.reads(Y[v0, v1, v2])
                                  T.writes(Y_shared[v0, v1, v2])
                                  Y_shared[v0, v1, v2] = Y[v0, v1, v2]
                      for (
                          i3_1,
                          i0_3,
                          i1_3,
                          i2_3,
                          i3_2,
                          i0_4,
                          i1_4,
                          i2_4,
                      ) in T.grid(1, 1, 4, 1, 32, 1, 1, 2):
                          with T.block("Z_update"):
                              b = T.axis.spatial(1, 0)
                              i = T.axis.spatial(
                                  128,
                                  i0_0_i1_0_i2_0_fused // 4 * 32
                                  + i0_2_i1_2_i2_2_fused // 16 * 4
                                  + i1_3,
                              )
                              j = T.axis.spatial(
                                  128,
                                  i0_0_i1_0_i2_0_fused % 4 * 32
                                  + i0_2_i1_2_i2_2_fused % 16 * 2
                                  + i2_4,
                              )
                              k = T.axis.reduce(128, i3_0 * 32 + i3_2)
                              T.block_attr(
                                  {
                                      "meta_schedule.thread_extent_low_inclusive": 0,
                                      "meta_schedule.thread_extent_high_inclusive": 32,
                                  }
                              )
                              T.reads(
                                  Z_local[b, i, j],
                                  X_shared[b, i, k],
                                  Y_shared[b, k, j],
                              )
                              T.writes(Z_local[b, i, j])
                              Z_local[b, i, j] = (
                                  Z_local[b, i, j]
                                  + X_shared[b, i, k]
                                  * Y_shared[b, k, j]
                              )
                  for ax0, ax1, ax2 in T.grid(1, 4, 2):
                      with T.block("Z_local"):
                          v0 = T.axis.spatial(1, ax0)
                          v1 = T.axis.spatial(
                              128,
                              i0_0_i1_0_i2_0_fused // 4 * 32
                              + i0_2_i1_2_i2_2_fused // 16 * 4
                              + ax1,
                          )
                          v2 = T.axis.spatial(
                              128,
                              i0_0_i1_0_i2_0_fused % 4 * 32
                              + i0_2_i1_2_i2_2_fused % 16 * 2
                              + ax2,
                          )
                          T.reads(Z_local[v0, v1, v2])
                          T.writes(Z[v0, v1, v2])
                          Z[v0, v1, v2] = Z_local[v0, v1, v2]
  
imports: import pytest
  import tvm
  import tvm.testing
  from tvm import meta_schedule
  from tvm import tir
  from tvm.script import tir
  from tvm.target import Target
label_desc: Use image node_id 4 for calling a global function with example usage: GmmCuda1(X, Y, Z) without return types