code: def test_cpu_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.parallel": 288, "meta_schedule.unroll_explicit": 64, "meta_schedule.vectorize": 64})
              query_T = T.alloc_buffer((1, 12, 128, 64))
              value_T = T.alloc_buffer((1, 12, 64, 128))
              C_global = T.alloc_buffer((1, 12, 128, 128))
              for b_0, h_0, i_0, j_0, b_1, h_1, i_1 in T.grid(1, 1, 1, 2, 1, 6, 2):
                  for ax0, ax1, ax2, ax3 in T.grid(1, 2, 64, 64):
                      with T.block("value_T"):
                          v_b = T.axis.spatial(1, ax0)
                          v_h = T.axis.spatial(12, h_1 * 2 + ax1)
                          v_d = T.axis.spatial(64, ax2)
                          v_l = T.axis.spatial(128, j_0 * 64 + ax3)
                          T.reads(value[v_b, v_l, v_h, v_d])
                          T.writes(value_T[v_b, v_h, v_d, v_l])
                          value_T[v_b, v_h, v_d, v_l] = value[v_b, v_l, v_h, v_d]
                  for ax0, ax1, ax2, ax3 in T.grid(1, 2, 64, 64):
                      with T.block("query_T"):
                          v_b = T.axis.spatial(1, ax0)
                          v_h = T.axis.spatial(12, h_1 * 2 + ax1)
                          v_l = T.axis.spatial(128, i_1 * 64 + ax2)
                          v_d = T.axis.spatial(64, ax3)
                          T.reads(query[v_b, v_l, v_h, v_d])
                          T.writes(query_T[v_b, v_h, v_l, v_d])
                          query_T[v_b, v_h, v_l, v_d] = query[v_b, v_l, v_h, v_d]
                  for j_1 in range(8):
                      for k_0, b_2, h_2, i_2, j_2, k_1, b_3, h_3, i_3, j_3 in T.grid(1, 1, 2, 2, 4, 64, 1, 1, 32, 2):
                          with T.block("C"):
                              v_b = T.axis.spatial(1, b_0 + b_1 + b_2 + b_3)
                              v_h = T.axis.spatial(12, h_0 * 12 + h_1 * 2 + h_2 + h_3)
                              v_i = T.axis.spatial(128, i_0 * 128 + i_1 * 64 + i_2 * 32 + i_3)
                              v_j = T.axis.spatial(128, j_0 * 64 + j_1 * 8 + j_2 * 2 + j_3)
                              v_k = T.axis.reduce(64, k_0 * 64 + k_1)
                              T.reads(query_T[v_b, v_h, v_i, v_k], value_T[v_b, v_h, v_k, v_j])
                              T.writes(C_global[v_b, v_h, v_i, v_j])
                              T.block_attr({"meta_schedule.tiling_structure": "SSRSRS"})
                              with T.init():
                                  C_global[v_b, v_h, v_i, v_j] = T.float32(0)
                              C_global[v_b, v_h, v_i, v_j] = C_global[v_b, v_h, v_i, v_j] + query_T[v_b, v_h, v_i, v_k] * value_T[v_b, v_h, v_k, v_j]
                      for ax0, ax1, ax2, ax3 in T.grid(1, 2, 64, 8):
                          with T.block("C_global"):
                              v0 = T.axis.spatial(1, ax0)
                              v1 = T.axis.spatial(12, h_1 * 2 + ax1)
                              v2 = T.axis.spatial(128, i_1 * 64 + ax2)
                              v3 = T.axis.spatial(128, j_0 * 64 + j_1 * 8 + ax3)
                              T.reads(C_global[v0, v1, v2, v3])
                              T.writes(C[v0, v1, v2, v3])
                              C[v0, v1, v2, v3] = C_global[v0, v1, v2, v3]
      @T.prim_func
      def tbg_1(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.parallel": 288, "meta_schedule.unroll_explicit": 64, "meta_schedule.vectorize": 64})
              query_T = T.alloc_buffer((1, 12, 128, 64))
              value_T = T.alloc_buffer((1, 12, 64, 128))
              C_global = T.alloc_buffer((1, 12, 128, 128))
              for b, h, l, d in T.grid(1, 12, 128, 64):
                  with T.block("query_T"):
                      v_b, v_h, v_l, v_d = T.axis.remap("SSSS", [b, h, l, d])
                      T.reads(query[v_b, v_l, v_h, v_d])
                      T.writes(query_T[v_b, v_h, v_l, v_d])
                      query_T[v_b, v_h, v_l, v_d] = query[v_b, v_l, v_h, v_d]
              for b_0, h_0, i_0, j_0 in T.grid(1, 1, 1, 2):
                  for b_1, h_1, i_1, j_1, k_0, b_2, h_2, i_2, j_2, k_1 in T.grid(1, 6, 2, 8, 1, 1, 2, 2, 4, 64):
                      for ax0, ax1, ax2, ax3 in T.grid(1, 1, 1, 2):
                          with T.block("value_T"):
                              v_b = T.axis.spatial(1, ax0)
                              v_h = T.axis.spatial(12, h_1 * 2 + h_2 + ax1)
                              v_d = T.axis.spatial(64, k_1 + ax2)
                              v_l = T.axis.spatial(128, j_0 * 64 + j_1 * 8 + j_2 * 2 + ax3)
                              T.reads(value[v_b, v_l, v_h, v_d])
                              T.writes(value_T[v_b, v_h, v_d, v_l])
                              value_T[v_b, v_h, v_d, v_l] = value[v_b, v_l, v_h, v_d]
                      for b_3, h_3, i_3, j_3 in T.grid(1, 1, 32, 2):
                          with T.block("C"):
                              v_b = T.axis.spatial(1, b_0 + b_1 + b_2 + b_3)
                              v_h = T.axis.spatial(12, h_0 * 12 + h_1 * 2 + h_2 + h_3)
                              v_i = T.axis.spatial(128, i_0 * 128 + i_1 * 64 + i_2 * 32 + i_3)
                              v_j = T.axis.spatial(128, j_0 * 64 + j_1 * 8 + j_2 * 2 + j_3)
                              v_k = T.axis.reduce(64, k_0 * 64 + k_1)
                              T.reads(query_T[v_b, v_h, v_i, v_k], value_T[v_b, v_h, v_k, v_j])
                              T.writes(C_global[v_b, v_h, v_i, v_j])
                              T.block_attr({"meta_schedule.tiling_structure": "SSRSRS"})
                              with T.init():
                                  C_global[v_b, v_h, v_i, v_j] = T.float32(0)
                              C_global[v_b, v_h, v_i, v_j] = C_global[v_b, v_h, v_i, v_j] + query_T[v_b, v_h, v_i, v_k] * value_T[v_b, v_h, v_k, v_j]
                  for ax0, ax1, ax2, ax3 in T.grid(1, 12, 128, 64):
                      with T.block("C_global"):
                          v0, v1, v2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                          v3 = T.axis.spatial(128, j_0 * 64 + ax3)
                          T.reads(C_global[v0, v1, v2, v3])
                          T.writes(C[v0, v1, v2, v3])
                          C[v0, v1, v2, v3] = C_global[v0, v1, v2, v3]
      @T.prim_func
      def tbg_2(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.parallel": 288, "meta_schedule.unroll_explicit": 512, "meta_schedule.vectorize": 64})
              value_T = T.alloc_buffer((1, 12, 64, 128))
              for b_0, h_0, i_0, j_0, b_1, h_1, i_1, j_1 in T.grid(1, 1, 1, 2, 1, 6, 2, 8):
                  for ax0, ax1, ax2, ax3 in T.grid(1, 2, 64, 8):
                      with T.block("value_T"):
                          v_b = T.axis.spatial(1, ax0)
                          v_h = T.axis.spatial(12, h_1 * 2 + ax1)
                          v_d = T.axis.spatial(64, ax2)
                          v_l = T.axis.spatial(128, j_0 * 64 + j_1 * 8 + ax3)
                          T.reads(value[v_b, v_l, v_h, v_d])
                          T.writes(value_T[v_b, v_h, v_d, v_l])
                          value_T[v_b, v_h, v_d, v_l] = value[v_b, v_l, v_h, v_d]
                  for k_0, b_2, h_2, i_2, j_2, k_1, b_3, h_3, i_3, j_3 in T.grid(1, 1, 2, 2, 4, 64, 1, 1, 32, 2):
                      with T.block("C"):
                          v_b = T.axis.spatial(1, b_0 + b_1 + b_2 + b_3)
                          v_h = T.axis.spatial(12, h_0 * 12 + h_1 * 2 + h_2 + h_3)
                          v_i = T.axis.spatial(128, i_0 * 128 + i_1 * 64 + i_2 * 32 + i_3)
                          v_j = T.axis.spatial(128, j_0 * 64 + j_1 * 8 + j_2 * 2 + j_3)
                          v_k = T.axis.reduce(64, k_0 * 64 + k_1)
                          T.reads(query[v_b, v_i, v_h, v_k], value_T[v_b, v_h, v_k, v_j])
                          T.writes(C[v_b, v_h, v_i, v_j])
                          T.block_attr({"meta_schedule.tiling_structure": "SSRSRS"})
                          with T.init():
                              C[v_b, v_h, v_i, v_j] = T.float32(0)
                          C[v_b, v_h, v_i, v_j] = C[v_b, v_h, v_i, v_j] + query[v_b, v_i, v_h, v_k] * value_T[v_b, v_h, v_k, v_j]
      # fmt: on
      decision_0 = [
          ("SamplePerfectTile", [1, 1, 1, 1]),
          ("SamplePerfectTile", [1, 6, 2, 1]),
          ("SamplePerfectTile", [1, 2, 2, 32]),
          ("SamplePerfectTile", [2, 8, 4, 2]),
          ("SamplePerfectTile", [1, 64]),
          ("SampleCategorical", 2),
          ("SampleComputeLocation", 6),
          ("SampleComputeLocation", 6),
      ]
      decision_1 = [
          ("SamplePerfectTile", [1, 1, 1, 1]),
          ("SamplePerfectTile", [1, 6, 2, 1]),
          ("SamplePerfectTile", [1, 2, 2, 32]),
          ("SamplePerfectTile", [2, 8, 4, 2]),
          ("SamplePerfectTile", [1, 64]),
          ("SampleCategorical", 2),
          ("SampleComputeLocation", 13),
          ("SampleComputeLocation", -1),
      ]
      decision_2 = [
          ("SamplePerfectTile", [1, 1, 1, 1]),
          ("SamplePerfectTile", [1, 6, 2, 1]),
          ("SamplePerfectTile", [1, 2, 2, 32]),
          ("SamplePerfectTile", [2, 8, 4, 2]),
          ("SamplePerfectTile", [1, 64]),
          ("SampleCategorical", 3),
          ("SampleComputeLocation", 7),
          ("SampleComputeLocation", -2),
      ]
      mod = create_te_workload("TBG", 0)
      actual = _design_space(mod)
      check_sketches(
          mod,
          sketches=actual,
          expected_mods=[tbg_0, tbg_1, tbg_2],
          expected_decisions=[decision_0, decision_1, decision_2],
      )
  
file_path: tvm/tests/python/meta_schedule/test_meta_schedule_space_cpu.py
imports: from tvm import meta_schedule
  from tvm.meta_schedule.testing.space_generation import check_sketches, print_sketches, generate_design_space
  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_cpu_tbg() without return types