code: def test_conv2d(shared_scope):
      intrin_suffix = shared_scope.replace(".", "_")
      # fmt: off
      @T.prim_func
      def conv2d_0(inputs: T.Buffer((1, 16, 16, 32), "float16"), weight: T.Buffer((3, 3, 32, 32), "float16"), conv2d_nhwc: T.Buffer((1, 16, 16, 32), "float32")):
          T.func_attr({"global_symbol": "main", "tir.noalias": T.bool(True)})
          # with T.block("root"):
          PadInput = T.alloc_buffer((1, 18, 18, 32), "float16")
          conv2d_nhwc_reindex_shared_dyn = T.alloc_buffer((16, 2, 1, 1, 16, 16), scope=shared_scope)
          conv2d_nhwc_reindex_shared_dyn_wmma_accumulator = T.alloc_buffer((16, 2, 1, 1, 16, 16), scope="wmma.accumulator")
          PadInput_reindex_shared_dyn = T.alloc_buffer((256, 288), "float16", scope=shared_scope)
          weight_reindex_shared_dyn = T.alloc_buffer((288, 32), "float16", scope=shared_scope)
          PadInput_reindex_shared_dyn_wmma_matrix_a = T.alloc_buffer((256, 288), "float16", scope="wmma.matrix_a")
          weight_reindex_shared_dyn_wmma_matrix_b = T.alloc_buffer((288, 32), "float16", scope="wmma.matrix_b")
          for i0, i1, i2, i3 in T.grid(1, 18, 18, 32):
              with T.block("PadInput"):
                  v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                  T.reads(inputs[v_i0, v_i1 - 1, v_i2 - 1, v_i3])
                  T.writes(PadInput[v_i0, v_i1, v_i2, v_i3])
                  PadInput[v_i0, v_i1, v_i2, v_i3] = T.if_then_else(1 <= v_i1 and v_i1 < 17 and 1 <= v_i2 and v_i2 < 17, inputs[v_i0, v_i1 - 1, v_i2 - 1, v_i3], T.float16(0))
          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(16, thread="blockIdx.x"):
                  for ax0_0_2_ax1_0_2_fused in T.thread_binding(1, thread="threadIdx.y"):
                      for ax2_0_0 in range(1):
                          for ax0_ax1_fused in range(4608):
                              with T.block("PadInput_reindex_shared.dyn"):
                                  v0 = T.axis.spatial(256, ax0_0_1_ax1_0_1_fused * 16 + ax0_ax1_fused // 288)
                                  v1 = T.axis.spatial(288, ax0_ax1_fused % 288)
                                  T.reads(PadInput[0, v0 // 16 + v1 // 96, v0 % 16 + v1 % 96 // 32, v1 % 32])
                                  T.writes(PadInput_reindex_shared_dyn[v0, v1])
                                  T.block_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 2})
                                  PadInput_reindex_shared_dyn[v0, v1] = PadInput[0, v0 // 16 + v1 // 96, v0 % 16 + v1 % 96 // 32, v1 % 32]
                          for ax0_ax1_fused in range(4608):
                              with T.block("weight_reindex_shared.dyn"):
                                  v0 = T.axis.spatial(288, ax0_ax1_fused // 16)
                                  v1 = T.axis.spatial(32, ax0_0_0_ax1_0_0_fused * 16 + ax0_ax1_fused % 16)
                                  T.reads(weight[v0 // 96, v0 % 96 // 32, v0 % 32, v1])
                                  T.writes(weight_reindex_shared_dyn[v0, v1])
                                  T.block_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 8})
                                  weight_reindex_shared_dyn[v0, v1] = weight[v0 // 96, v0 % 96 // 32, v0 % 32, v1]
                          for ax2_0_1 in range(18):
                              for ax0_0, ax1_0 in T.grid(1, 1):
                                  with T.block("PadInput_reindex_shared.dyn_wmma.matrix_a_o"):
                                      v0_o = T.axis.spatial(16, ax0_0_1_ax1_0_1_fused + ax0_0)
                                      v1_o = T.axis.spatial(18, ax2_0_1 + ax1_0)
                                      T.reads(PadInput_reindex_shared_dyn[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
                                      T.writes(PadInput_reindex_shared_dyn_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
                                      T.block_attr({"meta_schedule.auto_tensorize": f"wmma_load_16x16x16_f16_a_{intrin_suffix}"})
                                      for ax0_1, ax1_1 in T.grid(16, 16):
                                          with T.block("PadInput_reindex_shared.dyn_wmma.matrix_a"):
                                              v0_i, v1_i = T.axis.remap("SS", [ax0_1, ax1_1])
                                              T.reads(PadInput_reindex_shared_dyn[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
                                              T.writes(PadInput_reindex_shared_dyn_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
                                              PadInput_reindex_shared_dyn_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = PadInput_reindex_shared_dyn[v0_o * 16 + v0_i, v1_o * 16 + v1_i]
                              for ax0_0, ax1_0 in T.grid(1, 1):
                                  with T.block("weight_reindex_shared.dyn_wmma.matrix_b_o"):
                                      v0_o = T.axis.spatial(18, ax2_0_1 + ax0_0)
                                      v1_o = T.axis.spatial(2, ax0_0_0_ax1_0_0_fused + ax1_0)
                                      T.reads(weight_reindex_shared_dyn[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
                                      T.writes(weight_reindex_shared_dyn_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
                                      T.block_attr({"meta_schedule.auto_tensorize": f"wmma_load_16x16x16_f16_b_{intrin_suffix}"})
                                      for ax0_1, ax1_1 in T.grid(16, 16):
                                          with T.block("weight_reindex_shared.dyn_wmma.matrix_b"):
                                              v0_i, v1_i = T.axis.remap("SS", [ax0_1, ax1_1])
                                              T.reads(weight_reindex_shared_dyn[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
                                              T.writes(weight_reindex_shared_dyn_wmma_matrix_b[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
                                              weight_reindex_shared_dyn_wmma_matrix_b[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = weight_reindex_shared_dyn[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, 1, 1, 1):
                                  with T.block("conv2d_nhwc_o"):
                                      v0_o = T.axis.spatial(16, ax0_0_1_ax1_0_1_fused + ax0_0_3 + ax0_0_4)
                                      v1_o = T.axis.spatial(2, ax0_0_0_ax1_0_0_fused + ax1_0_3 + ax1_0_4)
                                      v2_o = T.axis.reduce(18, ax2_0_0 * 18 + ax2_0_1 + ax2_0_2)
                                      T.reads(PadInput_reindex_shared_dyn_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], weight_reindex_shared_dyn_wmma_matrix_b[v2_o * 16:v2_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
                                      T.writes(conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o, 0, 0, 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("conv2d_nhwc_init"):
                                                  v0_i_init, v1_i_init = T.axis.remap("SS", [ax0_1, ax1_1])
                                                  T.reads()
                                                  T.writes(conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o, 0, 0, v0_i_init, v1_i_init])
                                                  conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o, 0, 0, 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("conv2d_nhwc"):
                                              v0_i, v1_i, v2_i = T.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
                                              T.reads(conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o, 0, 0, v0_i, v1_i], PadInput_reindex_shared_dyn_wmma_matrix_a[v0_o * 16 + v0_i, v2_o * 16 + v2_i], weight_reindex_shared_dyn_wmma_matrix_b[v2_o * 16 + v2_i, v1_o * 16 + v1_i])
                                              T.writes(conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o, 0, 0, v0_i, v1_i])
                                              T.block_attr({"meta_schedule.tiling_structure": "SSSRRSRS"})
                                              conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o, 0, 0, v0_i, v1_i] = conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o, 0, 0, v0_i, v1_i] + T.Cast("float32", PadInput_reindex_shared_dyn_wmma_matrix_a[v0_o * 16 + v0_i, v2_o * 16 + v2_i]) * T.Cast("float32", weight_reindex_shared_dyn_wmma_matrix_b[v2_o * 16 + v2_i, v1_o * 16 + v1_i])
                  for ax2 in range(1):
                      for ax0_ax1_fused in T.thread_binding(1, thread="threadIdx.y"):
                          for ax2_1, ax3 in T.grid(1, 1):
                              with T.block("conv2d_nhwc_reindex_shared.dyn_wmma.accumulator_o"):
                                  v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0_0_1_ax1_0_1_fused, ax0_0_0_ax1_0_0_fused, ax2_1, ax3])
                                  v4_o = T.axis.spatial(1, 0)
                                  v5_o = T.axis.spatial(1, 0)
                                  T.reads(conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0, v1, v2, v3, 0:16, 0:16])
                                  T.writes(conv2d_nhwc_reindex_shared_dyn[v0, v1, v2, v3, 0:16, 0:16])
                                  T.block_attr({"meta_schedule.auto_tensorize": f"wmma_store_16x16x16_f32_{intrin_suffix}"})
                                  for ax4, ax5 in T.grid(16, 16):
                                      with T.block("conv2d_nhwc_reindex_shared.dyn_wmma.accumulator"):
                                          v4_i, v5_i = T.axis.remap("SS", [ax4, ax5])
                                          T.reads(conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0, v1, v2, v3, v4_i, v5_i])
                                          T.writes(conv2d_nhwc_reindex_shared_dyn[v0, v1, v2, v3, v4_i, v5_i])
                                          conv2d_nhwc_reindex_shared_dyn[v0, v1, v2, v3, v4_i, v5_i] = conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0, v1, v2, v3, v4_i, v5_i]
                      for ax0_ax1_ax3_ax4_ax5_fused in range(256):
                          with T.block("conv2d_nhwc_reindex_shared.dyn"):
                              v0, v1, v2 = T.axis.remap("SSS", [ax0_0_1_ax1_0_1_fused, ax0_0_0_ax1_0_0_fused, ax2])
                              v3 = T.axis.spatial(1, 0)
                              v4 = T.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused // 16)
                              v5 = T.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
                              T.reads(conv2d_nhwc_reindex_shared_dyn[v0, v1, v2, v3, v4, v5])
                              T.writes(conv2d_nhwc[0, (v4 + v0 * 16) // 16, (v4 + v0 * 16) % 16, v5 + v1 * 16])
                              T.block_attr({"meta_schedule.cooperative_fetch": 3})
                              conv2d_nhwc[0, (v4 + v0 * 16) // 16, (v4 + v0 * 16) % 16, v5 + v1 * 16] = conv2d_nhwc_reindex_shared_dyn[v0, v1, v2, v3, v4, v5]
      # fmt: on
      decision_0 = [
          ("SamplePerfectTile", [1, 16, 1, 1, 1]),
          ("SamplePerfectTile", [2, 1, 1, 1, 1]),
          ("SamplePerfectTile", [1, 18, 1]),
          ("SampleCategorical", 2),
          ("SampleCategorical", 1),
          ("SampleCategorical", 3),
      ]
      mod = te.create_prim_func(
          te_workload.conv2d_nhwc(
              N=1,
              H=16,
              W=16,
              CI=32,
              CO=32,
              kernel_size=3,
              stride=1,
              padding=1,
              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(
                  read_reuse_scope=shared_scope,
                  write_reuse_scope=shared_scope,
              ),
          ],
      )
      check_sketches(
          mod,
          sketches=actual,
          expected_mods=[conv2d_0],
          expected_decisions=[decision_0],
      )
  
      # Test adding inapplicable tensor intrinsics doesn't change the search space
      # This test case uses the same workload, decision and the expected sketch as above
      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(
                  read_reuse_scope=shared_scope,
                  write_reuse_scope=shared_scope,
                  in_dtype="float16",
                  out_dtype=["float16", "float32"],
              ),
          ],
      )
      check_sketches(
          mod,
          sketches=actual,
          expected_mods=[conv2d_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 4 for calling a global function with example usage: test_conv2d(shared_scope) without return types