code: def test_cuda_tbg():
      # fmt: off
      @T.prim_func
      def tbg_0(query: T.Buffer((1, 128, 12, 64), "float32"), value: T.Buffer((1, 128, 12, 64), "float32"), C: T.Buffer((1, 12, 128, 128), "float32")) -> None:
          T.func_attr({"global_symbol": "main", "tir.noalias": T.bool(True)})
          with T.block("root"):
              T.reads()
              T.writes()
              T.block_attr({"meta_schedule.unroll_explicit": 1024})
              C_local = T.alloc_buffer((1, 12, 128, 128), scope="local")
              query_T_shared = T.alloc_buffer((1, 12, 128, 64), scope="shared")
              value_T_shared = T.alloc_buffer((1, 12, 64, 128), scope="shared")
              for b_0_h_0_i_0_j_0_fused in T.thread_binding(4, thread="blockIdx.x"):
                  for b_1_h_1_i_1_j_1_fused in T.thread_binding(192, thread="vthread.x"):
                      for b_2_h_2_i_2_j_2_fused in T.thread_binding(32, thread="threadIdx.x"):
                          for k_0 in range(8):
                              for ax0_ax1_ax2_ax3_fused in range(12288):
                                  with T.block("query_T_shared"):
                                      v0 = T.axis.spatial(1, 0)
                                      v1 = T.axis.spatial(12, ax0_ax1_ax2_ax3_fused // 1024)
                                      v2 = T.axis.spatial(128, ax0_ax1_ax2_ax3_fused % 1024 // 8)
                                      v3 = T.axis.spatial(64, k_0 * 8 + ax0_ax1_ax2_ax3_fused % 8)
                                      T.reads(query[v0, v2, v1, v3])
                                      T.writes(query_T_shared[v0, v1, v2, v3])
                                      T.block_attr({"meta_schedule.cooperative_fetch": 3})
                                      query_T_shared[v0, v1, v2, v3] = query[v0, v2, v1, v3]
                              for ax0_ax1_ax2_ax3_fused in range(3072):
                                  with T.block("value_T_shared"):
                                      v0 = T.axis.spatial(1, 0)
                                      v1 = T.axis.spatial(12, ax0_ax1_ax2_ax3_fused // 256)
                                      v2 = T.axis.spatial(64, k_0 * 8 + ax0_ax1_ax2_ax3_fused % 256 // 32)
                                      v3 = T.axis.spatial(128, b_0_h_0_i_0_j_0_fused * 32 + ax0_ax1_ax2_ax3_fused % 32)
                                      T.reads(value[v0, v3, v1, v2])
                                      T.writes(value_T_shared[v0, v1, v2, v3])
                                      T.block_attr({"meta_schedule.cooperative_fetch": 4})
                                      value_T_shared[v0, v1, v2, v3] = value[v0, v3, v1, v2]
                              for k_1, b_3, h_3, i_3, j_3, k_2, b_4, h_4, i_4, j_4 in T.grid(4, 1, 2, 1, 1, 2, 1, 1, 4, 1):
                                  with T.block("C"):
                                      v_b = T.axis.spatial(1, b_3 + b_4)
                                      v_h = T.axis.spatial(12, b_1_h_1_i_1_j_1_fused // 32 * 2 + h_3 + h_4)
                                      v_i = T.axis.spatial(128, b_1_h_1_i_1_j_1_fused % 32 // 8 * 32 + b_2_h_2_i_2_j_2_fused // 4 * 4 + i_3 * 4 + i_4)
                                      v_j = T.axis.spatial(128, b_0_h_0_i_0_j_0_fused * 32 + b_1_h_1_i_1_j_1_fused % 8 * 4 + b_2_h_2_i_2_j_2_fused % 4 + j_3 + j_4)
                                      v_k = T.axis.reduce(64, k_0 * 8 + k_1 * 2 + k_2)
                                      T.reads(query_T_shared[v_b, v_h, v_i, v_k], value_T_shared[v_b, v_h, v_k, v_j])
                                      T.writes(C_local[v_b, v_h, v_i, v_j])
                                      T.block_attr({"meta_schedule.thread_extent_high_inclusive": 1024, "meta_schedule.thread_extent_low_inclusive": 32, "meta_schedule.tiling_structure": "SSSRRSRS"})
                                      with T.init():
                                          C_local[v_b, v_h, v_i, v_j] = T.float32(0)
                                      C_local[v_b, v_h, v_i, v_j] = C_local[v_b, v_h, v_i, v_j] + query_T_shared[v_b, v_h, v_i, v_k] * value_T_shared[v_b, v_h, v_k, v_j]
                          for ax0, ax1, ax2, ax3 in T.grid(1, 2, 4, 1):
                              with T.block("C_local"):
                                  v0 = T.axis.spatial(1, ax0)
                                  v1 = T.axis.spatial(12, b_1_h_1_i_1_j_1_fused // 32 * 2 + ax1)
                                  v2 = T.axis.spatial(128, b_1_h_1_i_1_j_1_fused % 32 // 8 * 32 + b_2_h_2_i_2_j_2_fused // 4 * 4 + ax2)
                                  v3 = T.axis.spatial(128, b_0_h_0_i_0_j_0_fused * 32 + b_1_h_1_i_1_j_1_fused % 8 * 4 + b_2_h_2_i_2_j_2_fused % 4 + ax3)
                                  T.reads(C_local[v0, v1, v2, v3])
                                  T.writes(C[v0, v1, v2, v3])
                                  C[v0, v1, v2, v3] = C_local[v0, v1, v2, v3]
      # fmt: on
      decision_0 = [
          ("SamplePerfectTile", [1, 1, 1, 1, 1]),
          ("SamplePerfectTile", [1, 6, 1, 2, 1]),
          ("SamplePerfectTile", [1, 4, 8, 1, 4]),
          ("SamplePerfectTile", [4, 8, 4, 1, 1]),
          ("SamplePerfectTile", [8, 4, 2]),
          ("SampleCategorical", 2),
          ("SampleCategorical", 3),
          ("SampleCategorical", 4),
      ]
      mod = create_te_workload("TBG", 0)
      actual = _design_space(mod)
      check_sketches(
          mod,
          sketches=actual,
          expected_mods=[tbg_0],
          expected_decisions=[decision_0],
      )
  
file_path: tvm/tests/python/meta_schedule/test_meta_schedule_space_cuda.py
imports: from tvm import meta_schedule
  from tvm.meta_schedule.testing.space_generation import check_sketches, generate_design_space, print_sketches
  from tvm.meta_schedule.testing.te_workload import create_te_workload
  from tvm.script import tir
  from tvm.target import Target
label_desc: Use image node_id 15 for calling a global function with example usage: test_cuda_tbg() without return types