code: def test_matmul_relu_with_fallback():
      # fmt: off
      @T.prim_func
      def matmul_relu_fallback_0(A: T.Buffer((128, 128), "float16"), B: T.Buffer((128, 128), "float16"), compute: T.Buffer((128, 128), "float32")) -> None:
          T.func_attr({"global_symbol": "main", "tir.noalias": True})
          # with T.block("root"):
          C_reindex_shared = T.alloc_buffer((4, 2, 2, 4, 16, 16), scope="shared")
          C_reindex_shared_wmma_accumulator = T.alloc_buffer((4, 2, 2, 4, 16, 16), scope="wmma.accumulator")
          A_reindex_shared = T.alloc_buffer((128, 128), "float16", scope="shared")
          B_reindex_shared = T.alloc_buffer((128, 128), "float16", scope="shared")
          A_reindex_shared_wmma_matrix_a = T.alloc_buffer((128, 128), "float16", scope="wmma.matrix_a")
          B_reindex_shared_wmma_matrix_b = T.alloc_buffer((128, 128), "float16", scope="wmma.matrix_b")
          for ax0_0_0_ax1_0_0_fused in T.thread_binding(2, thread="blockIdx.y"):
              for ax0_0_1_ax1_0_1_fused in T.thread_binding(2, thread="blockIdx.x"):
                  for ax0_0_2_ax1_0_2_fused in T.thread_binding(2, thread="threadIdx.y"):
                      for ax2_0_0 in range(2):
                          for ax0_ax1_fused in range(2048):
                              with T.block("A_reindex_shared"):
                                  v0 = T.axis.spatial(128, ax0_0_0_ax1_0_0_fused * 64 + ax0_0_1_ax1_0_1_fused * 32 + ax0_ax1_fused // 64)
                                  v1 = T.axis.spatial(128, ax2_0_0 * 64 + ax0_ax1_fused % 64)
                                  T.reads(A[v0, v1])
                                  T.writes(A_reindex_shared[v0, v1])
                                  T.block_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 4})
                                  A_reindex_shared[v0, v1] = A[v0, v1]
                          for ax0_ax1_fused in range(8192):
                              with T.block("B_reindex_shared"):
                                  v0 = T.axis.spatial(128, ax2_0_0 * 64 + ax0_ax1_fused // 128)
                                  v1 = T.axis.spatial(128, ax0_ax1_fused % 128)
                                  T.reads(B[v0, v1])
                                  T.writes(B_reindex_shared[v0, v1])
                                  T.block_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 2})
                                  B_reindex_shared[v0, v1] = B[v0, v1]
                          for ax2_0_1 in range(1):
                              for ax0_0, ax1_0 in T.grid(2, 4):
                                  with T.block("A_reindex_shared_wmma.matrix_a_o"):
                                      v0_o = T.axis.spatial(8, ax0_0_0_ax1_0_0_fused * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_0)
                                      v1_o = T.axis.spatial(8, ax2_0_0 * 4 + ax1_0)
                                      T.reads(A_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
                                      T.writes(A_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
                                      T.block_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_a_shared"})
                                      for ax0_1, ax1_1 in T.grid(16, 16):
                                          with T.block("A_reindex_shared_wmma.matrix_a"):
                                              v0_i, v1_i = T.axis.remap("SS", [ax0_1, ax1_1])
                                              T.reads(A_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
                                              T.writes(A_reindex_shared_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
                                              A_reindex_shared_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = A_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i]
                              for ax0_0, ax1_0 in T.grid(4, 4):
                                  with T.block("B_reindex_shared_wmma.matrix_b_o"):
                                      v0_o = T.axis.spatial(8, ax2_0_0 * 4 + ax0_0)
                                      v1_o = T.axis.spatial(8, ax0_0_2_ax1_0_2_fused * 4 + ax1_0)
                                      T.reads(B_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
                                      T.writes(B_reindex_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
                                      T.block_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_b_shared"})
                                      for ax0_1, ax1_1 in T.grid(16, 16):
                                          with T.block("B_reindex_shared_wmma.matrix_b"):
                                              v0_i, v1_i = T.axis.remap("SS", [ax0_1, ax1_1])
                                              T.reads(B_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
                                              T.writes(B_reindex_shared_wmma_matrix_b[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
                                              B_reindex_shared_wmma_matrix_b[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = B_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i]
                              for ax0_0_3, ax1_0_3, ax2_0_2, ax0_0_4, ax1_0_4 in T.grid(1, 1, 4, 2, 4):
                                  with T.block("C_o"):
                                      v0_o = T.axis.spatial(8, ax0_0_0_ax1_0_0_fused * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_0_3 * 2 + ax0_0_4)
                                      v1_o = T.axis.spatial(8, ax0_0_2_ax1_0_2_fused * 4 + ax1_0_3 * 4 + ax1_0_4)
                                      v2_o = T.axis.reduce(8, ax2_0_0 * 4 + ax2_0_1 * 4 + ax2_0_2)
                                      T.reads(A_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], B_reindex_shared_wmma_matrix_b[v2_o * 16:v2_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
                                      T.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, 0:16, 0:16])
                                      T.block_attr({"meta_schedule.auto_tensorize": "wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init": "wmma_fill_16x16x16_f32", "warp_execution": 1})
                                      with T.init():
                                          for ax0_1, ax1_1 in T.grid(16, 16):
                                              with T.block("C_init"):
                                                  v0_i_init, v1_i_init = T.axis.remap("SS", [ax0_1, ax1_1])
                                                  T.reads()
                                                  T.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, v0_i_init, v1_i_init])
                                                  C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, v0_i_init, v1_i_init] = T.float32(0)
                                      for ax0_1, ax1_1, ax2_1 in T.grid(16, 16, 16):
                                          with T.block("C"):
                                              v0_i, v1_i, v2_i = T.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
                                              T.reads(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, v0_i, v1_i], A_reindex_shared_wmma_matrix_a[v0_o * 16 + v0_i, v2_o * 16 + v2_i], B_reindex_shared_wmma_matrix_b[v2_o * 16 + v2_i, v1_o * 16 + v1_i])
                                              T.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, v0_i, v1_i])
                                              T.block_attr({"meta_schedule.tiling_structure": "SSSRRSRS"})
                                              C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, v0_i, v1_i] = C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, v0_i, v1_i] + T.Cast("float32", A_reindex_shared_wmma_matrix_a[v0_o * 16 + v0_i, v2_o * 16 + v2_i]) * T.Cast("float32", B_reindex_shared_wmma_matrix_b[v2_o * 16 + v2_i, v1_o * 16 + v1_i])
                  for ax2 in range(2):
                      for ax0_ax1_fused in T.thread_binding(2, thread="threadIdx.y"):
                          for ax2_1, ax3 in T.grid(1, 4):
                              with T.block("C_reindex_shared_wmma.accumulator_o"):
                                  v0 = T.axis.spatial(4, ax0_0_0_ax1_0_0_fused * 2 + ax0_0_1_ax1_0_1_fused)
                                  v1 = T.axis.spatial(2, ax0_ax1_fused)
                                  v2 = T.axis.spatial(2, ax2 + ax2_1)
                                  v3 = T.axis.spatial(4, ax3)
                                  v4_o = T.axis.spatial(1, 0)
                                  v5_o = T.axis.spatial(1, 0)
                                  T.reads(C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, 0:16, 0:16])
                                  T.writes(C_reindex_shared[v0, v1, v2, v3, 0:16, 0:16])
                                  T.block_attr({"meta_schedule.auto_tensorize": "wmma_store_16x16x16_f32_shared"})
                                  for ax4, ax5 in T.grid(16, 16):
                                      with T.block("C_reindex_shared_wmma.accumulator"):
                                          v4_i, v5_i = T.axis.remap("SS", [ax4, ax5])
                                          T.reads(C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, v4_i, v5_i])
                                          T.writes(C_reindex_shared[v0, v1, v2, v3, v4_i, v5_i])
                                          C_reindex_shared[v0, v1, v2, v3, v4_i, v5_i] = C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, v4_i, v5_i]
                      for ax0_ax1_ax3_ax4_ax5_fused in range(2048):
                          with T.block("C_reindex_shared"):
                              v0 = T.axis.spatial(4, ax0_0_0_ax1_0_0_fused * 2 + ax0_0_1_ax1_0_1_fused)
                              v1 = T.axis.spatial(2, ax0_ax1_ax3_ax4_ax5_fused // 1024)
                              v2 = T.axis.spatial(2, ax2)
                              v3 = T.axis.spatial(4, ax0_ax1_ax3_ax4_ax5_fused % 1024 // 256)
                              v4 = T.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 256 // 16)
                              v5 = T.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
                              T.reads(C_reindex_shared[v0, v1, v2, v3, v4, v5])
                              T.writes(compute[v4 + v2 * 16 + v0 * 32, v5 + v3 * 16 + v1 * 64])
                              T.block_attr({"meta_schedule.cooperative_fetch": 4})
                              compute[v4 + v2 * 16 + v0 * 32, v5 + v3 * 16 + v1 * 64] = T.max(C_reindex_shared[v0, v1, v2, v3, v4, v5], T.float32(0))
      # fmt: on
      decision_0 = [
          ("SamplePerfectTile", [2, 2, 1, 1, 2]),
          ("SamplePerfectTile", [1, 1, 2, 1, 4]),
          ("SamplePerfectTile", [2, 1, 4]),
          ("SampleCategorical", 3),
          ("SampleCategorical", 2),
          ("SampleCategorical", 1),
      ]
  
      mod = te.create_prim_func(
          te_workload.matmul_relu(
              n=128,
              m=128,
              k=128,
              in_dtype="float16",
              out_dtype="float32",
          )
      )
      actual = generate_design_space(
          kind="cuda",
          mod=mod,
          target=tvm.target.Target("cuda --arch=sm_70"),
          types=None,
          sch_rules=[
              multi_level_tiling_tensor_core(),
          ]
          + get_rules(
              "cuda",
              (
                  ms.schedule_rule.MultiLevelTiling,
                  ms.schedule_rule.AutoInline,
              ),
          ),
      )
      check_sketches(
          mod,
          sketches=actual,
          expected_mods=[matmul_relu_fallback_0],
          expected_decisions=[decision_0],
      )
  
file_path: tvm/tests/python/meta_schedule/test_meta_schedule_schedule_rule_mlt_tc.py
imports: import pytest
  import tvm
  import tvm.testing
  from tvm import meta_schedule
  from tvm import te
  from tvm.meta_schedule.testing import te_workload
  from tvm.meta_schedule.testing.space_generation import check_sketches, generate_design_space, get_rules, print_sketches
  from tvm.script import tir
  from tvm.tir.tensor_intrin.cuda import get_wmma_intrin_group
label_desc: Use image node_id 3 for calling a global function with example usage: test_matmul_relu_with_fallback() without return types