code: def main(a: T.handle, b: T.handle):
      A = T.match_buffer(a, (1024, 1024))
      B = T.match_buffer(b, (1024, 1024))
      with T.block("root"):
          T.block_attr({"warp_execution": True})
          for bx in T.thread_binding(8, thread="blockIdx.x"):
              for by in T.thread_binding(8, thread="blockIdx.y"):
                  for ty in T.thread_binding(8, thread="threadIdx.y"):
                      with T.block(""):
                          T.reads(
                              A[
                                  bx * 128 : bx * 128 + 128,
                                  by * 128 : by * 128 + 128,
                              ]
                          )
                          T.writes(
                              B[
                                  bx * 128 : bx * 128 + 128,
                                  by * 128 : by * 128 + 128,
                              ]
                          )
                          A_shared_dyn = T.alloc_buffer(
                              (128, 128),
                              strides=(128, 1),
                              scope="shared.dyn",
                          )
                          with T.block("A_shared"):
                              T.reads(
                                  A[
                                      bx * 128 : bx * 128 + 128,
                                      by * 128 : by * 128 + 128,
                                  ]
                              )
                              T.writes(A_shared_dyn[0:128, 0:128])
                              T.block_attr(
                                  {
                                      "auto_copy": 1,
                                      "local_stage": True,
                                      "vector_bytes": 16,
                                  }
                              )
                              A_shared_dyn_local = T.alloc_buffer(
                                  (16, 4), scope="local"
                              )
                              for ax0_ax1_fused_1 in T.thread_binding(
                                  8, thread="threadIdx.y"
                              ):
                                  for (
                                      ax0_ax1_fused_2
                                  ) in T.thread_binding(
                                      32, thread="threadIdx.x"
                                  ):
                                      for (
                                          ax0_ax1_fused_0_cache
                                      ) in range(16):
                                          for (
                                              ax0_ax1_fused_3_cache
                                          ) in T.vectorized(4):
                                              A_shared_dyn_local[
                                                  ax0_ax1_fused_0_cache
                                                  * 8
                                                  * 32
                                                  * 4
                                                  // 128
                                                  % 128
                                                  // 8,
                                                  ax0_ax1_fused_3_cache
                                                  % 128,
                                              ] = A[
                                                  bx * 128
                                                  + (
                                                      (
                                                          (
                                                              ax0_ax1_fused_0_cache
                                                              * 8
                                                              + ax0_ax1_fused_1
                                                          )
                                                          * 32
                                                          + ax0_ax1_fused_2
                                                      )
                                                      * 4
                                                      + ax0_ax1_fused_3_cache
                                                  )
                                                  // 128
                                                  % 128,
                                                  by * 128
                                                  + (
                                                      (
                                                          (
                                                              ax0_ax1_fused_0_cache
                                                              * 8
                                                              + ax0_ax1_fused_1
                                                          )
                                                          * 32
                                                          + ax0_ax1_fused_2
                                                      )
                                                      * 4
                                                      + ax0_ax1_fused_3_cache
                                                  )
                                                  % 128,
                                              ]
                                      for ax0_ax1_fused_0 in range(16):
                                          for (
                                              ax0_ax1_fused_3
                                          ) in T.vectorized(4):
                                              A_shared_dyn[
                                                  (
                                                      (
                                                          (
                                                              ax0_ax1_fused_0
                                                              * 8
                                                              + ax0_ax1_fused_1
                                                          )
                                                          * 32
                                                          + ax0_ax1_fused_2
                                                      )
                                                      * 4
                                                      + ax0_ax1_fused_3
                                                  )
                                                  // 128
                                                  % 128,
                                                  (
                                                      (
                                                          (
                                                              ax0_ax1_fused_0
                                                              * 8
                                                              + ax0_ax1_fused_1
                                                          )
                                                          * 32
                                                          + ax0_ax1_fused_2
                                                      )
                                                      * 4
                                                      + ax0_ax1_fused_3
                                                  )
                                                  % 128,
                                              ] = A_shared_dyn_local[
                                                  ax0_ax1_fused_0
                                                  * 8
                                                  * 32
                                                  * 4
                                                  // 128
                                                  % 128
                                                  // 8,
                                                  ax0_ax1_fused_3 % 128,
                                              ]
                          with T.block("B"):
                              T.reads(A_shared_dyn[0:128, 0:128])
                              T.writes(
                                  B[
                                      bx * 128 : bx * 128 + 128,
                                      by * 128 : by * 128 + 128,
                                  ]
                              )
                              for ax0 in range(128):
                                  for ax1 in range(128):
                                      B[
                                          bx * 128 + ax0, by * 128 + ax1
                                      ] = A_shared_dyn[ax0, ax1]
  
imports: import tvm
  from tvm import te
  from tvm.script import tir
  import sys
  import pytest
label_desc: Use image node_id 1 for calling the TransformedGlobalToSharedWithLocalStage obj's underlying member method code with example usage: obj.main(a, b) without return types