code: def test_cpu_gmm():
      # fmt: off
      @T.prim_func
      def gmm_0(X: T.Buffer((1, 128, 128), "float32"), Y: T.Buffer((1, 128, 128), "float32"), Z: T.Buffer((1, 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": 16, "meta_schedule.vectorize": 64})
              Z_global = T.alloc_buffer((1, 128, 128))
              for b_0, i_0, j_0, b_1, i_1, j_1 in T.grid(1, 4, 2, 1, 1, 8):
                  for k_0, b_2, i_2, j_2, k_1, b_3, i_3, j_3 in T.grid(128, 1, 16, 1, 1, 1, 2, 8):
                      with T.block("Z"):
                          v_b = T.axis.spatial(1, b_0 + b_1 + b_2 + b_3)
                          v_i = T.axis.spatial(128, i_0 * 32 + i_1 * 32 + i_2 * 2 + i_3)
                          v_j = T.axis.spatial(128, j_0 * 64 + j_1 * 8 + j_2 * 8 + j_3)
                          v_k = T.axis.reduce(128, k_0 + k_1)
                          T.reads(X[v_b, v_i, v_k], Y[v_b, v_k, v_j])
                          T.writes(Z_global[v_b, v_i, v_j])
                          T.block_attr({"meta_schedule.tiling_structure": "SSRSRS"})
                          with T.init():
                              Z_global[v_b, v_i, v_j] = T.float32(0)
                          Z_global[v_b, v_i, v_j] = Z_global[v_b, v_i, v_j] + X[v_b, v_i, v_k] * Y[v_b, v_k, v_j]
                  for ax0, ax1, ax2 in T.grid(1, 32, 8):
                      with T.block("Z_global"):
                          v0 = T.axis.spatial(1, ax0)
                          v1 = T.axis.spatial(128, i_0 * 32 + ax1)
                          v2 = T.axis.spatial(128, j_0 * 64 + j_1 * 8 + ax2)
                          T.reads(Z_global[v0, v1, v2])
                          T.writes(Z[v0, v1, v2])
                          Z[v0, v1, v2] = Z_global[v0, v1, v2]
      @T.prim_func
      def gmm_1(X: T.Buffer((1, 128, 128), "float32"), Y: T.Buffer((1, 128, 128), "float32"), Z: T.Buffer((1, 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": 16, "meta_schedule.vectorize": 64})
              Z_global = T.alloc_buffer((1, 128, 128))
              for b_0, i_0, j_0 in T.grid(1, 4, 2):
                  for b_1, i_1, j_1, k_0, b_2, i_2, j_2, k_1, b_3, i_3, j_3 in T.grid(1, 1, 8, 128, 1, 16, 1, 1, 1, 2, 8):
                      with T.block("Z"):
                          v_b = T.axis.spatial(1, b_0 + b_1 + b_2 + b_3)
                          v_i = T.axis.spatial(128, i_0 * 32 + i_1 * 32 + i_2 * 2 + i_3)
                          v_j = T.axis.spatial(128, j_0 * 64 + j_1 * 8 + j_2 * 8 + j_3)
                          v_k = T.axis.reduce(128, k_0 + k_1)
                          T.reads(X[v_b, v_i, v_k], Y[v_b, v_k, v_j])
                          T.writes(Z_global[v_b, v_i, v_j])
                          T.block_attr({"meta_schedule.tiling_structure": "SSRSRS"})
                          with T.init():
                              Z_global[v_b, v_i, v_j] = T.float32(0)
                          Z_global[v_b, v_i, v_j] = Z_global[v_b, v_i, v_j] + X[v_b, v_i, v_k] * Y[v_b, v_k, v_j]
                  for ax0, ax1, ax2 in T.grid(1, 32, 64):
                      with T.block("Z_global"):
                          v0 = T.axis.spatial(1, ax0)
                          v1 = T.axis.spatial(128, i_0 * 32 + ax1)
                          v2 = T.axis.spatial(128, j_0 * 64 + ax2)
                          T.reads(Z_global[v0, v1, v2])
                          T.writes(Z[v0, v1, v2])
                          Z[v0, v1, v2] = Z_global[v0, v1, v2]
      @T.prim_func
      def gmm_2(X: T.Buffer((1, 128, 128), "float32"), Y: T.Buffer((1, 128, 128), "float32"), Z: T.Buffer((1, 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": 16, "meta_schedule.vectorize": 64})
              for b_0, i_0, j_0, b_1, i_1, j_1, k_0, b_2, i_2, j_2, k_1, b_3, i_3, j_3 in T.grid(1, 4, 2, 1, 1, 8, 128, 1, 16, 1, 1, 1, 2, 8):
                  with T.block("Z"):
                      v_b = T.axis.spatial(1, b_0 + b_1 + b_2 + b_3)
                      v_i = T.axis.spatial(128, i_0 * 32 + i_1 * 32 + i_2 * 2 + i_3)
                      v_j = T.axis.spatial(128, j_0 * 64 + j_1 * 8 + j_2 * 8 + j_3)
                      v_k = T.axis.reduce(128, k_0 + k_1)
                      T.reads(X[v_b, v_i, v_k], Y[v_b, v_k, v_j])
                      T.writes(Z[v_b, v_i, v_j])
                      T.block_attr({"meta_schedule.tiling_structure": "SSRSRS"})
                      with T.init():
                          Z[v_b, v_i, v_j] = T.float32(0)
                      Z[v_b, v_i, v_j] = Z[v_b, v_i, v_j] + X[v_b, v_i, v_k] * Y[v_b, v_k, v_j]
      # fmt: on
      decision_0 = [
          ("SamplePerfectTile", [1, 1, 1, 1]),
          ("SamplePerfectTile", [4, 1, 16, 2]),
          ("SamplePerfectTile", [2, 8, 1, 8]),
          ("SamplePerfectTile", [128, 1]),
          ("SampleCategorical", 1),
      ]
      decision_1 = [
          ("SamplePerfectTile", [1, 1, 1, 1]),
          ("SamplePerfectTile", [4, 1, 16, 2]),
          ("SamplePerfectTile", [2, 8, 1, 8]),
          ("SamplePerfectTile", [128, 1]),
          ("SampleCategorical", 1),
      ]
      decision_2 = [
          ("SamplePerfectTile", [1, 1, 1, 1]),
          ("SamplePerfectTile", [4, 1, 16, 2]),
          ("SamplePerfectTile", [2, 8, 1, 8]),
          ("SamplePerfectTile", [128, 1]),
          ("SampleCategorical", 1),
      ]
      mod = create_te_workload("GMM", 0)
      actual = _design_space(mod)
      check_sketches(
          mod,
          sketches=actual,
          expected_mods=[gmm_0, gmm_1, gmm_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 9 for calling a global function with example usage: test_cpu_gmm() without return types