You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/09 19:44:43 UTC
(tvm) branch main updated: [REFACTOR][CUDA] Retire legacy WMMA tensorization and emission (#20631)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new c9c2d13754 [REFACTOR][CUDA] Retire legacy WMMA tensorization and emission (#20631)
c9c2d13754 is described below
commit c9c2d1375484140093b1485113b38deb6ac64f0c
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Fri Oct 9 15:44:37 2026 -0400
[REFACTOR][CUDA] Retire legacy WMMA tensorization and emission (#20631)
Remove the legacy CUDA WMMA fragment emitter and tensor-intrinsic
catalogue, together with DLight's WMMA scheduling paths and
MetaSchedule's default WMMA groups. DLight retains its generic CUDA
matmul schedule, and MetaSchedule retains PTX MMA and ordinary CUDA
scheduling.
Shared GPU operations, fragment inference, Vulkan consumers, and modern
CUDA tensor-core paths remain available.
---
python/tvm/s_tir/dlight/gpu/matmul.py | 470 +-----
python/tvm/s_tir/tensor_intrin/cuda.py | 614 --------
src/backend/cuda/codegen/codegen_cuda.cc | 184 +--
src/backend/cuda/codegen/codegen_cuda.h | 3 -
.../meta_schedule/schedule_rule/schedule_rule.cc | 62 -
.../s_tir/dlight/test_gpu_matmul_tensorize.py | 672 +--------
.../test_meta_schedule_postproc_verify_gpu_code.py | 378 +----
.../test_meta_schedule_schedule_rule_mlt_tc.py | 1511 --------------------
.../test_meta_schedule_trace_apply.py | 544 -------
.../s_tir/schedule/test_tir_schedule_analysis.py | 18 +-
.../script/test_tirx_script_source_locations.py | 4 +-
11 files changed, 26 insertions(+), 4434 deletions(-)
diff --git a/python/tvm/s_tir/dlight/gpu/matmul.py b/python/tvm/s_tir/dlight/gpu/matmul.py
index 679ebe4343..4ad2a45d7a 100644
--- a/python/tvm/s_tir/dlight/gpu/matmul.py
+++ b/python/tvm/s_tir/dlight/gpu/matmul.py
@@ -327,21 +327,6 @@ def get_reduction_blocks(sch, blocks) -> bool:
return reduction_blocks
-def get_in_out_dtypes(block: s_tir.SBlock) -> tuple[str]:
- """
- Detect In/Out data types for the given block based on the analysis if read/write buffers.
- """
- assert len(block.reads) > 0 and len(block.writes) > 0
- in_dtype = block.reads[0].source.dtype
- out_dtype = block.writes[0].source.dtype
- return (in_dtype, out_dtype)
-
-
-def check_sm_version(arch: str) -> int:
- sm_version = arch.replace("sm_", "")
- return int(sm_version) if sm_version.isdigit() else -1
-
-
class MetalMatmul(GPUScheduleRule):
"""
The schedule rule for Metal matmul computation.
@@ -435,7 +420,7 @@ class MetalMatmul(GPUScheduleRule):
auto_inline_producers(sch, a_g2s)
auto_inline_producers(sch, b_g2s)
- # create read cache to load matrix from shared memory to wmma fragments
+ # create read cache to load matrix from shared memory to SIMD-group fragments
A_simdgroup = sch.cache_read(main_block, 0, "metal.simdgroup")
B_simdgroup = sch.cache_read(main_block, 1, "metal.simdgroup")
sch.compute_at(A_simdgroup, k1)
@@ -482,432 +467,6 @@ class MetalMatmul(GPUScheduleRule):
return sch
-class MatmulTensorization(GPUScheduleRule):
- """
- The schedule rule for float16 tensor core matmul computation.
- func with attr 'dlight.do_not_tensorize' will not be tensorized.
- """
-
- def apply( # pylint: disable=too-many-locals,missing-docstring
- self,
- func: tirx.Function,
- target: Target,
- _: bool,
- ) -> s_tir.Schedule | None:
- from tvm.s_tir.tensor_intrin.cuda import ( # pylint: disable=import-outside-toplevel
- get_wmma_intrin_group,
- )
-
- if not isinstance(func, tirx.Function) or not self.is_target_available(target):
- return None
- sch = s_tir.Schedule(func)
- root_block = get_root_block(sch)
- blocks = sch.get_child_blocks(root_block)
-
- if "dlight.do_not_tensorize" in func.attrs.keys():
- return None
-
- reduction_blocks = get_reduction_blocks(sch, blocks)
- if reduction_blocks is None:
- return None
-
- # Start Schedule
- # Step 0. Get schedule config.
- # NOTE: we can analyze the config by the hardware spec in the future
-
- # tensor core intrinsic size
- micro_size_x = 16
- micro_size_y = 16
- micro_size_k = 16
-
- warp_size = 32
- vector_size = 4
-
- i_factors, j_factors, k_factors = (
- [None, 1, 4, 2],
- [1, None, 4, 2],
- [None, 4],
- )
-
- num_ty = i_factors[2] * j_factors[2]
- x_pad_factor = i_factors[2] * i_factors[3]
- y_pad_factor = j_factors[2] * j_factors[3]
- k_pad_factor = k_factors[1]
-
- # Step 1. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K]
- # Reindex first and than analyze the index map
- main_block = reduction_blocks[0]
- reindex_a = sch.reindex(main_block, ("read", 0))
- reindex_b = sch.reindex(main_block, ("read", 1))
- reindex_c = sch.reindex(main_block, ("write", 0))
-
- index_maps = get_index_map(sch.get(main_block))
- assert index_maps is not None
- matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps
-
- sch.transform_layout(reindex_a, ("write", 0), a_index_map)
- sch.transform_layout(reindex_b, ("write", 0), b_index_map)
- sch.transform_layout(reindex_c, ("read", 0), c_index_map)
- sch.transform_block_layout(main_block, matmul_index_map)
-
- # Step 2. Padding for dynamic shape kernels
- sch.pad_einsum(
- main_block,
- [
- 1,
- micro_size_x * x_pad_factor,
- micro_size_y * y_pad_factor,
- micro_size_k * k_pad_factor,
- ],
- )
-
- # Step 3. Schedule matmul to use tensor core
- block = main_block
-
- batch, i, j, k = sch.get_loops(block)
-
- # inner loops for tensor core computation
- i, i_inner = sch.split(i, factors=[None, micro_size_x])
- j, j_inner = sch.split(j, factors=[None, micro_size_y])
- k, k_inner = sch.split(k, factors=[None, micro_size_k])
-
- sch.reorder(i, j, k, i_inner, j_inner, k_inner)
-
- block_inner = block
- block_outer = sch.blockize(i_inner)
-
- i0, i1, i2, i3 = sch.split(i, factors=i_factors)
- j0, j1, j2, j3 = sch.split(j, factors=j_factors)
- k0, k1 = sch.split(k, k_factors)
- sch.annotate(k0, "software_pipeline_order", [0, 3, 1, 4, 5, 2, 6])
- sch.annotate(k0, "software_pipeline_stage", [0, 0, 0, 0, 0, 1, 1])
- sch.annotate(k1, "software_pipeline_order", [0, 1, 2])
- sch.annotate(k1, "software_pipeline_stage", [0, 0, 1])
-
- sch.reorder(i0, j0, i1, j1, j2, i2, k0, k1, i3, j3)
-
- block_idx = sch.fuse(i0, j0)
- block_idy = sch.fuse(i1, j1)
- thread_idy = sch.fuse(j2, i2)
- sch.bind(batch, "blockIdx.z")
- sch.bind(block_idx, "blockIdx.x")
- sch.bind(block_idy, "blockIdx.y")
- sch.bind(thread_idy, "threadIdx.y")
-
- def fetch_to_shared(block, idx, ndim):
- block_read = sch.cache_read(block, idx, "shared.dyn")
- sch.compute_at(block_read, k0)
- fused = sch.fuse(*sch.get_loops(block_read)[-ndim:])
-
- _, f_1, f_2, f_3 = sch.split(fused, factors=[None, num_ty, warp_size, vector_size])
-
- sch.bind(f_2, "threadIdx.x")
- sch.bind(f_1, "threadIdx.y")
- sch.vectorize(f_3)
-
- sch.storage_align(block_read, 0, axis=-2, factor=16, offset=8)
- sch.annotate(block_read, "tirx.manifest_shared_memory_local_stage", 1)
- sch.annotate(block_read, "double_buffer_scope", 0)
- return block_read
-
- a_g2s = fetch_to_shared(block_outer, 0, 2)
- b_g2s = fetch_to_shared(block_outer, 1, 2)
-
- auto_inline_producers(sch, a_g2s)
- auto_inline_producers(sch, b_g2s)
-
- # create read cache to load matrix from shared memory to wmma fragments
- A_mat = sch.cache_read(block_outer, 0, "wmma.matrix_a")
- B_mat = sch.cache_read(block_outer, 1, "wmma.matrix_b")
- sch.compute_at(A_mat, k1)
- sch.compute_at(B_mat, k1)
-
- # create write cache to store matrix from wmma fragments to shared memory and global memory
- accumulator_shared_to_global = sch.cache_write(block_outer, 0, "shared.dyn")
- sch.storage_align(accumulator_shared_to_global, 0, -2, 16, 4)
-
- store = sch.cache_write(block_outer, 0, "wmma.accumulator")
- sch.reverse_compute_at(store, thread_idy)
- sch.reverse_compute_at(accumulator_shared_to_global, thread_idy)
-
- # split the store loop to match hardware intrinsic pattern
- i, j = sch.get_loops(store)[-2:]
- i0, i1 = sch.split(i, factors=[None, 16])
- j0, j1 = sch.split(j, factors=[None, 16])
- sch.reorder(i0, j0, i1, j1)
-
- block_init_c = sch.decompose_reduction(block_outer, k0)
- block_init_c_inner = sch.get_child_blocks(block_init_c)[0]
-
- # Tensorization by hardware intrinsics
- intrin_group = get_wmma_intrin_group(
- load_scope="shared.dyn",
- store_scope="shared.dyn",
- in_dtype="float16",
- out_dtype="float32",
- trans_b=True,
- )
-
- try:
- i, j = sch.get_loops(A_mat)[-2:]
- i0, i1 = sch.split(i, factors=[None, 16])
- j0, j1 = sch.split(j, factors=[None, 16])
- sch.reorder(i0, j0, i1, j1)
- sch.unroll(i0)
- sch.unroll(j0)
- sch.tensorize(i1, intrin_group["load_a"])
-
- i, j = sch.get_loops(B_mat)[-2:]
- i0, i1 = sch.split(i, factors=[None, 16])
- j0, j1 = sch.split(j, factors=[None, 16])
- sch.reorder(i0, j0, i1, j1)
- sch.unroll(i0)
- sch.unroll(j0)
- sch.tensorize(i1, intrin_group["load_b"])
- except: # pylint: disable=bare-except
- return None
-
- # Try to tensorize the init, store and compute block with f16 or f32 intrinsics
- tensorize_success: bool = False
-
- def tensorize_init_store_compute():
- sch.tensorize(sch.get_loops(block_init_c_inner)[-2], intrin_group["init"])
- sch.tensorize(sch.get_loops(store)[-2], intrin_group["store"])
- sch.tensorize(sch.get_loops(block_inner)[-3], intrin_group["compute"])
-
- try:
- tensorize_init_store_compute()
- tensorize_success = True
- except: # pylint: disable=bare-except
- intrin_group = get_wmma_intrin_group(
- load_scope="shared.dyn",
- store_scope="shared.dyn",
- in_dtype="float16",
- out_dtype="float16",
- trans_b=True,
- )
-
- if not tensorize_success:
- try:
- tensorize_init_store_compute()
- tensorize_success = True
- except: # pylint: disable=bare-except
- return None
- auto_inline_consumer_chain(sch, accumulator_shared_to_global)
-
- fused = sch.fuse(*sch.get_loops(accumulator_shared_to_global)[-2:])
- _, f1, f2 = sch.split(fused, factors=[None, warp_size, vector_size])
- sch.bind(f1, "threadIdx.x")
- sch.vectorize(f2)
-
- return sch if tensorize_success else None
-
-
-class MatmulInt8Tensorization(GPUScheduleRule):
- """
- The schedule rule for int8 tensor core matmul computation.
- func with attr 'dlight.do_not_tensorize' will not be tensorized.
- """
-
- def apply( # pylint: disable=too-many-locals,missing-docstring
- self,
- func: tirx.Function,
- target: Target,
- _: bool,
- ) -> s_tir.Schedule | None:
- from tvm.s_tir.tensor_intrin.cuda import ( # pylint: disable=import-outside-toplevel
- get_wmma_intrin_group,
- )
-
- if not isinstance(func, tirx.Function) or not self.is_target_available(target):
- return None
- sch = s_tir.Schedule(func)
- root_block = get_root_block(sch)
- blocks = sch.get_child_blocks(root_block)
-
- if "dlight.do_not_tensorize" in func.attrs.keys():
- return None
-
- reduction_blocks = get_reduction_blocks(sch, blocks)
- if reduction_blocks is None:
- return None
-
- # Start Schedule
- # Step 0. Get schedule config.
- # NOTE: we can analyze the config by the hardware spec in the future
-
- # tensor core intrinsic size
- micro_size_x = 16
- micro_size_y = 16
- micro_size_k = 16
-
- warp_size = 32
- vector_size = 4
-
- i_factors, j_factors, k_factors = (
- [None, 1, 4, 2],
- [1, None, 4, 2],
- [None, 1],
- )
-
- num_ty = i_factors[2] * j_factors[2]
- x_pad_factor = i_factors[2] * i_factors[3]
- y_pad_factor = j_factors[2] * j_factors[3]
- k_pad_factor = k_factors[1]
-
- # Step 1. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K]
- # Reindex first and than analyze the index map
- main_block = reduction_blocks[0]
- reindex_a = sch.reindex(main_block, ("read", 0))
- reindex_b = sch.reindex(main_block, ("read", 1))
- reindex_c = sch.reindex(main_block, ("write", 0))
-
- index_maps = get_index_map(sch.get(main_block))
- assert index_maps is not None
- matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps
-
- sch.transform_layout(reindex_a, ("write", 0), a_index_map)
- sch.transform_layout(reindex_b, ("write", 0), b_index_map)
- sch.transform_layout(reindex_c, ("read", 0), c_index_map)
- sch.transform_block_layout(main_block, matmul_index_map)
-
- # Step 2. Padding for dynamic shape kernels
- sch.pad_einsum(
- main_block,
- [
- 1,
- micro_size_x * x_pad_factor,
- micro_size_y * y_pad_factor,
- micro_size_k * k_pad_factor,
- ],
- )
-
- # Step 3. Schedule matmul to use tensor core
- block = main_block
-
- batch, i, j, k = sch.get_loops(block)
-
- # inner loops for tensor core computation
- i, i_inner = sch.split(i, factors=[None, micro_size_x])
- j, j_inner = sch.split(j, factors=[None, micro_size_y])
- k, k_inner = sch.split(k, factors=[None, micro_size_k])
-
- sch.reorder(i, j, k, i_inner, j_inner, k_inner)
-
- block_inner = block
- block_outer = sch.blockize(i_inner)
-
- i0, i1, i2, i3 = sch.split(i, factors=i_factors)
- j0, j1, j2, j3 = sch.split(j, factors=j_factors)
- k0, k1 = sch.split(k, k_factors)
- sch.annotate(k0, "software_pipeline_order", [0, 3, 1, 4, 5, 2, 6])
- sch.annotate(k0, "software_pipeline_stage", [0, 0, 0, 0, 0, 1, 1])
- sch.annotate(k1, "software_pipeline_order", [0, 1, 2])
- sch.annotate(k1, "software_pipeline_stage", [0, 0, 1])
-
- sch.reorder(i0, j0, i1, j1, j2, i2, k0, k1, i3, j3)
-
- block_idx = sch.fuse(i0, j0)
- block_idy = sch.fuse(i1, j1)
- thread_idy = sch.fuse(j2, i2)
- sch.bind(batch, "blockIdx.z")
- sch.bind(block_idx, "blockIdx.x")
- sch.bind(block_idy, "blockIdx.y")
- sch.bind(thread_idy, "threadIdx.y")
-
- def fetch_to_shared(block, idx, ndim):
- block_read = sch.cache_read(block, idx, "shared.dyn")
- sch.compute_at(block_read, k0)
- fused = sch.fuse(*sch.get_loops(block_read)[-ndim:])
-
- _, f_1, f_2, f_3 = sch.split(fused, factors=[None, num_ty, warp_size, vector_size])
-
- sch.bind(f_2, "threadIdx.x")
- sch.bind(f_1, "threadIdx.y")
- sch.vectorize(f_3)
-
- sch.storage_align(block_read, 0, axis=-2, factor=32, offset=16)
- sch.annotate(block_read, "tirx.manifest_shared_memory_local_stage", 1)
- sch.annotate(block_read, "double_buffer_scope", 0)
- return block_read
-
- a_g2s = fetch_to_shared(block_outer, 0, 2)
- b_g2s = fetch_to_shared(block_outer, 1, 2)
-
- auto_inline_producers(sch, a_g2s)
- auto_inline_producers(sch, b_g2s)
-
- # create read cache to load matrix from shared memory to wmma fragments
- A_mat = sch.cache_read(block_outer, 0, "wmma.matrix_a")
- B_mat = sch.cache_read(block_outer, 1, "wmma.matrix_b")
- sch.compute_at(A_mat, k1)
- sch.compute_at(B_mat, k1)
-
- # create write cache to store matrix from wmma fragments to shared memory and global memory
- accumulator_shared_to_global = sch.cache_write(block_outer, 0, "shared.dyn")
- sch.storage_align(accumulator_shared_to_global, 0, -2, 16, 4)
-
- store = sch.cache_write(block_outer, 0, "wmma.accumulator")
- sch.reverse_compute_at(store, thread_idy)
- sch.reverse_compute_at(accumulator_shared_to_global, thread_idy)
-
- # split the store loop to match hardware intrinsic pattern
- i, j = sch.get_loops(store)[-2:]
- i0, i1 = sch.split(i, factors=[None, 16])
- j0, j1 = sch.split(j, factors=[None, 16])
- sch.reorder(i0, j0, i1, j1)
-
- block_init_c = sch.decompose_reduction(block_outer, k0)
- block_init_c_inner = sch.get_child_blocks(block_init_c)[0]
-
- # Tensorization by hardware intrinsics
- intrin_group = get_wmma_intrin_group(
- load_scope="shared.dyn",
- store_scope="shared.dyn",
- in_dtype="int8",
- out_dtype="int32",
- trans_b=True,
- )
-
- try:
- i, j = sch.get_loops(A_mat)[-2:]
- i0, i1 = sch.split(i, factors=[None, 16])
- j0, j1 = sch.split(j, factors=[None, 16])
- sch.reorder(i0, j0, i1, j1)
- sch.unroll(i0)
- sch.unroll(j0)
- sch.tensorize(i1, intrin_group["load_a"])
-
- i, j = sch.get_loops(B_mat)[-2:]
- i0, i1 = sch.split(i, factors=[None, 16])
- j0, j1 = sch.split(j, factors=[None, 16])
- sch.reorder(i0, j0, i1, j1)
- sch.unroll(i0)
- sch.unroll(j0)
- sch.tensorize(i1, intrin_group["load_b"])
- except: # pylint: disable=bare-except
- return None
-
- def tensorize_init_store_compute():
- sch.tensorize(sch.get_loops(block_init_c_inner)[-2], intrin_group["init"])
- sch.tensorize(sch.get_loops(store)[-2], intrin_group["store"])
- sch.tensorize(sch.get_loops(block_inner)[-3], intrin_group["compute"])
-
- try:
- tensorize_init_store_compute()
- except: # pylint: disable=bare-except
- return None
-
- auto_inline_consumer_chain(sch, accumulator_shared_to_global)
-
- fused = sch.fuse(*sch.get_loops(accumulator_shared_to_global)[-2:])
- _, f1, f2 = sch.split(fused, factors=[None, warp_size, vector_size])
- sch.bind(f1, "threadIdx.x")
- sch.vectorize(f2)
-
- return sch
-
-
class Matmul(GPUScheduleRule):
"""The schedule rule for matmul-like computation"""
@@ -1016,31 +575,8 @@ class Matmul(GPUScheduleRule):
sch.transform_layout(reindex_c, ("read", 0), c_index_map)
sch.transform_block_layout(main_block, matmul_index_map)
- # Step 1. Check Tensor Core support
- # Tensorization config:
- # If any value of I, J, K is fixed and less than this threshold,
- # tensorization rule will not be applied.
- minimal_tensorize_threshold = 64
- block_stmt = sch.get(main_block)
- if target.kind.name == "cuda" and check_sm_version(str(target.attrs.get("arch", ""))) >= 70:
- apply_tensorization: bool = True
- # the batch dimension is not taken into consideration.
- for item_var in block_stmt.iter_vars[1:]:
- extent = item_var.dom.extent
- if isinstance(extent, tirx.expr.IntImm):
- if extent.value <= minimal_tensorize_threshold:
- apply_tensorization = False
- if apply_tensorization:
- # Analyze read/write buffers and choose correct tensorizer: int8 or fp16.
- in_dtype, out_dtype = get_in_out_dtypes(block_stmt)
- tensorize_sch = None
- if in_dtype == "int8" and out_dtype == "int32":
- tensorize_sch = MatmulInt8Tensorization().apply(func, target, _)
- elif in_dtype == "float16" and out_dtype in ["float16", "float32"]:
- tensorize_sch = MatmulTensorization().apply(func, target, _)
- if tensorize_sch is not None:
- return tensorize_sch
- elif target.kind.name == "metal":
+ # Step 1. Use the specialized Metal schedule when available.
+ if target.kind.name == "metal":
try:
return MetalMatmul().apply(func, target, _)
except: # pylint: disable=bare-except
diff --git a/python/tvm/s_tir/tensor_intrin/cuda.py b/python/tvm/s_tir/tensor_intrin/cuda.py
index 186ee6a3ad..f985bb93ea 100644
--- a/python/tvm/s_tir/tensor_intrin/cuda.py
+++ b/python/tvm/s_tir/tensor_intrin/cuda.py
@@ -763,620 +763,6 @@ def get_mma_intrin_group(
}
-######## WMMA intrinsics ########
-
-
-def get_wmma_fragment_index(buffer, stride, m_dim, n_dim):
- """Compute wmma fragment index using elem_offset of the buffer"""
- frag_index_m = buffer.elem_offset // stride // m_dim
- frag_index_n = buffer.elem_offset % stride // n_dim
-
- num_fragments_per_row = stride // n_dim
- return frag_index_m * num_fragments_per_row + frag_index_n
-
-
-def get_wmma_load_intrin(
- m_dim: int,
- n_dim: int,
- k_dim: int,
- dtype: str,
- shared_scope: str,
- is_b: bool,
- is_col_major: bool,
-) -> tuple[Function, Function]:
- """Generator of wmma_load intrins"""
- wmma_fragment_scope = f"wmma.matrix_{'b' if is_b else 'a'}"
- layout = "col_major" if is_col_major else "row_major"
-
- frag_m, frag_n = (k_dim, n_dim) if is_b else (m_dim, k_dim)
- if is_col_major:
- frag_m, frag_n = frag_n, frag_m
- offset_factor = frag_n
-
- @Ts.function
- def wmma_load_desc(
- A: T.Tensor(
- (frag_m, frag_n), dtype, align=64, offset_factor=offset_factor, scope=shared_scope
- ),
- C: T.Tensor(
- (frag_m, frag_n),
- dtype,
- align=64,
- offset_factor=offset_factor,
- scope=wmma_fragment_scope,
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(A[0:frag_m, 0:frag_n])
- Ts.writes(C[0:frag_m, 0:frag_n])
- for i, j in T.grid(frag_m, frag_n):
- with Ts.sblock("load"):
- vii, vjj = Ts.axis.remap("SS", [i, j])
- C[vii, vjj] = A[vii, vjj]
-
- s1 = T.dynamic("s1", "int32")
- s0 = T.dynamic("s0", "int32")
- d1 = T.dynamic("d1", "int32")
- d0 = T.dynamic("d0", "int32")
-
- @Ts.function
- def wmma_load_impl(
- A: T.Tensor(
- (frag_m, frag_n),
- dtype,
- align=64,
- offset_factor=offset_factor,
- scope=shared_scope,
- strides=[s1, s0],
- ),
- C: T.Tensor(
- (frag_m, frag_n),
- dtype,
- align=64,
- offset_factor=offset_factor,
- scope=wmma_fragment_scope,
- strides=[d1, d0],
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(A[0:frag_m, 0:frag_n])
- Ts.writes(C[0:frag_m, 0:frag_n])
- T.evaluate(
- T.gpu_load_matrix_sync(
- C.data,
- m_dim,
- n_dim,
- k_dim,
- get_wmma_fragment_index(C, d1, frag_m, frag_n),
- A.access_ptr("r", ptr_type=dtype),
- s1,
- layout,
- )
- )
-
- return wmma_load_desc, wmma_load_impl
-
-
-def get_wmma_fill_intrin(
- m_dim: int, n_dim: int, k_dim: int, dtype: str
-) -> tuple[Function, Function]:
- """Generator of wmma_fill intrins"""
- zero = IntImm("int32", 0).astype(dtype)
- offset_factor = n_dim
-
- @Ts.function
- def wmma_fill_desc(
- C: T.Tensor(
- (m_dim, n_dim), dtype, align=64, offset_factor=offset_factor, scope="wmma.accumulator"
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads()
- Ts.writes(C[0:m_dim, 0:n_dim])
- for i, j in T.grid(m_dim, n_dim):
- with Ts.sblock("init"):
- vii, vjj = Ts.axis.remap("SS", [i, j])
- C[vii, vjj] = zero
-
- d1 = T.dynamic("d1", "int32")
- d0 = T.dynamic("d0", "int32")
-
- @Ts.function
- def wmma_fill_impl(
- C: T.Tensor(
- (m_dim, n_dim),
- dtype,
- align=64,
- offset_factor=offset_factor,
- scope="wmma.accumulator",
- strides=[d1, d0],
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads()
- Ts.writes(C[0:m_dim, 0:n_dim])
- T.evaluate(
- T.gpu_fill_fragment(
- C.data,
- m_dim,
- n_dim,
- k_dim,
- get_wmma_fragment_index(C, d1, m_dim, n_dim),
- T.float32(0),
- )
- )
-
- return wmma_fill_desc, wmma_fill_impl
-
-
-def get_wmma_store_intrin(
- m_dim: int, n_dim: int, k_dim: int, dtype: str, scope: str
-) -> tuple[Function, Function]:
- """Generator of wmma_store intrins"""
- offset_factor = n_dim
-
- @Ts.function
- def wmma_store_desc(
- A: T.Tensor(
- (m_dim, n_dim), dtype, align=64, offset_factor=offset_factor, scope="wmma.accumulator"
- ),
- C: T.Tensor((m_dim, n_dim), dtype, align=64, offset_factor=offset_factor, scope=scope),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(A[0:m_dim, 0:n_dim])
- Ts.writes(C[0:m_dim, 0:n_dim])
- for i, j in T.grid(m_dim, n_dim):
- with Ts.sblock("store"):
- vii, vjj = Ts.axis.remap("SS", [i, j])
- C[vii, vjj] = A[vii, vjj]
-
- s1 = T.dynamic("s1", "int32")
- s0 = T.dynamic("s0", "int32")
- d1 = T.dynamic("d1", "int32")
- d0 = T.dynamic("d0", "int32")
-
- @Ts.function
- def wmma_store_impl(
- A: T.Tensor(
- (m_dim, n_dim),
- dtype,
- align=64,
- offset_factor=offset_factor,
- scope="wmma.accumulator",
- strides=[d1, d0],
- ),
- C: T.Tensor(
- (m_dim, n_dim),
- dtype,
- align=64,
- offset_factor=offset_factor,
- scope=scope,
- strides=[s1, s0],
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(A[0:m_dim, 0:n_dim])
- Ts.writes(C[0:m_dim, 0:n_dim])
- T.evaluate(
- T.gpu_store_matrix_sync(
- A.data,
- m_dim,
- n_dim,
- k_dim,
- get_wmma_fragment_index(A, d1, m_dim, n_dim),
- C.access_ptr("w", ptr_type=dtype),
- s1,
- "row_major",
- )
- )
-
- return wmma_store_desc, wmma_store_impl
-
-
-def get_wmma_sync_intrin(
- m_dim: int, n_dim: int, k_dim: int, in_dtype: str, out_dtype: str, b_transposed: bool
-) -> tuple[Function, Function]:
- """Generator of wmma_sync intrins"""
-
- def maybe_cast(v):
- if in_dtype != out_dtype:
- return Cast(out_dtype, v)
- return v
-
- def maybe_swap(i, j):
- if b_transposed:
- return j, i
- return i, j
-
- b_shape_0, b_shape_1 = maybe_swap(k_dim, n_dim)
-
- A_offset_factor = k_dim
- B_offset_factor = b_shape_1
- out_offset_factor = n_dim
-
- @Ts.function
- def wmma_sync_desc(
- A: T.Tensor(
- (m_dim, k_dim), in_dtype, align=64, offset_factor=A_offset_factor, scope="wmma.matrix_a"
- ),
- B: T.Tensor(
- maybe_swap(k_dim, n_dim),
- in_dtype,
- align=64,
- offset_factor=B_offset_factor,
- scope="wmma.matrix_b",
- ),
- C: T.Tensor(
- (m_dim, n_dim),
- out_dtype,
- align=64,
- offset_factor=out_offset_factor,
- scope="wmma.accumulator",
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:b_shape_0, 0:b_shape_1])
- Ts.writes(C[0:m_dim, 0:n_dim])
- for i, j, k in T.grid(m_dim, n_dim, k_dim):
- with Ts.sblock(""):
- vii, vjj, vkk = Ts.axis.remap("SSR", [i, j, k])
- b_indices = maybe_swap(vkk, vjj)
- C[vii, vjj] = C[vii, vjj] + maybe_cast(A[vii, vkk]) * maybe_cast(
- B[b_indices[0], b_indices[1]]
- )
-
- a1 = T.dynamic("a1", "int32")
- a0 = T.dynamic("a0", "int32")
- b1 = T.dynamic("b1", "int32")
- b0 = T.dynamic("b0", "int32")
- c1 = T.dynamic("c1", "int32")
- c0 = T.dynamic("c0", "int32")
-
- @Ts.function
- def wmma_sync_impl(
- A: T.Tensor(
- (m_dim, k_dim),
- in_dtype,
- align=64,
- offset_factor=A_offset_factor,
- scope="wmma.matrix_a",
- strides=[a1, a0],
- ),
- B: T.Tensor(
- maybe_swap(k_dim, n_dim),
- in_dtype,
- align=64,
- offset_factor=B_offset_factor,
- scope="wmma.matrix_b",
- strides=[b1, b0],
- ),
- C: T.Tensor(
- (m_dim, n_dim),
- out_dtype,
- align=64,
- offset_factor=out_offset_factor,
- scope="wmma.accumulator",
- strides=[c1, c0],
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:b_shape_0, 0:b_shape_1])
- Ts.writes(C[0:m_dim, 0:n_dim])
- T.evaluate(
- T.gpu_mma_sync(
- C.data,
- get_wmma_fragment_index(C, c1, m_dim, n_dim),
- A.data,
- get_wmma_fragment_index(A, a1, m_dim, k_dim),
- B.data,
- get_wmma_fragment_index(B, b1, b_shape_0, b_shape_1),
- C.data,
- get_wmma_fragment_index(C, c1, m_dim, n_dim),
- )
- )
-
- return wmma_sync_desc, wmma_sync_impl
-
-
-WMMA_SYNC_16x16x16_f16f16f32_INTRIN = "wmma_sync_16x16x16_f16f16f32"
-TensorIntrin.register(
- WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
- *get_wmma_sync_intrin(16, 16, 16, "float16", "float32", False),
-)
-
-WMMA_SYNC_16x16x16_f16f16f32_TRANS_INTRIN = "wmma_sync_16x16x16_f16f16f32_trans"
-TensorIntrin.register(
- WMMA_SYNC_16x16x16_f16f16f32_TRANS_INTRIN,
- *get_wmma_sync_intrin(16, 16, 16, "float16", "float32", True),
-)
-
-WMMA_SYNC_16x16x16_f16f16f16_INTRIN = "wmma_sync_16x16x16_f16f16f16"
-TensorIntrin.register(
- WMMA_SYNC_16x16x16_f16f16f16_INTRIN,
- *get_wmma_sync_intrin(16, 16, 16, "float16", "float16", False),
-)
-
-WMMA_SYNC_16x16x16_f16f16f16_TRANS_INTRIN = "wmma_sync_16x16x16_f16f16f16_trans"
-TensorIntrin.register(
- WMMA_SYNC_16x16x16_f16f16f16_TRANS_INTRIN,
- *get_wmma_sync_intrin(16, 16, 16, "float16", "float16", True),
-)
-
-WMMA_SYNC_16x16x16_s8s8s32_INTRIN = "wmma_sync_16x16x16_s8s8s32"
-TensorIntrin.register(
- WMMA_SYNC_16x16x16_s8s8s32_INTRIN, *get_wmma_sync_intrin(16, 16, 16, "int8", "int32", False)
-)
-
-WMMA_SYNC_16x16x16_s8s8s32_TRANS_INTRIN = "wmma_sync_16x16x16_s8s8s32_trans"
-TensorIntrin.register(
- WMMA_SYNC_16x16x16_s8s8s32_TRANS_INTRIN,
- *get_wmma_sync_intrin(16, 16, 16, "int8", "int32", True),
-)
-
-WMMA_SYNC_8x8x32_s4s4s32_TRANS_INTRIN = "wmma_sync_8x8x32_s4s4s32_trans"
-TensorIntrin.register(
- WMMA_SYNC_8x8x32_s4s4s32_TRANS_INTRIN, *get_wmma_sync_intrin(8, 8, 32, "int4", "int32", True)
-)
-
-WMMA_LOAD_16x16x16_F16_A_INTRIN = "wmma_load_16x16x16_f16_a_shared"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_F16_A_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "float16", "shared", False, False),
-)
-
-WMMA_LOAD_16x16x16_F16_A_DYN_INTRIN = "wmma_load_16x16x16_f16_a_shared_dyn"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_F16_A_DYN_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "float16", "shared.dyn", False, False),
-)
-
-WMMA_LOAD_16x16x16_F16_B_INTRIN = "wmma_load_16x16x16_f16_b_shared"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_F16_B_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "float16", "shared", True, False),
-)
-
-WMMA_LOAD_16x16x16_F16_B_DYN_INTRIN = "wmma_load_16x16x16_f16_b_shared_dyn"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_F16_B_DYN_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "float16", "shared.dyn", True, False),
-)
-
-WMMA_LOAD_16x16x16_F16_A_TRANS_INTRIN = "wmma_load_16x16x16_f16_a_trans_shared"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_F16_A_TRANS_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "float16", "shared", False, True),
-)
-
-WMMA_LOAD_16x16x16_F16_A_TRANS_DYN_INTRIN = "wmma_load_16x16x16_f16_a_trans_shared_dyn"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_F16_A_TRANS_DYN_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "float16", "shared.dyn", False, True),
-)
-
-WMMA_LOAD_16x16x16_F16_B_TRANS_INTRIN = "wmma_load_16x16x16_f16_b_trans_shared"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_F16_B_TRANS_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "float16", "shared", True, True),
-)
-
-WMMA_LOAD_16x16x16_F16_B_TRANS_DYN_INTRIN = "wmma_load_16x16x16_f16_b_trans_shared_dyn"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_F16_B_TRANS_DYN_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "float16", "shared.dyn", True, True),
-)
-
-WMMA_LOAD_16x16x16_S8_A_INTRIN = "wmma_load_16x16x16_s8_a_shared"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_S8_A_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "int8", "shared", False, False),
-)
-
-WMMA_LOAD_16x16x16_S8_A_DYN_INTRIN = "wmma_load_16x16x16_s8_a_shared_dyn"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_S8_A_DYN_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "int8", "shared.dyn", False, False),
-)
-
-WMMA_LOAD_16x16x16_S8_B_INTRIN = "wmma_load_16x16x16_s8_b_shared"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_S8_B_INTRIN, *get_wmma_load_intrin(16, 16, 16, "int8", "shared", True, False)
-)
-
-WMMA_LOAD_16x16x16_S8_B_DYN_INTRIN = "wmma_load_16x16x16_s8_b_shared_dyn"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_S8_B_DYN_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "int8", "shared.dyn", True, False),
-)
-
-WMMA_LOAD_16x16x16_S8_A_TRANS_INTRIN = "wmma_load_16x16x16_s8_a_trans_shared"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_S8_A_TRANS_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "int8", "shared", False, True),
-)
-
-WMMA_LOAD_16x16x16_S8_A_TRANS_DYN_INTRIN = "wmma_load_16x16x16_s8_a_trans_shared_dyn"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_S8_A_TRANS_DYN_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "int8", "shared.dyn", False, True),
-)
-
-WMMA_LOAD_16x16x16_S8_B_TRANS_INTRIN = "wmma_load_16x16x16_s8_b_trans_shared"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_S8_B_TRANS_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "int8", "shared", True, True),
-)
-
-WMMA_LOAD_16x16x16_S8_B_TRANS_DYN_INTRIN = "wmma_load_16x16x16_s8_b_trans_shared_dyn"
-TensorIntrin.register(
- WMMA_LOAD_16x16x16_S8_B_TRANS_DYN_INTRIN,
- *get_wmma_load_intrin(16, 16, 16, "int8", "shared.dyn", True, True),
-)
-
-WMMA_LOAD_8x8x32_S4_A_INTRIN = "wmma_load_8x8x32_s4_a_shared"
-TensorIntrin.register(
- WMMA_LOAD_8x8x32_S4_A_INTRIN, *get_wmma_load_intrin(8, 8, 32, "int4", "shared", False, False)
-)
-
-WMMA_LOAD_8x8x32_S4_A_DYN_INTRIN = "wmma_load_8x8x32_s4_a_shared_dyn"
-TensorIntrin.register(
- WMMA_LOAD_8x8x32_S4_A_DYN_INTRIN,
- *get_wmma_load_intrin(8, 8, 32, "int4", "shared.dyn", False, False),
-)
-
-WMMA_LOAD_8x8x32_S4_B_TRANS_INTRIN = "wmma_load_8x8x32_s4_b_trans_shared"
-TensorIntrin.register(
- WMMA_LOAD_8x8x32_S4_B_TRANS_INTRIN,
- *get_wmma_load_intrin(8, 8, 32, "int4", "shared", True, True),
-)
-
-WMMA_LOAD_8x8x32_S4_B_TRANS_DYN_INTRIN = "wmma_load_8x8x32_s4_b_trans_shared_dyn"
-TensorIntrin.register(
- WMMA_LOAD_8x8x32_S4_B_TRANS_DYN_INTRIN,
- *get_wmma_load_intrin(8, 8, 32, "int4", "shared.dyn", True, True),
-)
-
-WMMA_FILL_16x16x16_F32_INTRIN = "wmma_fill_16x16x16_f32"
-TensorIntrin.register(WMMA_FILL_16x16x16_F32_INTRIN, *get_wmma_fill_intrin(16, 16, 16, "float32"))
-
-WMMA_FILL_16x16x16_F16_INTRIN = "wmma_fill_16x16x16_f16"
-TensorIntrin.register(WMMA_FILL_16x16x16_F16_INTRIN, *get_wmma_fill_intrin(16, 16, 16, "float16"))
-
-WMMA_FILL_16x16x16_S32_INTRIN = "wmma_fill_16x16x16_s32"
-TensorIntrin.register(WMMA_FILL_16x16x16_S32_INTRIN, *get_wmma_fill_intrin(16, 16, 16, "int32"))
-
-WMMA_FILL_8x8x32_S32_INTRIN = "wmma_fill_8x8x32_s32"
-TensorIntrin.register(WMMA_FILL_8x8x32_S32_INTRIN, *get_wmma_fill_intrin(8, 8, 32, "int32"))
-
-WMMA_STORE_16x16x16_F32_SHARED_INTRIN = "wmma_store_16x16x16_f32_shared"
-TensorIntrin.register(
- WMMA_STORE_16x16x16_F32_SHARED_INTRIN, *get_wmma_store_intrin(16, 16, 16, "float32", "shared")
-)
-
-WMMA_STORE_16x16x16_F32_SHARED_DYN_INTRIN = "wmma_store_16x16x16_f32_shared_dyn"
-TensorIntrin.register(
- WMMA_STORE_16x16x16_F32_SHARED_DYN_INTRIN,
- *get_wmma_store_intrin(16, 16, 16, "float32", "shared.dyn"),
-)
-
-WMMA_STORE_16x16x16_F16_SHARED_INTRIN = "wmma_store_16x16x16_f16_shared"
-TensorIntrin.register(
- WMMA_STORE_16x16x16_F16_SHARED_INTRIN, *get_wmma_store_intrin(16, 16, 16, "float16", "shared")
-)
-
-WMMA_STORE_16x16x16_F16_SHARED_DYN_INTRIN = "wmma_store_16x16x16_f16_shared_dyn"
-TensorIntrin.register(
- WMMA_STORE_16x16x16_F16_SHARED_DYN_INTRIN,
- *get_wmma_store_intrin(16, 16, 16, "float16", "shared.dyn"),
-)
-
-WMMA_STORE_16x16x16_S32_SHARED_INTRIN = "wmma_store_16x16x16_s32_shared"
-TensorIntrin.register(
- WMMA_STORE_16x16x16_S32_SHARED_INTRIN, *get_wmma_store_intrin(16, 16, 16, "int32", "shared")
-)
-
-WMMA_STORE_16x16x16_S32_SHARED_DYN_INTRIN = "wmma_store_16x16x16_s32_shared_dyn"
-TensorIntrin.register(
- WMMA_STORE_16x16x16_S32_SHARED_DYN_INTRIN,
- *get_wmma_store_intrin(16, 16, 16, "int32", "shared.dyn"),
-)
-
-WMMA_STORE_8x8x32_S32_SHARED_INTRIN = "wmma_store_8x8x32_s32_shared"
-TensorIntrin.register(
- WMMA_STORE_8x8x32_S32_SHARED_INTRIN, *get_wmma_store_intrin(8, 8, 32, "int32", "shared")
-)
-
-WMMA_STORE_8x8x32_S32_SHARED_DYN_INTRIN = "wmma_store_8x8x32_s32_shared_dyn"
-TensorIntrin.register(
- WMMA_STORE_8x8x32_S32_SHARED_DYN_INTRIN, *get_wmma_store_intrin(8, 8, 32, "int32", "shared.dyn")
-)
-
-WMMA_STORE_16x16x16_F32_GLOBAL_INTRIN = "wmma_store_16x16x16_f32_global"
-TensorIntrin.register(
- WMMA_STORE_16x16x16_F32_GLOBAL_INTRIN, *get_wmma_store_intrin(16, 16, 16, "float32", "global")
-)
-
-WMMA_STORE_16x16x16_F16_GLOBAL_INTRIN = "wmma_store_16x16x16_f16_global"
-TensorIntrin.register(
- WMMA_STORE_16x16x16_F16_GLOBAL_INTRIN, *get_wmma_store_intrin(16, 16, 16, "float16", "global")
-)
-
-WMMA_STORE_16x16x16_S32_GLOBAL_INTRIN = "wmma_store_16x16x16_s32_global"
-TensorIntrin.register(
- WMMA_STORE_16x16x16_S32_GLOBAL_INTRIN, *get_wmma_store_intrin(16, 16, 16, "int32", "global")
-)
-
-WMMA_STORE_8x8x32_S32_GLOBAL_INTRIN = "wmma_store_8x8x32_s32_global"
-TensorIntrin.register(
- WMMA_STORE_8x8x32_S32_GLOBAL_INTRIN, *get_wmma_store_intrin(8, 8, 32, "int32", "global")
-)
-
-
-def get_wmma_intrin_group(
- load_scope: Literal["shared", "shared.dyn"],
- store_scope: Literal["global", "shared", "shared.dyn"],
- in_dtype: str,
- out_dtype: str,
- trans_b: bool,
-) -> dict[str, str]:
- """Get a group of intrinsics for wmma tensor core with the given configurations
-
- Parameters
- ----------
- load_scope : Literal["shared", "shared.dyn"]
- The memory scope of the input buffer.
-
- store_scope : Literal["global", "shared", "shared.dyn"]
- The memory scope of the result buffer.
-
- in_dtype : str
- The input data type.
-
- out_dtype : str
- The output data dtype.
-
- trans_b : bool
- Whether the input matrix B is transposed.
-
- Returns
- -------
- ret : Dict[str, str]
- A group of tensor intrinsics.
- """
- assert load_scope in ["shared", "shared.dyn"]
- assert store_scope in ["global", "shared", "shared.dyn"]
- assert in_dtype in ["float16", "int8"]
- assert out_dtype in ["float16", "float32", "int32"]
-
- shape = "16x16x16"
- in_dtype = "f16" if in_dtype == "float16" else "s8"
- out_dtype = "f16" if out_dtype == "float16" else "f32" if out_dtype == "float32" else "s32"
- # convert "shared.dyn" to "shared_dyn"
- load_scope = load_scope.replace(".", "_")
- store_scope = store_scope.replace(".", "_")
- trans_a = ""
- trans_b = "_trans" if trans_b else ""
-
- # e.g. wmma_load_16x16x16_f16_a_shared
- load_a_intrin = f"wmma_load_{shape}_{in_dtype}_a{trans_a}_{load_scope}"
- # e.g. wmma_load_16x16x16_f16_b_trans_shared_dyn
- load_b_intrin = f"wmma_load_{shape}_{in_dtype}_b{trans_b}_{load_scope}"
- # e.g. wmma_sync_16x16x16_f16f16f32_trans
- compute_intrin = f"wmma_sync_{shape}_{in_dtype}{in_dtype}{out_dtype}{trans_b}"
- # e.g. wmma_fill_16x16x16_f16
- init_intrin = f"wmma_fill_{shape}_{out_dtype}"
- # e.g. wmma_store_16x16x16_f16_shared_dyn
- store_intrin = f"wmma_store_{shape}_{out_dtype}_{store_scope}"
-
- return {
- "init": init_intrin,
- "load_a": load_a_intrin,
- "load_b": load_b_intrin,
- "compute": compute_intrin,
- "store": store_intrin,
- }
-
-
######## MMA intrinsics ########
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc
index 3704226b7d..3f746a2e4f 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -26,7 +26,6 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/op.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/index_map.h>
#include <tvm/tirx/op/gpu.h>
@@ -42,7 +41,6 @@
#include "../../../runtime/thread_storage_scope.h"
#include "../../../target/build_common.h"
-#include "../../../tirx/transform/ir_utils.h"
#include "cuda_fallback_module.h"
#include "literal/cuda_half_t.h"
#include "literal/cuda_int8_t.h"
@@ -1085,11 +1083,6 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
}
}
- static const Op gpu_fill_fragment_op = Op::Get("tirx.gpu_fill_fragment");
- static const Op gpu_load_matrix_sync_op = Op::Get("tirx.gpu_load_matrix_sync");
- static const Op gpu_store_matrix_sync_op = Op::Get("tirx.gpu_store_matrix_sync");
- static const Op gpu_mma_sync_op = Op::Get("tirx.gpu_mma_sync");
- static const Op bmma_sync_op = Op::Get("tirx.cuda.bmma_sync");
static const Op mma_store_op = Op::Get("tirx.mma_store");
static const Op mma_fill_op = Op::Get("tirx.mma_fill");
static const Op ptx_mma_legacy_op = Op::Get("tirx.ptx_legacy.mma");
@@ -1098,66 +1091,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
static const Op mma_fill_legacy_op = Op::Get("tirx.mma_fill_legacy");
static const Op cuda_func_call_op = Op::Get("tirx.cuda.func_call");
- if (op->op.same_as(gpu_fill_fragment_op)) {
- codegen_tags_.insert("mma");
- TVM_FFI_ICHECK_EQ(op->args.size(), 6U);
- os << "nvcuda::wmma::fill_fragment(";
- this->PrintExpr(op->args[0], os);
- os << "[";
- this->PrintExpr(op->args[4], os);
- os << "], ";
- this->PrintExpr(op->args[5], os);
- os << ")";
- } else if (op->op.same_as(gpu_load_matrix_sync_op)) {
- codegen_tags_.insert("mma");
- TVM_FFI_ICHECK_EQ(op->args.size(), 8U);
- os << "nvcuda::wmma::load_matrix_sync(";
- this->PrintExpr(op->args[0], os);
- os << "[";
- this->PrintExpr(op->args[4], os);
- os << "], ";
- this->PrintExpr(op->args[5], os);
- os << ", ";
- this->PrintExpr(op->args[6], os);
- os << ")";
- } else if (op->op.same_as(gpu_store_matrix_sync_op)) {
- codegen_tags_.insert("mma");
- TVM_FFI_ICHECK_EQ(op->args.size(), 8U);
- os << "nvcuda::wmma::store_matrix_sync(";
- this->PrintExpr(op->args[5], os);
- os << ", ";
- this->PrintExpr(op->args[0], os);
- os << "[";
- this->PrintExpr(op->args[4], os);
- os << "], ";
- this->PrintExpr(op->args[6], os);
- if (const StringImmNode* str = op->args[7].as<StringImmNode>()) {
- os << ", nvcuda::wmma::mem_" << str->value;
- } else {
- TVM_FFI_THROW(InternalError) << "Invalid parameters";
- }
- os << ")";
- } else if (op->op.same_as(gpu_mma_sync_op)) {
- codegen_tags_.insert("mma");
- TVM_FFI_ICHECK_EQ(op->args.size(), 8U);
- os << "nvcuda::wmma::mma_sync(";
- for (int i = 0; i < 4; ++i) {
- this->PrintExpr(op->args[i * 2], os);
- os << "[";
- this->PrintExpr(op->args[i * 2 + 1], os);
- os << "]" << ((i < 3) ? ", " : ")");
- }
- } else if (op->op.same_as(bmma_sync_op)) {
- codegen_tags_.insert("mma");
- TVM_FFI_ICHECK_EQ(op->args.size(), 8U);
- os << "nvcuda::wmma::bmma_sync(";
- for (int i = 0; i < 4; ++i) {
- this->PrintExpr(op->args[i * 2], os);
- os << "[";
- this->PrintExpr(op->args[i * 2 + 1], os);
- os << "]" << ((i < 3) ? ", " : ")");
- }
- } else if (op->op.same_as(mma_store_op)) {
+ if (op->op.same_as(mma_store_op)) {
int m = op->args[0].as_or_throw<IntImm>()->value.as<int>().value();
int n = op->args[1].as_or_throw<IntImm>()->value.as<int>().value();
std::string dst = this->PrintExpr(op->args[2]);
@@ -1460,51 +1394,21 @@ void CodeGenCUDA::DispatchAllocTensor(const BindNode* op, const CallNode* buffer
TVM_FFI_ICHECK(buffer.defined());
std::string vid = AllocVarID(buffer.get(), buffer.name() + "_ptr");
- std::string fragment_shape;
- std::string fragment_layout;
- if (std::string(scope).find("wmma.") == 0) {
- auto shape = annotations->dict.Get(s_tir::attr::fragment_shape);
- TVM_FFI_ICHECK(shape.has_value()) << "Cannot find shape of the wmma fragment " << buffer.name();
- fragment_shape = shape.value().as_or_throw<ffi::String>();
- if (auto layout = annotations->dict.Get(s_tir::attr::fragment_layout)) {
- fragment_layout = layout.value().as_or_throw<ffi::String>();
- }
- }
this->PrintIndent();
- if (std::string(scope).find("wmma.") == 0) {
- if (scope == "wmma.matrix_a" || scope == "wmma.matrix_b") {
- bool supported_wmma_input_dtype =
- PrimType(dtype) == PrimType::Float(16) || PrimType(dtype) == PrimType::Int(8) ||
- PrimType(dtype) == PrimType::UInt(8) || PrimType(dtype) == PrimType::Int(4) ||
- PrimType(dtype) == PrimType::UInt(4) || PrimType(dtype) == PrimType::Int(1) ||
- PrimType(dtype) == PrimType::BFloat(16);
- TVM_FFI_ICHECK(supported_wmma_input_dtype)
- << "Matrix_a and matrix_b only support half or char or unsigned char "
- << "or uint4 or int4 or int1 type for now";
- } else {
- bool supported_wmma_accumulator_dtype = PrimType(dtype) == PrimType::Float(16) ||
- PrimType(dtype) == PrimType::Float(32) ||
- PrimType(dtype) == PrimType::Int(32);
- TVM_FFI_ICHECK(supported_wmma_accumulator_dtype)
- << "Accumulator only support half, float and int type for now";
- }
- PrintWmmaScope(scope, PrimType(dtype), fragment_shape, fragment_layout, stream);
- } else {
- PrintStorageScope(scope, stream);
- int align = buffer->data_alignment;
- auto it = annotations->dict.find(tirx::attr::buffer_data_alignment);
- if (it != annotations->dict.end()) {
- if (const auto* n = (*it).second.as<IntImmNode>()) {
- align = n->value.as<int>().value();
- }
- }
- if (align > 0 && scope == "shared.dyn") {
- stream << "__align__(" << align << ") ";
- } else if (align > 0) {
- stream << "alignas(" << align << ") ";
+ PrintStorageScope(scope, stream);
+ int align = buffer->data_alignment;
+ auto it = annotations->dict.find(tirx::attr::buffer_data_alignment);
+ if (it != annotations->dict.end()) {
+ if (const auto* n = (*it).second.as<IntImmNode>()) {
+ align = n->value.as<int>().value();
}
- PrintType(PrimType(dtype), stream);
}
+ if (align > 0 && scope == "shared.dyn") {
+ stream << "__align__(" << align << ") ";
+ } else if (align > 0) {
+ stream << "alignas(" << align << ") ";
+ }
+ PrintType(PrimType(dtype), stream);
if (scope == "shared.dyn") {
stream << ' ' << vid << "[];\n";
@@ -1518,9 +1422,6 @@ void CodeGenCUDA::DispatchAllocTensor(const BindNode* op, const CallNode* buffer
}
TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation for now";
- if (std::string(scope).find("wmma.") == 0) {
- constant_size = GetWmmaFragmentSize(scope, fragment_shape, constant_size);
- }
bool is_packed_integer_dtype = PrimType(dtype) == PrimType::Int(4) ||
PrimType(dtype) == PrimType::UInt(4) ||
PrimType(dtype) == PrimType::Int(1);
@@ -1847,65 +1748,6 @@ void CodeGenCUDA::Dispatch_(const FloatImmNode* op, std::ostream& os) { // NOLI
PrintConst(op, os, this);
}
-void CodeGenCUDA::PrintWmmaScope(const std::string& scope, const PrimType& t,
- const std::string& shape_str, const std::string& layout_str,
- std::ostream& os) {
- std::stringstream type;
- PrintType(t, type);
- if ((t.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) && t.bits() < 8 &&
- t.lanes() == 1) {
- type.str(std::string());
- if (t.MatchesCode(DLDataTypeCode::kDLInt)) {
- if (t.bits() == 4) {
- type << "nvcuda::wmma::experimental::precision::s4";
- } else if (t.bits() == 1) {
- type << "nvcuda::wmma::experimental::precision::b1";
- } else {
- TVM_FFI_THROW(InternalError) << "Unhandled interger type for wmma fragment!";
- }
- } else if (t.MatchesCode(DLDataTypeCode::kDLUInt)) {
- if (t.bits() == 4) {
- type << "nvcuda::wmma::experimental::precision::u4";
- } else {
- TVM_FFI_THROW(InternalError) << "Unhandled interger type for wmma fragment!";
- }
- }
- }
- if (scope == "wmma.matrix_a") {
- codegen_tags_.insert("mma");
- TVM_FFI_ICHECK_NE(layout_str, "") << "Layout must be defined for matrix_a";
- os << "nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, " << shape_str << ", " << type.str()
- << ", nvcuda::wmma::" << layout_str << ">";
- } else if (scope == "wmma.matrix_b") {
- codegen_tags_.insert("mma");
- TVM_FFI_ICHECK_NE(layout_str, "") << "Layout must be defined for matrix_b";
- os << "nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, " << shape_str << ", " << type.str()
- << ", nvcuda::wmma::" << layout_str << ">";
- } else if (scope == "wmma.accumulator") {
- codegen_tags_.insert("mma");
- os << "nvcuda::wmma::fragment<nvcuda::wmma::accumulator, " << shape_str << ", " << type.str()
- << ">";
- }
-}
-
-int stoi(const std::string& str) {
- try {
- return std::stoi(str);
- } catch (std::invalid_argument& e) {
- TVM_FFI_THROW(InternalError) << "Cannot convert \"" << str << "\" to int";
- throw;
- }
-}
-
-int32_t CodeGenCUDA::GetWmmaFragmentSize(const std::string& scope, const std::string& shape_str,
- int32_t size) {
- std::pair<int32_t, int32_t> dim = GetWmmaFragmentDimSize(shape_str, scope);
- if (dim.first * dim.second != 0)
- return size / dim.first / dim.second;
- else
- return 0;
-}
-
void CodeGenCUDA::HandleVolatileLoads(const std::string& value, const TensorLoadNode* op,
std::ostream& os) {
// Cast away volatile qualifier for fp16 types. That is, only loads and
diff --git a/src/backend/cuda/codegen/codegen_cuda.h b/src/backend/cuda/codegen/codegen_cuda.h
index fdcdcf7506..aee905b397 100644
--- a/src/backend/cuda/codegen/codegen_cuda.h
+++ b/src/backend/cuda/codegen/codegen_cuda.h
@@ -137,9 +137,6 @@ class CodeGenCUDA final : public CodeGenC {
const std::string cuda_barrier_arrival_token_name_ = "cubar_tok";
friend void PrintConst(const FloatImmNode* op, std::ostream& os, CodeGenCUDA* p);
- void PrintWmmaScope(const std::string& scope, const PrimType& t, const std::string& shape,
- const std::string& layout, std::ostream& os);
- int32_t GetWmmaFragmentSize(const std::string& scope, const std::string& shape, int32_t size);
};
} // namespace codegen
diff --git a/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc b/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc
index c918b0cf42..2d2fe67bc0 100644
--- a/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc
@@ -174,53 +174,6 @@ ffi::Array<ScheduleRule> ScheduleRule::DefaultCUDA() {
}
ffi::Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
- ffi::Array<ffi::Map<ffi::String, ffi::String>> wmma_intrin_groups = {
- // Tensor Cores f32 += f16 * f16
- {
- {"init", "wmma_fill_16x16x16_f32"},
- {"load_a", "wmma_load_16x16x16_f16_a_shared_dyn"},
- {"load_b", "wmma_load_16x16x16_f16_b_shared_dyn"},
- {"compute", "wmma_sync_16x16x16_f16f16f32"},
- {"store", "wmma_store_16x16x16_f32_shared_dyn"},
- },
- {
- {"init", "wmma_fill_16x16x16_f32"},
- {"load_a", "wmma_load_16x16x16_f16_a_shared_dyn"},
- {"load_b", "wmma_load_16x16x16_f16_b_trans_shared_dyn"},
- {"compute", "wmma_sync_16x16x16_f16f16f32_trans"},
- {"store", "wmma_store_16x16x16_f32_shared_dyn"},
- },
- // Tensor Cores f16 += f16 * f16
- {
- {"init", "wmma_fill_16x16x16_f16"},
- {"load_a", "wmma_load_16x16x16_f16_a_shared_dyn"},
- {"load_b", "wmma_load_16x16x16_f16_b_shared_dyn"},
- {"compute", "wmma_sync_16x16x16_f16f16f16"},
- {"store", "wmma_store_16x16x16_f16_shared_dyn"},
- },
- {
- {"init", "wmma_fill_16x16x16_f16"},
- {"load_a", "wmma_load_16x16x16_f16_a_shared_dyn"},
- {"load_b", "wmma_load_16x16x16_f16_b_trans_shared_dyn"},
- {"compute", "wmma_sync_16x16x16_f16f16f16_trans"},
- {"store", "wmma_store_16x16x16_f16_shared_dyn"},
- },
- // Tensor Cores s32 += s8 * s8
- {
- {"init", "wmma_fill_16x16x16_s32"},
- {"load_a", "wmma_load_16x16x16_s8_a_shared_dyn"},
- {"load_b", "wmma_load_16x16x16_s8_b_shared_dyn"},
- {"compute", "wmma_sync_16x16x16_s8s8s32"},
- {"store", "wmma_store_16x16x16_s32_shared_dyn"},
- },
- {
- {"init", "wmma_fill_16x16x16_s32"},
- {"load_a", "wmma_load_16x16x16_s8_a_shared_dyn"},
- {"load_b", "wmma_load_16x16x16_s8_b_trans_shared_dyn"},
- {"compute", "wmma_sync_16x16x16_s8s8s32_trans"},
- {"store", "wmma_store_16x16x16_s32_shared_dyn"},
- },
- };
ffi::Array<ffi::Map<ffi::String, ffi::String>> mma_intrin_groups = {
// Tensor Core MMA
{
@@ -240,21 +193,6 @@ ffi::Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
};
ffi::Array<ScheduleRule> results{
ScheduleRule::ApplyCustomRule(),
- ScheduleRule::MultiLevelTilingTensorCore(
- /*intrin_groups=*/wmma_intrin_groups,
- /*structure=*/"SSSRRSRS",
- /*tile_binds=*/ffi::Array<ffi::String>{"blockIdx.y", "blockIdx.x", "threadIdx.y"},
- /*max_innermost_factor=*/static_cast<int64_t>(4),
- /*vector_load_lens=*/ffi::Array<int64_t>{1, 2, 3, 4, 8, 16},
- /*reuse_read=*/
- ffi::Map<ffi::String, ffi::Any>{{"req", ffi::String("must")},
- {"levels", ffi::Array<int64_t>{4}}, //
- {"scope", ffi::String("shared.dyn")}},
- /*reuse_write=*/
- ffi::Map<ffi::String, ffi::Any>{{"req", ffi::String("must")},
- {"levels", ffi::Array<int64_t>{2}}, //
- {"scope", ffi::String("shared.dyn")}},
- /*use_software_pipeline=*/false), //
ScheduleRule::MultiLevelTilingTensorCore(
/*intrin_groups=*/mma_intrin_groups,
/*structure=*/"SSSRRSRS",
diff --git a/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
index 9cc0a72369..7668210885 100644
--- a/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
+++ b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=missing-docstring, unused-variable, invalid-name
-# ruff: noqa: E501, F841
+# ruff: noqa: E501
import tvm
import tvm.testing
@@ -25,166 +25,6 @@ from tvm.script import tirx as T
from tvm.target import Target
-def test_matmul_tensorize():
- # fmt: off
- @Ts.function(private=True)
- def before(X: T.Tensor((256, 256), "float16"), W: T.Tensor((256, 256), "float16"), compute: T.Tensor((256, 256), "float16")):
- T.func_attr({"tirx.noalias": True})
- # with Ts.sblock("root"):
- for i, j, k in T.grid(256, 256, 256):
- with Ts.sblock("compute"):
- v_i, v_j, v_k = Ts.axis.remap("SSR", [i, j, k])
- Ts.reads(X[v_i, v_k], W[v_j, v_k])
- Ts.writes(compute[v_i, v_j])
- with Ts.init():
- compute[v_i, v_j] = T.float16(0)
- compute[v_i, v_j] = compute[v_i, v_j] + X[v_i, v_k] * W[v_j, v_k]
-
- A_3_s0 = T.dynamic("A_3_s0", "int32")
- A_3_s1 = T.dynamic("A_3_s1", "int32")
- C_4_s0 = T.dynamic("C_4_s0", "int32")
- C_4_s1 = T.dynamic("C_4_s1", "int32")
- C_s0 = T.dynamic("C_s0", "int32")
- C_s1 = T.dynamic("C_s1", "int32")
- A_s0 = T.dynamic("A_s0", "int32")
- A_s1 = T.dynamic("A_s1", "int32")
- C_1_s0 = T.dynamic("C_1_s0", "int32")
- C_1_s1 = T.dynamic("C_1_s1", "int32")
- A_1_s0 = T.dynamic("A_1_s0", "int32")
- A_1_s1 = T.dynamic("A_1_s1", "int32")
- C_2_s0 = T.dynamic("C_2_s0", "int32")
- C_2_s1 = T.dynamic("C_2_s1", "int32")
- A_2_s0 = T.dynamic("A_2_s0", "int32")
- A_2_s1 = T.dynamic("A_2_s1", "int32")
- B_s0 = T.dynamic("B_s0", "int32")
- B_s1 = T.dynamic("B_s1", "int32")
- C_3_s0 = T.dynamic("C_3_s0", "int32")
- C_3_s1 = T.dynamic("C_3_s1", "int32")
-
- @Ts.function(private=True)
- def expected(X: T.Tensor((256, 256), "float16"), W: T.Tensor((256, 256), "float16"), compute: T.Tensor((256, 256), "float16")):
- T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
- # with Ts.sblock("root"):
- X_reindex_shared_dyn = Ts.sblock_alloc_buffer((1, 256, 256), "float16", scope="shared.dyn")
- W_reindex_shared_dyn = Ts.sblock_alloc_buffer((1, 256, 256), "float16", scope="shared.dyn")
- X_reindex_shared_dyn_wmma_matrix_a = Ts.sblock_alloc_buffer((1, 256, 256), "float16", scope="wmma.matrix_a")
- W_reindex_shared_dyn_wmma_matrix_b = Ts.sblock_alloc_buffer((1, 256, 256), "float16", scope="wmma.matrix_b")
- compute_reindex_shared_dyn = Ts.sblock_alloc_buffer((1, 256, 256), "float16", scope="shared.dyn")
- compute_reindex_shared_dyn_wmma_accumulator = Ts.sblock_alloc_buffer((1, 256, 256), "float16", scope="wmma.accumulator")
- for ax0 in T.thread_binding(1, thread="blockIdx.z"):
- for ax1_0_0_ax2_0_0_fused in T.thread_binding(2, thread="blockIdx.x"):
- for ax1_0_1_ax2_0_1_fused in T.thread_binding(2, thread="blockIdx.y"):
- for ax2_0_2_ax1_0_2_fused in T.thread_binding(16, thread="threadIdx.y"):
- for ax1_0_3_init, ax2_0_3_init in T.grid(2, 2):
- with Ts.sblock("compute_o_init"):
- v0_o = Ts.axis.spatial(1, ax0)
- v1_o = Ts.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3_init)
- v2_o = Ts.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3_init)
- Ts.reads()
- Ts.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- with Ts.sblock("compute_init_o"):
- v1_i_init_o = Ts.axis.spatial(1, 0)
- v2_i_init_o = Ts.axis.spatial(1, 0)
- Ts.reads()
- Ts.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- C = Ts.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_s0, C_s1), scope="wmma.accumulator", offset_factor=16)
- T.gpu_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0))
- for ax3_0_0 in range(4, annotations={"software_pipeline_order": [0, 3, 1, 4, 5, 2, 6], "software_pipeline_stage": [0, 0, 0, 0, 0, 1, 1]}):
- for ax0_ax1_fused_0 in range(4):
- for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
- for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_3 in T.vectorized(4):
- with Ts.sblock("X_reindex_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial(256, ax1_0_0_ax2_0_0_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 64)
- v2 = Ts.axis.spatial(256, ax3_0_0 * 64 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 64)
- Ts.reads(X[v1, v2])
- Ts.writes(X_reindex_shared_dyn[v0, v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 16, 8]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
- X_reindex_shared_dyn[v0, v1, v2] = X[v1, v2]
- for ax0_ax1_fused_0 in range(4):
- for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
- for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_3 in T.vectorized(4):
- with Ts.sblock("W_reindex_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 64)
- v2 = Ts.axis.spatial(256, ax3_0_0 * 64 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 64)
- Ts.reads(W[v1, v2])
- Ts.writes(W_reindex_shared_dyn[v0, v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 16, 8]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
- W_reindex_shared_dyn[v0, v1, v2] = W[v1, v2]
- for ax3_0_1 in range(4, annotations={"software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, 1]}):
- for ax0_0 in T.unroll(2):
- for ax1_0 in T.unroll(1):
- with Ts.sblock("X_reindex_shared.dyn_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(16, ax3_0_0 * 4 + ax3_0_1 + ax1_0)
- Ts.reads(X_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(X_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(A_s0, A_s1), scope="shared.dyn", offset_factor=16)
- C = Ts.match_buffer(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_1_s0, C_1_s1), scope="wmma.matrix_a", offset_factor=16)
- T.gpu_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.access_ptr("float16", A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
- for ax0_0 in T.unroll(2):
- for ax1_0 in T.unroll(1):
- with Ts.sblock("W_reindex_shared.dyn_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(16, ax3_0_0 * 4 + ax3_0_1 + ax1_0)
- Ts.reads(W_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(W_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(A_1_s0, A_1_s1), scope="shared.dyn", offset_factor=16)
- C = Ts.match_buffer(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_2_s0, C_2_s1), scope="wmma.matrix_b", offset_factor=16)
- T.gpu_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.access_ptr("float16", A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
- for ax1_0_3, ax2_0_3 in T.grid(2, 2):
- with Ts.sblock("compute_o_update"):
- v0_o = Ts.axis.spatial(1, ax0)
- v1_o = Ts.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3)
- v2_o = Ts.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3)
- v3_o = Ts.axis.reduce(16, ax3_0_0 * 4 + ax3_0_1)
- Ts.reads(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- with Ts.sblock("compute_o"):
- v1_i_o = Ts.axis.spatial(1, 0)
- v2_i_o = Ts.axis.spatial(1, 0)
- v3_i_o = Ts.axis.reduce(1, 0)
- Ts.reads(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "float16", strides=(A_2_s0, A_2_s1), scope="wmma.matrix_a", offset_factor=16)
- B = Ts.match_buffer(W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "float16", strides=(B_s0, B_s1), scope="wmma.matrix_b", offset_factor=16)
- C = Ts.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_3_s0, C_3_s1), scope="wmma.accumulator", offset_factor=16)
- T.gpu_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
- for ax0_0, ax1_0 in T.grid(2, 2):
- with Ts.sblock("compute_reindex_shared.dyn_wmma.accumulator_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax1_0)
- Ts.reads(compute_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(compute_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(A_3_s0, A_3_s1), scope="wmma.accumulator", offset_factor=16)
- C = Ts.match_buffer(compute_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_4_s0, C_4_s1), scope="shared.dyn", offset_factor=16)
- T.gpu_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.access_ptr("float16", C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
- for ax0_ax1_fused_0 in range(8):
- for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_2 in T.vectorized(4):
- with Ts.sblock("compute_reindex_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial(256, ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 32)
- v2 = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 128 + ax2_0_2_ax1_0_2_fused // 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 32)
- Ts.reads(compute_reindex_shared_dyn[v0, v1, v2])
- Ts.writes(compute[v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 16, 4]]})
- compute[v1, v2] = compute_reindex_shared_dyn[v0, v1, v2]
-
- # fmt: on
-
- mod = tvm.IRModule({"main": before})
- with Target("nvidia/geforce-rtx-2080-ti"):
- mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
- tvm.ir.assert_structural_equal(mod["main"], expected)
-
-
def test_matmul_tensorize_too_small():
# fmt: off
m = T.dynamic("m", "int32")
@@ -281,516 +121,6 @@ def test_matmul_tensorize_too_small():
tvm.ir.assert_structural_equal(mod["main"], expected)
-def test_matmul_tensorize_epilogue():
- # fmt: off
- n = T.dynamic("n", "int32")
-
- @Ts.function(private=True)
- def before(lv686: T.Tensor((T.int32(4096), T.int32(256)), "uint32"), lv687: T.Tensor((T.int32(4096), T.int32(64)), "float16"), lv42: T.Tensor((T.int32(1), n, T.int32(2048)), 'float16'), lv3: T.Tensor((T.int32(1), n, T.int32(4096)), 'float16'), p_output0_intermediate: T.Tensor((T.int32(1), n, T.int32(4096)), 'float16')):
- T.func_attr({"tirx.noalias": True})
-
- # with Ts.sblock("root"):
- p_output0_intermediate_1 = Ts.sblock_alloc_buffer((T.int32(4096), T.int32(2048)), "float16")
- var_NT_matmul_intermediate = Ts.sblock_alloc_buffer((T.int32(1), n, T.int32(4096)), "float16")
- var_T_divide_intermediate = Ts.sblock_alloc_buffer((T.int32(1), n, T.int32(4096)), "float16")
- for i, j in T.grid(T.int32(4096), T.int32(2048)):
- with Ts.sblock("decode"):
- v_i, v_j = Ts.axis.remap("SS", [i, j])
- Ts.reads(lv686[v_i, v_j // T.int32(8)], lv687[v_i, v_j // T.int32(32)])
- Ts.writes(p_output0_intermediate_1[v_i, v_j])
- p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv686[v_i, v_j // T.int32(8)], T.Cast("uint32", v_j % T.int32(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv687[v_i, v_j // T.int32(32)]
- for i0, i1, i2, k in T.grid(T.int32(1), n, T.int32(4096), T.int32(2048)):
- with Ts.sblock("NT_matmul"):
- v_i0, v_i1, v_i2, v_k = Ts.axis.remap("SSSR", [i0, i1, i2, k])
- Ts.reads(lv42[v_i0, v_i1, v_k], p_output0_intermediate_1[v_i2, v_k])
- Ts.writes(var_NT_matmul_intermediate[v_i0, v_i1, v_i2])
- with Ts.init():
- var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
- var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2] + lv42[v_i0, v_i1, v_k] * p_output0_intermediate_1[v_i2, v_k]
- for ax0, ax1, ax2 in T.grid(T.int32(1), n, T.int32(4096)):
- with Ts.sblock("T_divide"):
- v_ax0, v_ax1, v_ax2 = Ts.axis.remap("SSS", [ax0, ax1, ax2])
- Ts.reads(lv3[v_ax0, v_ax1, v_ax2])
- Ts.writes(var_T_divide_intermediate[v_ax0, v_ax1, v_ax2])
- var_T_divide_intermediate[v_ax0, v_ax1, v_ax2] = lv3[v_ax0, v_ax1, v_ax2] * T.float16(0.5)
- for ax0, ax1, ax2 in T.grid(T.int32(1), n, T.int32(4096)):
- with Ts.sblock("T_add"):
- v_ax0, v_ax1, v_ax2 = Ts.axis.remap("SSS", [ax0, ax1, ax2])
- Ts.reads(var_T_divide_intermediate[v_ax0, v_ax1, v_ax2], var_NT_matmul_intermediate[v_ax0, v_ax1, v_ax2])
- Ts.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
- p_output0_intermediate[v_ax0, v_ax1, v_ax2] = var_T_divide_intermediate[v_ax0, v_ax1, v_ax2] + var_NT_matmul_intermediate[v_ax0, v_ax1, v_ax2]
-
- n = T.dynamic("n", "int32")
- A_3_s0 = T.dynamic("A_3_s0", "int32")
- A_3_s1 = T.dynamic("A_3_s1", "int32")
- C_4_s0 = T.dynamic("C_4_s0", "int32")
- C_4_s1 = T.dynamic("C_4_s1", "int32")
- C_s0 = T.dynamic("C_s0", "int32")
- C_s1 = T.dynamic("C_s1", "int32")
- A_s0 = T.dynamic("A_s0", "int32")
- A_s1 = T.dynamic("A_s1", "int32")
- C_1_s0 = T.dynamic("C_1_s0", "int32")
- C_1_s1 = T.dynamic("C_1_s1", "int32")
- A_1_s0 = T.dynamic("A_1_s0", "int32")
- A_1_s1 = T.dynamic("A_1_s1", "int32")
- C_2_s0 = T.dynamic("C_2_s0", "int32")
- C_2_s1 = T.dynamic("C_2_s1", "int32")
- A_2_s0 = T.dynamic("A_2_s0", "int32")
- A_2_s1 = T.dynamic("A_2_s1", "int32")
- B_s0 = T.dynamic("B_s0", "int32")
- B_s1 = T.dynamic("B_s1", "int32")
- C_3_s0 = T.dynamic("C_3_s0", "int32")
- C_3_s1 = T.dynamic("C_3_s1", "int32")
-
- @Ts.function(private=True)
- def expected(lv686: T.Tensor((4096, 256), "uint32"), lv687: T.Tensor((4096, 64), "float16"), lv42: T.Tensor((1, n, 2048), 'float16'), lv3: T.Tensor((1, n, 4096), 'float16'), p_output0_intermediate: T.Tensor((1, n, 4096), 'float16')):
- T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
-
- # with Ts.sblock("root"):
- lv42_reindex_pad_shared_dyn = Ts.sblock_alloc_buffer((1, (n + 127) // 128 * 128, 2048), "float16", scope="shared.dyn")
- p_output0_intermediate_1_reindex_shared_dyn = Ts.sblock_alloc_buffer((1, 4096, 2048), "float16", scope="shared.dyn")
- lv42_reindex_pad_shared_dyn_wmma_matrix_a = Ts.sblock_alloc_buffer((1, (n + 127) // 128 * 128, 2048), "float16", scope="wmma.matrix_a")
- p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b = Ts.sblock_alloc_buffer((1, 4096, 2048), "float16", scope="wmma.matrix_b")
- var_NT_matmul_intermediate_reindex_pad_shared_dyn = Ts.sblock_alloc_buffer((1, (n + 127) // 128 * 128, 4096), "float16", scope="shared.dyn")
- var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator = Ts.sblock_alloc_buffer((1, (n + 127) // 128 * 128, 4096), "float16", scope="wmma.accumulator")
- for ax0 in T.thread_binding(1, thread="blockIdx.z"):
- for ax1_0_0_ax2_0_0_fused in T.thread_binding((n + 127) // 128, thread="blockIdx.x"):
- for ax1_0_1_ax2_0_1_fused in T.thread_binding(32, thread="blockIdx.y"):
- for ax2_0_2_ax1_0_2_fused in T.thread_binding(16, thread="threadIdx.y"):
- for ax1_0_3_init, ax2_0_3_init in T.grid(2, 2):
- with Ts.sblock("NT_matmul_o_init"):
- v0_o = Ts.axis.spatial(1, ax0)
- v1_o = Ts.axis.spatial((n + 127) // 128 * 8, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3_init)
- v2_o = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3_init)
- Ts.reads()
- Ts.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- with Ts.sblock("NT_matmul_init_o"):
- v1_i_init_o = Ts.axis.spatial(1, 0)
- v2_i_init_o = Ts.axis.spatial(1, 0)
- Ts.reads()
- Ts.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- C = Ts.match_buffer(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_s0, C_s1), scope="wmma.accumulator", offset_factor=16)
- T.gpu_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0))
- for ax3_0_0 in range(32, annotations={"software_pipeline_order": [0, 3, 1, 4, 5, 2, 6], "software_pipeline_stage": [0, 0, 0, 0, 0, 1, 1]}):
- for ax0_ax1_fused_0 in range(4):
- for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
- for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_3 in T.vectorized(4):
- with Ts.sblock("lv42_reindex_pad_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial((n + 127) // 128 * 128, ax1_0_0_ax2_0_0_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 64)
- v2 = Ts.axis.spatial(2048, ax3_0_0 * 64 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 64)
- Ts.reads(lv42[v0, v1, v2])
- Ts.writes(lv42_reindex_pad_shared_dyn[v0, v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 16, 8]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
- lv42_reindex_pad_shared_dyn[v0, v1, v2] = T.if_then_else(v1 < n, lv42[v0, v1, v2], T.float16(0))
- for ax0_ax1_fused_0 in range(4):
- for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
- for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_3 in T.vectorized(4):
- with Ts.sblock("p_output0_intermediate_1_reindex_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial(4096, ax1_0_1_ax2_0_1_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 64)
- v2 = Ts.axis.spatial(2048, ax3_0_0 * 64 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 64)
- Ts.reads(lv686[v1, v2 // 8], lv687[v1, v2 // 32])
- Ts.writes(p_output0_intermediate_1_reindex_shared_dyn[v0, v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 16, 8]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
- p_output0_intermediate_1_reindex_shared_dyn[v0, v1, v2] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv686[v1, v2 // 8], T.Cast("uint32", v2 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv687[v1, v2 // 32]
- for ax3_0_1 in range(4, annotations={"software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, 1]}):
- for ax0_0 in T.unroll(2):
- for ax1_0 in T.unroll(1):
- with Ts.sblock("lv42_reindex_pad_shared.dyn_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(8 * ((n + 127) // 128), ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(128, ax3_0_0 * 4 + ax3_0_1 + ax1_0)
- Ts.reads(lv42_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(lv42_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(lv42_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(A_s0, A_s1), scope="shared.dyn", offset_factor=16)
- C = Ts.match_buffer(lv42_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_1_s0, C_1_s1), scope="wmma.matrix_a", offset_factor=16)
- T.gpu_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.access_ptr("float16", A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
- for ax0_0 in T.unroll(2):
- for ax1_0 in T.unroll(1):
- with Ts.sblock("p_output0_intermediate_1_reindex_shared.dyn_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(128, ax3_0_0 * 4 + ax3_0_1 + ax1_0)
- Ts.reads(p_output0_intermediate_1_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(p_output0_intermediate_1_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(A_1_s0, A_1_s1), scope="shared.dyn", offset_factor=16)
- C = Ts.match_buffer(p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_2_s0, C_2_s1), scope="wmma.matrix_b", offset_factor=16)
- T.gpu_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.access_ptr("float16", A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
- for ax1_0_3, ax2_0_3 in T.grid(2, 2):
- with Ts.sblock("NT_matmul_o_update"):
- v0_o = Ts.axis.spatial(1, ax0)
- v1_o = Ts.axis.spatial((n + 127) // 128 * 8, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3)
- v2_o = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3)
- v3_o = Ts.axis.reduce(128, ax3_0_0 * 4 + ax3_0_1)
- Ts.reads(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], lv42_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- with Ts.sblock("NT_matmul_o"):
- v1_i_o = Ts.axis.spatial(1, 0)
- v2_i_o = Ts.axis.spatial(1, 0)
- v3_i_o = Ts.axis.reduce(1, 0)
- Ts.reads(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], lv42_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(lv42_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "float16", strides=(A_2_s0, A_2_s1), scope="wmma.matrix_a", offset_factor=16)
- B = Ts.match_buffer(p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "float16", strides=(B_s0, B_s1), scope="wmma.matrix_b", offset_factor=16)
- C = Ts.match_buffer(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_3_s0, C_3_s1), scope="wmma.accumulator", offset_factor=16)
- T.gpu_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
- for ax0_0, ax1_0 in T.grid(2, 2):
- with Ts.sblock("var_NT_matmul_intermediate_reindex_pad_shared.dyn_wmma.accumulator_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(8 * ((n + 127) // 128), ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax1_0)
- Ts.reads(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(A_3_s0, A_3_s1), scope="wmma.accumulator", offset_factor=16)
- C = Ts.match_buffer(var_NT_matmul_intermediate_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=(C_4_s0, C_4_s1), scope="shared.dyn", offset_factor=16)
- T.gpu_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.access_ptr("float16", C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
- for ax0_ax1_fused_0 in range(8):
- for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_2 in T.vectorized(4):
- with Ts.sblock("var_NT_matmul_intermediate_reindex_pad_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial((n + 127) // 128 * 128, ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 32)
- v2 = Ts.axis.spatial(4096, ax1_0_1_ax2_0_1_fused * 128 + ax2_0_2_ax1_0_2_fused // 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 32)
- Ts.where(ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + ((ax0_ax1_fused_0 * 32 + ax0_ax1_fused_1) * 4 + ax0_ax1_fused_2) // 32 < n)
- Ts.reads(lv3[0, v1, v2], var_NT_matmul_intermediate_reindex_pad_shared_dyn[v0, v1, v2])
- Ts.writes(p_output0_intermediate[0, v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 16, 4]]})
- p_output0_intermediate[0, v1, v2] = lv3[0, v1, v2] * T.float16(0.5) + var_NT_matmul_intermediate_reindex_pad_shared_dyn[v0, v1, v2]
- # fmt: on
-
- mod = tvm.IRModule({"main": before})
- with Target("nvidia/geforce-rtx-2080-ti"):
- mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
- tvm.ir.assert_structural_equal(mod["main"], expected)
-
-
-def test_matmul_int8_tensorize():
- # fmt: off
- @Ts.function(private=True)
- def before(X: T.Tensor((256, 256), "int8"), W: T.Tensor((256, 256), "int8"), compute: T.Tensor((256, 256), "int32")):
- T.func_attr({"tirx.noalias": True})
- # with Ts.sblock("root"):
- for i, j, r in T.grid(256, 256, 256):
- with Ts.sblock("compute"):
- v_i, v_j, v_k = Ts.axis.remap("SSR", [i, j, r])
- Ts.reads(X[v_i, v_k], W[v_j, v_k])
- Ts.writes(compute[v_i, v_j])
- with Ts.init():
- compute[v_i, v_j] = 0
- compute[v_i, v_j] = compute[v_i, v_j] + T.Cast("int32", X[v_i, v_k]) * T.Cast("int32", W[v_j, v_k])
-
- A_3_s0 = T.dynamic("A_3_s0", "int32")
- A_3_s1 = T.dynamic("A_3_s1", "int32")
- C_4_s0 = T.dynamic("C_4_s0", "int32")
- C_4_s1 = T.dynamic("C_4_s1", "int32")
- C_s0 = T.dynamic("C_s0", "int32")
- C_s1 = T.dynamic("C_s1", "int32")
- A_s0 = T.dynamic("A_s0", "int32")
- A_s1 = T.dynamic("A_s1", "int32")
- C_1_s0 = T.dynamic("C_1_s0", "int32")
- C_1_s1 = T.dynamic("C_1_s1", "int32")
- A_1_s0 = T.dynamic("A_1_s0", "int32")
- A_1_s1 = T.dynamic("A_1_s1", "int32")
- C_2_s0 = T.dynamic("C_2_s0", "int32")
- C_2_s1 = T.dynamic("C_2_s1", "int32")
- A_2_s0 = T.dynamic("A_2_s0", "int32")
- A_2_s1 = T.dynamic("A_2_s1", "int32")
- B_s0 = T.dynamic("B_s0", "int32")
- B_s1 = T.dynamic("B_s1", "int32")
- C_3_s0 = T.dynamic("C_3_s0", "int32")
- C_3_s1 = T.dynamic("C_3_s1", "int32")
-
- @Ts.function(private=True)
- def expected(X: T.Tensor((256, 256), "int8"), W: T.Tensor((256, 256), "int8"), compute: T.Tensor((256, 256), "int32")):
- T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
- # with Ts.sblock("root"):
- X_reindex_shared_dyn = Ts.sblock_alloc_buffer((1, 256, 256), "int8", scope="shared.dyn")
- W_reindex_shared_dyn = Ts.sblock_alloc_buffer((1, 256, 256), "int8", scope="shared.dyn")
- X_reindex_shared_dyn_wmma_matrix_a = Ts.sblock_alloc_buffer((1, 256, 256), "int8", scope="wmma.matrix_a")
- W_reindex_shared_dyn_wmma_matrix_b = Ts.sblock_alloc_buffer((1, 256, 256), "int8", scope="wmma.matrix_b")
- compute_reindex_shared_dyn = Ts.sblock_alloc_buffer((1, 256, 256), "int32", scope="shared.dyn")
- compute_reindex_shared_dyn_wmma_accumulator = Ts.sblock_alloc_buffer((1, 256, 256), "int32", scope="wmma.accumulator")
- for ax0 in T.thread_binding(1, thread="blockIdx.z"):
- for ax1_0_0_ax2_0_0_fused in T.thread_binding(2, thread="blockIdx.x"):
- for ax1_0_1_ax2_0_1_fused in T.thread_binding(2, thread="blockIdx.y"):
- for ax2_0_2_ax1_0_2_fused in T.thread_binding(16, thread="threadIdx.y"):
- for ax1_0_3_init, ax2_0_3_init in T.grid(2, 2):
- with Ts.sblock("compute_o_init"):
- v0_o = Ts.axis.spatial(1, ax0)
- v1_o = Ts.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3_init)
- v2_o = Ts.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3_init)
- Ts.reads()
- Ts.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- with Ts.sblock("compute_init_o"):
- v1_i_init_o = Ts.axis.spatial(1, 0)
- v2_i_init_o = Ts.axis.spatial(1, 0)
- Ts.reads()
- Ts.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- C = Ts.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=(C_s0, C_s1), scope="wmma.accumulator", offset_factor=16)
- T.gpu_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0))
- for ax3_0_0 in T.serial(16, annotations={"software_pipeline_order": [0, 3, 1, 4, 5, 2, 6], "software_pipeline_stage": [0, 0, 0, 0, 0, 1, 1]}):
- for ax0_ax1_fused_0 in range(1):
- for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
- for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_3 in T.vectorized(4):
- with Ts.sblock("X_reindex_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial(256, ax1_0_0_ax2_0_0_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 16)
- v2 = Ts.axis.spatial(256, ax3_0_0 * 16 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 16)
- Ts.reads(X[v1, v2])
- Ts.writes(X_reindex_shared_dyn[v0, v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 32, 16]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
- X_reindex_shared_dyn[v0, v1, v2] = X[v1, v2]
- for ax0_ax1_fused_0 in range(1):
- for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
- for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_3 in T.vectorized(4):
- with Ts.sblock("W_reindex_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 16)
- v2 = Ts.axis.spatial(256, ax3_0_0 * 16 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 16)
- Ts.reads(W[v1, v2])
- Ts.writes(W_reindex_shared_dyn[v0, v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 32, 16]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
- W_reindex_shared_dyn[v0, v1, v2] = W[v1, v2]
- for ax3_0_1 in T.serial(1, annotations={"software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, 1]}):
- for ax0_0 in T.unroll(2):
- for ax1_0 in T.unroll(1):
- with Ts.sblock("X_reindex_shared.dyn_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(16, ax3_0_0 + ax1_0)
- Ts.reads(X_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(X_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(A_s0, A_s1), scope="shared.dyn", offset_factor=16)
- C = Ts.match_buffer(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(C_1_s0, C_1_s1), scope="wmma.matrix_a", offset_factor=16)
- T.gpu_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.access_ptr("int8", A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
- for ax0_0 in T.unroll(2):
- for ax1_0 in T.unroll(1):
- with Ts.sblock("W_reindex_shared.dyn_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(16, ax3_0_0 + ax1_0)
- Ts.reads(W_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(W_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(A_1_s0, A_1_s1), scope="shared.dyn", offset_factor=16)
- C = Ts.match_buffer(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(C_2_s0, C_2_s1), scope="wmma.matrix_b", offset_factor=16)
- T.gpu_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.access_ptr("int8", A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
- for ax1_0_3, ax2_0_3 in T.grid(2, 2):
- with Ts.sblock("compute_o_update"):
- v0_o = Ts.axis.spatial(1, ax0)
- v1_o = Ts.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3)
- v2_o = Ts.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3)
- v3_o = Ts.axis.reduce(16, ax3_0_0 + ax3_0_1)
- Ts.reads(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- with Ts.sblock("compute_o"):
- v1_i_o = Ts.axis.spatial(1, 0)
- v2_i_o = Ts.axis.spatial(1, 0)
- v3_i_o = Ts.axis.reduce(1, 0)
- Ts.reads(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=(A_2_s0, A_2_s1), scope="wmma.matrix_a", offset_factor=16)
- B = Ts.match_buffer(W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=(B_s0, B_s1), scope="wmma.matrix_b", offset_factor=16)
- C = Ts.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=(C_3_s0, C_3_s1), scope="wmma.accumulator", offset_factor=16)
- T.gpu_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
- for ax0_0, ax1_0 in T.grid(2, 2):
- with Ts.sblock("compute_reindex_shared.dyn_wmma.accumulator_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax1_0)
- Ts.reads(compute_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(compute_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A = Ts.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=(A_3_s0, A_3_s1), scope="wmma.accumulator", offset_factor=16)
- C = Ts.match_buffer(compute_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=(C_4_s0, C_4_s1), scope="shared.dyn", offset_factor=16)
- T.gpu_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.access_ptr("int32", C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
- for ax0_ax1_fused_0 in range(8):
- for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_2 in T.vectorized(4):
- with Ts.sblock("compute_reindex_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial(256, ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 32)
- v2 = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 128 + ax2_0_2_ax1_0_2_fused // 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 32)
- Ts.reads(compute_reindex_shared_dyn[v0, v1, v2])
- Ts.writes(compute[v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 16, 4]]})
- compute[v1, v2] = compute_reindex_shared_dyn[v0, v1, v2]
- # fmt: on
-
- mod = tvm.IRModule({"main": before})
- with Target("nvidia/geforce-rtx-2080-ti"):
- mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
- tvm.ir.assert_structural_equal(mod["main"], expected)
-
-
-def test_matmul_int8_tensorize_3d2d_dyn():
- # fmt: off
- m = T.dynamic("m", "int32")
-
- @Ts.function(private=True)
- def before(A: T.Tensor((1, m, 22016), 'int8'), B: T.Tensor((4096, 22016), "int8"), matmul_1: T.Tensor((1, m, 4096), 'int32')):
- T.func_attr({"op_pattern": 4, "tirx.noalias": True})
-
- # with Ts.sblock("root"):
- for i0, i1, i2, k in T.grid(1, m, 4096, 22016):
- with Ts.sblock("matmul"):
- v_i0, v_i1, v_i2, v_k = Ts.axis.remap("SSSR", [i0, i1, i2, k])
- Ts.reads(A[v_i0, v_i1, v_k], B[v_i2, v_k])
- Ts.writes(matmul_1[v_i0, v_i1, v_i2])
- with Ts.init():
- matmul_1[v_i0, v_i1, v_i2] = 0
- matmul_1[v_i0, v_i1, v_i2] = matmul_1[v_i0, v_i1, v_i2] + T.Cast("int32", A[v_i0, v_i1, v_k]) * T.Cast("int32", B[v_i2, v_k])
-
- m = T.dynamic("m", "int32")
- A_3_s0 = T.dynamic("A_3_s0", "int32")
- A_3_s1 = T.dynamic("A_3_s1", "int32")
- C_4_s0 = T.dynamic("C_4_s0", "int32")
- C_4_s1 = T.dynamic("C_4_s1", "int32")
- C_s0 = T.dynamic("C_s0", "int32")
- C_s1 = T.dynamic("C_s1", "int32")
- A_s0 = T.dynamic("A_s0", "int32")
- A_s1 = T.dynamic("A_s1", "int32")
- C_1_s0 = T.dynamic("C_1_s0", "int32")
- C_1_s1 = T.dynamic("C_1_s1", "int32")
- A_1_s0 = T.dynamic("A_1_s0", "int32")
- A_1_s1 = T.dynamic("A_1_s1", "int32")
- C_2_s0 = T.dynamic("C_2_s0", "int32")
- C_2_s1 = T.dynamic("C_2_s1", "int32")
- A_2_s0 = T.dynamic("A_2_s0", "int32")
- A_2_s1 = T.dynamic("A_2_s1", "int32")
- B_s0 = T.dynamic("B_s0", "int32")
- B_s1 = T.dynamic("B_s1", "int32")
- C_3_s0 = T.dynamic("C_3_s0", "int32")
- C_3_s1 = T.dynamic("C_3_s1", "int32")
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((1, m, 22016), 'int8'), B: T.Tensor((4096, 22016), "int8"), matmul_1: T.Tensor((1, m, 4096), 'int32')):
- T.func_attr({"op_pattern": 4, "tirx.is_scheduled": True, "tirx.noalias": True})
-
- # with Ts.sblock("root"):
- A_reindex_pad_shared_dyn = Ts.sblock_alloc_buffer((1, (m + 127) // 128 * 128, 22016), "int8", scope="shared.dyn")
- B_reindex_shared_dyn = Ts.sblock_alloc_buffer((1, 4096, 22016), "int8", scope="shared.dyn")
- A_reindex_pad_shared_dyn_wmma_matrix_a = Ts.sblock_alloc_buffer((1, (m + 127) // 128 * 128, 22016), "int8", scope="wmma.matrix_a")
- B_reindex_shared_dyn_wmma_matrix_b = Ts.sblock_alloc_buffer((1, 4096, 22016), "int8", scope="wmma.matrix_b")
- matmul_1_reindex_pad_shared_dyn = Ts.sblock_alloc_buffer((1, (m + 127) // 128 * 128, 4096), "int32", scope="shared.dyn")
- matmul_1_reindex_pad_shared_dyn_wmma_accumulator = Ts.sblock_alloc_buffer((1, (m + 127) // 128 * 128, 4096), "int32", scope="wmma.accumulator")
- for ax0 in T.thread_binding(1, thread="blockIdx.z"):
- for ax1_0_0_ax2_0_0_fused in T.thread_binding((m + 127) // 128, thread="blockIdx.x"):
- for ax1_0_1_ax2_0_1_fused in T.thread_binding(32, thread="blockIdx.y"):
- for ax2_0_2_ax1_0_2_fused in T.thread_binding(16, thread="threadIdx.y"):
- for ax1_0_3_init, ax2_0_3_init in T.grid(2, 2):
- with Ts.sblock("matmul_o_init"):
- v0_o = Ts.axis.spatial(1, ax0)
- v1_o = Ts.axis.spatial((m + 127) // 128 * 8, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3_init)
- v2_o = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3_init)
- Ts.reads()
- Ts.writes(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- with Ts.sblock("matmul_init_o"):
- v1_i_init_o = Ts.axis.spatial(1, 0)
- v2_i_init_o = Ts.axis.spatial(1, 0)
- Ts.reads()
- Ts.writes(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- C = Ts.match_buffer(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=(C_s0, C_s1), scope="wmma.accumulator", offset_factor=16)
- T.gpu_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0))
- for ax3_0_0 in T.serial(1376, annotations={"software_pipeline_order": [0, 3, 1, 4, 5, 2, 6], "software_pipeline_stage": [0, 0, 0, 0, 0, 1, 1]}):
- for ax0_ax1_fused_0 in range(1):
- for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
- for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_3 in T.vectorized(4):
- with Ts.sblock("A_reindex_pad_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial((m + 127) // 128 * 128, ax1_0_0_ax2_0_0_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 16)
- v2 = Ts.axis.spatial(22016, ax3_0_0 * 16 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 16)
- Ts.reads(A[v0, v1, v2])
- Ts.writes(A_reindex_pad_shared_dyn[v0, v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 32, 16]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
- A_reindex_pad_shared_dyn[v0, v1, v2] = T.if_then_else(v1 < m, A[v0, v1, v2], T.int8(0))
- for ax0_ax1_fused_0 in range(1):
- for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
- for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_3 in T.vectorized(4):
- with Ts.sblock("B_reindex_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial(4096, ax1_0_1_ax2_0_1_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 16)
- v2 = Ts.axis.spatial(22016, ax3_0_0 * 16 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 16)
- Ts.reads(B[v1, v2])
- Ts.writes(B_reindex_shared_dyn[v0, v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 32, 16]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
- B_reindex_shared_dyn[v0, v1, v2] = B[v1, v2]
- for ax3_0_1 in T.serial(1, annotations={"software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, 1]}):
- for ax0_0 in T.unroll(2):
- for ax1_0 in T.unroll(1):
- with Ts.sblock("A_reindex_pad_shared.dyn_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(8 * ((m + 127) // 128), ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(1376, ax3_0_0 + ax1_0)
- Ts.reads(A_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(A_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A_1 = Ts.match_buffer(A_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(A_s0, A_s1), scope="shared.dyn", offset_factor=16)
- C = Ts.match_buffer(A_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(C_1_s0, C_1_s1), scope="wmma.matrix_a", offset_factor=16)
- T.gpu_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.access_ptr("int8", A_1.data, A_1.elem_offset, A_1.strides[0] * 16, 1), A_1.strides[0], "row_major")
- for ax0_0 in T.unroll(2):
- for ax1_0 in T.unroll(1):
- with Ts.sblock("B_reindex_shared.dyn_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(1376, ax3_0_0 + ax1_0)
- Ts.reads(B_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(B_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A_1 = Ts.match_buffer(B_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(A_1_s0, A_1_s1), scope="shared.dyn", offset_factor=16)
- C = Ts.match_buffer(B_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=(C_2_s0, C_2_s1), scope="wmma.matrix_b", offset_factor=16)
- T.gpu_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.access_ptr("int8", A_1.data, A_1.elem_offset, A_1.strides[0] * 16, 1), A_1.strides[0], "col_major")
- for ax1_0_3, ax2_0_3 in T.grid(2, 2):
- with Ts.sblock("matmul_o_update"):
- v0_o = Ts.axis.spatial(1, ax0)
- v1_o = Ts.axis.spatial((m + 127) // 128 * 8, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3)
- v2_o = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3)
- v3_o = Ts.axis.reduce(1376, ax3_0_0 + ax3_0_1)
- Ts.reads(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], A_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], B_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- with Ts.sblock("matmul_o"):
- v1_i_o = Ts.axis.spatial(1, 0)
- v2_i_o = Ts.axis.spatial(1, 0)
- v3_i_o = Ts.axis.reduce(1, 0)
- Ts.reads(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], A_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], B_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A_1 = Ts.match_buffer(A_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=(A_2_s0, A_2_s1), scope="wmma.matrix_a", offset_factor=16)
- B_1 = Ts.match_buffer(B_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=(B_s0, B_s1), scope="wmma.matrix_b", offset_factor=16)
- C = Ts.match_buffer(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=(C_3_s0, C_3_s1), scope="wmma.accumulator", offset_factor=16)
- T.gpu_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A_1.data, A_1.elem_offset // A_1.strides[0] // 16 * (A_1.strides[0] // 16) + A_1.elem_offset % A_1.strides[0] // 16, B_1.data, B_1.elem_offset // B_1.strides[0] // 16 * (B_1.strides[0] // 16) + B_1.elem_offset % B_1.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides [...]
- for ax0_0, ax1_0 in T.grid(2, 2):
- with Ts.sblock("matmul_1_reindex_pad_shared.dyn_wmma.accumulator_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(8 * ((m + 127) // 128), ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
- v2_o = Ts.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax1_0)
- Ts.reads(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- Ts.writes(matmul_1_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
- A_1 = Ts.match_buffer(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=(A_3_s0, A_3_s1), scope="wmma.accumulator", offset_factor=16)
- C = Ts.match_buffer(matmul_1_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=(C_4_s0, C_4_s1), scope="shared.dyn", offset_factor=16)
- T.gpu_store_matrix_sync(A_1.data, 16, 16, 16, A_1.elem_offset // A_1.strides[0] // 16 * (A_1.strides[0] // 16) + A_1.elem_offset % A_1.strides[0] // 16, T.access_ptr("int32", C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
- for ax0_ax1_fused_0 in range(8):
- for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_2 in T.vectorized(4):
- with Ts.sblock("matmul_1_reindex_pad_shared.dyn"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial((m + 127) // 128 * 128, ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 32)
- v2 = Ts.axis.spatial(4096, ax1_0_1_ax2_0_1_fused * 128 + ax2_0_2_ax1_0_2_fused // 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 32)
- Ts.where(ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + ((ax0_ax1_fused_0 * 32 + ax0_ax1_fused_1) * 4 + ax0_ax1_fused_2) // 32 < m)
- Ts.reads(matmul_1_reindex_pad_shared_dyn[v0, v1, v2])
- Ts.writes(matmul_1[0, v1, v2])
- Ts.sblock_attr({"buffer_dim_align": [[0, 1, 16, 4]]})
- matmul_1[0, v1, v2] = matmul_1_reindex_pad_shared_dyn[v0, v1, v2]
- # fmt: on
-
- mod = tvm.IRModule({"main": before})
- with Target("nvidia/geforce-rtx-2080-ti"):
- mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
- tvm.ir.assert_structural_equal(mod["main"], expected)
-
-
def test_matmul_metal():
# fmt: off
batch_size = T.dynamic("batch_size", "int32")
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
index 184fc2c10f..01842232a8 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
@@ -372,388 +372,12 @@ s1 = T.dynamic("s1", "int32")
s1_1 = T.dynamic("s1_1", "int32")
s1_2 = T.dynamic("s1_2", "int32")
-@Ts.function
-def GMMCUDATensorCore(
- X: T.Tensor((1024, 1024), "float16"),
- Y: T.Tensor((1024, 1024), "float16"),
- Z: T.Tensor((1024, 1024), "float32"),
-) -> None:
- # function attr dict
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- # body
- # with Ts.sblock("root")
- Z_wmma_accumulator = Ts.sblock_alloc_buffer([1024, 1024], dtype="float32", scope="wmma.accumulator")
- X_shared = Ts.sblock_alloc_buffer([1024, 1024], dtype="float16", scope="shared")
- Y_shared = Ts.sblock_alloc_buffer([1024, 1024], dtype="float16", scope="shared")
- X_shared_wmma_matrix_a = Ts.sblock_alloc_buffer([1024, 1024], dtype="float16", scope="wmma.matrix_a")
- Y_shared_wmma_matrix_b = Ts.sblock_alloc_buffer([1024, 1024], dtype="float16", scope="wmma.matrix_b")
- for ax0_0_ax1_0_0_ax2_0_0_fused in T.thread_binding(64, thread="blockIdx.x"):
- for ax0_1_ax1_0_1_ax2_0_1_fused in T.thread_binding(2, thread="blockIdx.y"):
- for ax0_2_ax1_0_2_ax2_0_2_fused in T.thread_binding(2, thread="threadIdx.y"):
- for ax1_0_3_init, ax2_0_3_init, ax1_0_4_init, ax2_0_4_init in T.grid(2, 1, 2, 4):
- with Ts.sblock("Z_o_init"):
- v0 = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(
- 64,
- ax0_0_ax1_0_0_ax2_0_0_fused % 64 // 16 * 16
- + ax0_1_ax1_0_1_ax2_0_1_fused % 2 * 8
- + ax0_2_ax1_0_2_ax2_0_2_fused % 2 * 4
- + ax1_0_3_init * 2
- + ax1_0_4_init,
- )
- v2_o = Ts.axis.spatial(
- 64,
- (ax0_0_ax1_0_0_ax2_0_0_fused % 16 + 0 + 0 + ax2_0_3_init) * 4
- + ax2_0_4_init,
- )
- Ts.reads()
- Ts.writes(
- Z_wmma_accumulator[
- v1_o * 16 : v1_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
- ]
- )
- Ts.sblock_attr(
- {
- "meta_schedule.thread_extent_high_inclusive": 1024,
- "meta_schedule.thread_extent_low_inclusive": 32,
- "warp_execution": 1,
- }
- )
- C = Ts.match_buffer(
- Z_wmma_accumulator[
- v1_o * 16 : v1_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
- ],
- [16, 16],
- dtype="float32",
- scope="wmma.accumulator",
- offset_factor=16,
- )
- T.evaluate(
- T.gpu_fill_fragment(
- C.data,
- 16,
- 16,
- 16,
- C.elem_offset // 256 + C.elem_offset % 256 // 16,
- T.float32(0),
-
- )
- )
- for ax3_0_0 in T.serial(32):
- for ax0_ax1_fused_0 in T.serial(16):
- for ax0_ax1_fused_1 in T.thread_binding(2, thread="threadIdx.y"):
- for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_3 in T.vectorized(4):
- with Ts.sblock("X_shared"):
- v0 = Ts.axis.spatial(
- 1024,
- ax0_0_ax1_0_0_ax2_0_0_fused // 16 * 256
- + ax0_1_ax1_0_1_ax2_0_1_fused * 128
- + (
- ax0_ax1_fused_0 * 256
- + ax0_ax1_fused_1 * 128
- + ax0_ax1_fused_2 * 4
- + ax0_ax1_fused_3
- )
- // 32,
- )
- v1 = Ts.axis.spatial(
- 1024,
- ax3_0_0 * 32
- + (
- ax0_ax1_fused_0 * 256
- + ax0_ax1_fused_1 * 128
- + ax0_ax1_fused_2 * 4
- + ax0_ax1_fused_3
- )
- % 32,
- )
- Ts.reads(X[v0, v1])
- Ts.writes(X_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]]})
- X_shared[v0, v1] = X[v0, v1]
- for ax0_ax1_fused_0 in T.serial(8):
- for ax0_ax1_fused_1 in T.thread_binding(2, thread="threadIdx.y"):
- for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
- for ax0_ax1_fused_3 in T.vectorized(4):
- with Ts.sblock("Y_shared"):
- v0 = Ts.axis.spatial(
- 1024,
- ax3_0_0 * 32
- + (
- ax0_ax1_fused_0 * 256
- + ax0_ax1_fused_1 * 128
- + ax0_ax1_fused_2 * 4
- + ax0_ax1_fused_3
- )
- // 64,
- )
- v1 = Ts.axis.spatial(
- 1024,
- ax0_0_ax1_0_0_ax2_0_0_fused % 16 * 64
- + (
- ax0_ax1_fused_0 * 256
- + ax0_ax1_fused_1 * 128
- + ax0_ax1_fused_2 * 4
- + ax0_ax1_fused_3
- )
- % 64,
- )
- Ts.reads(Y[v0, v1])
- Ts.writes(Y_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]]})
- Y_shared[v0, v1] = Y[v0, v1]
- for ax3_0_1 in T.serial(2):
- for ax0_0, ax1_0 in T.grid(4, 1):
- with Ts.sblock("X_shared_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(
- 64,
- ax0_0_ax1_0_0_ax2_0_0_fused // 16 * 16
- + ax0_1_ax1_0_1_ax2_0_1_fused * 8
- + ax0_2_ax1_0_2_ax2_0_2_fused * 4
- + ax0_0,
- )
- v1_o = Ts.axis.spatial(64, ax3_0_0 * 2 + ax3_0_1)
- Ts.reads(
- X_shared[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16]
- )
- Ts.writes(
- X_shared_wmma_matrix_a[
- v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
- ]
- )
- A = Ts.match_buffer(
- X_shared[
- v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
- ],
- [16, 16],
- dtype="float16",
- strides=[s1, s0],
- scope="shared",
- offset_factor=16,
- )
- C_1 = Ts.match_buffer(
- X_shared_wmma_matrix_a[
- v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
- ],
- [16, 16],
- dtype="float16",
- scope="wmma.matrix_a",
- offset_factor=16,
- )
- T.evaluate(
- T.gpu_load_matrix_sync(
- C_1.data,
- 16,
- 16,
- 16,
- C_1.elem_offset // 256 + C_1.elem_offset % 256 // 16,
- T.access_ptr(
- "float16",
- A.data,
- A.elem_offset,
- s1 * 16,
- 1,
-
- ),
- s1,
- "row_major",
-
- )
- )
- for ax0_0, ax1_0 in T.grid(1, 4):
- with Ts.sblock("Y_shared_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(64, ax3_0_0 * 2 + ax3_0_1)
- v1_o = Ts.axis.spatial(
- 64, ax0_0_ax1_0_0_ax2_0_0_fused % 16 * 4 + ax1_0
- )
- Ts.reads(
- Y_shared[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16]
- )
- Ts.writes(
- Y_shared_wmma_matrix_b[
- v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
- ]
- )
- A_1 = Ts.match_buffer(
- Y_shared[
- v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
- ],
- [16, 16],
- dtype="float16",
- strides=[s1_1, s0_1],
- scope="shared",
- offset_factor=16,
- )
- C_2 = Ts.match_buffer(
- Y_shared_wmma_matrix_b[
- v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
- ],
- [16, 16],
- dtype="float16",
- scope="wmma.matrix_b",
- offset_factor=16,
- )
- T.evaluate(
- T.gpu_load_matrix_sync(
- C_2.data,
- 16,
- 16,
- 16,
- C_2.elem_offset // 256 + C_2.elem_offset % 256 // 16,
- T.access_ptr(
- "float16",
- A_1.data,
- A_1.elem_offset,
- s1_1 * 16,
- 1,
-
- ),
- s1_1,
- "row_major",
-
- )
- )
- for ax0_3, ax1_0_3, ax2_0_3, ax3_0_2, ax0_4, ax1_0_4, ax2_0_4 in T.grid(
- 1, 2, 1, 1, 1, 2, 4
- ):
- with Ts.sblock("Z_o_update"):
- v0 = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(
- 64,
- ax0_0_ax1_0_0_ax2_0_0_fused % 64 // 16 * 16
- + ax0_1_ax1_0_1_ax2_0_1_fused % 2 * 8
- + ax0_2_ax1_0_2_ax2_0_2_fused % 2 * 4
- + ax1_0_3 * 2
- + ax1_0_4,
- )
- v2_o = Ts.axis.spatial(
- 64,
- (ax0_0_ax1_0_0_ax2_0_0_fused % 16 + 0 + 0 + ax2_0_3) * 4
- + ax2_0_4,
- )
- v3_o = Ts.axis.reduce(64, ax3_0_0 * 2 + ax3_0_1 + ax3_0_2)
- Ts.reads(
- Z_wmma_accumulator[
- v1_o * 16 : v1_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
- ],
- X_shared_wmma_matrix_a[
- v1_o * 16 : v1_o * 16 + 16, v3_o * 16 : v3_o * 16 + 16
- ],
- Y_shared_wmma_matrix_b[
- v3_o * 16 : v3_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
- ],
- )
- Ts.writes(
- Z_wmma_accumulator[
- v1_o * 16 : v1_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
- ]
- )
- Ts.sblock_attr(
- {
- "meta_schedule.thread_extent_high_inclusive": 1024,
- "meta_schedule.thread_extent_low_inclusive": 32,
- "warp_execution": 1,
- }
- )
- A_2 = Ts.match_buffer(
- X_shared_wmma_matrix_a[
- v1_o * 16 : v1_o * 16 + 16, v3_o * 16 : v3_o * 16 + 16
- ],
- [16, 16],
- dtype="float16",
- scope="wmma.matrix_a",
- offset_factor=16,
- )
- B = Ts.match_buffer(
- Y_shared_wmma_matrix_b[
- v3_o * 16 : v3_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
- ],
- [16, 16],
- dtype="float16",
- scope="wmma.matrix_b",
- offset_factor=16,
- )
- C_3 = Ts.match_buffer(
- Z_wmma_accumulator[
- v1_o * 16 : v1_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
- ],
- [16, 16],
- dtype="float32",
- scope="wmma.accumulator",
- offset_factor=16,
- )
- T.evaluate(
- T.gpu_mma_sync(
- C_3.data,
- C_3.elem_offset // 256 + C_3.elem_offset % 256 // 16,
- A_2.data,
- A_2.elem_offset // 256,
- B.data,
- B.elem_offset // 256,
- C_3.data,
- C_3.elem_offset // 256 + C_3.elem_offset % 256 // 16,
-
- )
- )
- for ax0_0, ax1_0 in T.grid(4, 4):
- with Ts.sblock("Z_wmma.accumulator_o"):
- v0_o = Ts.axis.spatial(
- 64,
- ax0_0_ax1_0_0_ax2_0_0_fused // 16 * 16
- + ax0_1_ax1_0_1_ax2_0_1_fused * 8
- + ax0_2_ax1_0_2_ax2_0_2_fused * 4
- + ax0_0,
- )
- v1_o = Ts.axis.spatial(64, ax0_0_ax1_0_0_ax2_0_0_fused % 16 * 4 + ax1_0)
- Ts.reads(
- Z_wmma_accumulator[
- v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
- ]
- )
- Ts.writes(Z[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16])
- A_3 = Ts.match_buffer(
- Z_wmma_accumulator[
- v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
- ],
- [16, 16],
- dtype="float32",
- scope="wmma.accumulator",
- offset_factor=16,
- )
- C_4 = Ts.match_buffer(
- Z[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16],
- [16, 16],
- dtype="float32",
- strides=[s1_2, s0_2],
- offset_factor=16,
- )
- T.evaluate(
- T.gpu_store_matrix_sync(
- A_3.data,
- 16,
- 16,
- 16,
- A_3.elem_offset // 256 + A_3.elem_offset % 256 // 16,
- T.access_ptr(
- "float32",
- C_4.data,
- C_4.elem_offset,
- s1_2 * 16,
- 2,
-
- ),
- s1_2,
- "row_major",
-
- )
- )
# fmt: on
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,not-callable,misplaced-comparison-constant
-@pytest.mark.parametrize("mod", [Conv2dCuda0, Conv2dCuda1, GmmCuda0, GMMCUDATensorCore])
+@pytest.mark.parametrize("mod", [Conv2dCuda0, Conv2dCuda1, GmmCuda0])
def test_postproc_check_pass(mod):
ctx = _create_context(mod, target=_target())
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_tc.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_tc.py
deleted file mode 100644
index 3c01392321..0000000000
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_tc.py
+++ /dev/null
@@ -1,1511 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring,line-too-long,invalid-name,too-many-locals,too-many-statements,too-many-nested-blocks,too-many-branches,too-many-lines,chained-comparison
-# ruff: noqa: E501, F401, F841, RUF005
-
-import pytest
-
-import tvm
-import tvm.testing
-from tvm import te
-from tvm.s_tir import meta_schedule as ms
-from tvm.s_tir.meta_schedule.testing import te_workload
-from tvm.s_tir.meta_schedule.testing.space_generation import (
- check_sketches,
- generate_design_space,
- get_rules,
- print_sketches,
-)
-from tvm.s_tir.tensor_intrin.cuda import get_wmma_intrin_group
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-
-def multi_level_tiling_tensor_core(
- *,
- read_reuse_scope="shared",
- write_reuse_scope="shared",
- in_dtype="float16",
- out_dtype="float32",
- trans_b=False,
- use_software_pipeline=False,
-) -> ms.schedule_rule.ScheduleRule:
- assert read_reuse_scope in ["shared", "shared.dyn"]
- assert write_reuse_scope in ["shared", "shared.dyn", "global"]
- if not isinstance(in_dtype, list):
- in_dtype = [in_dtype]
- if not isinstance(out_dtype, list):
- out_dtype = [out_dtype]
- if not isinstance(trans_b, list):
- trans_b = [trans_b]
- return ms.schedule_rule.MultiLevelTilingTensorCore(
- intrin_groups=[
- get_wmma_intrin_group(
- read_reuse_scope, write_reuse_scope, _in_dtype, _out_dtype, _trans_b
- )
- for _in_dtype in in_dtype
- for _out_dtype in out_dtype
- for _trans_b in trans_b
- ],
- structure="SSSRRSRS",
- tile_binds=["blockIdx.y", "blockIdx.x", "threadIdx.y"],
- max_innermost_factor=4, # 64 // tensor intrin size
- vector_load_lens=[1, 2, 3, 4, 8, 16],
- reuse_read=ms.schedule_rule.ReuseType(
- req="must",
- levels=[4],
- scope=read_reuse_scope,
- ),
- reuse_write=ms.schedule_rule.ReuseType(
- req="must" if write_reuse_scope.startswith("shared") else "no",
- levels=[2],
- scope=write_reuse_scope,
- ),
- use_software_pipeline=use_software_pipeline,
- )
-
-
-@pytest.mark.parametrize("shared_scope", ["shared", "shared.dyn"])
-def test_matmul_relu(shared_scope):
- intrin_suffix = shared_scope.replace(".", "_")
-
- # fmt: off
- @Ts.function
- def matmul_relu_0(A: T.Tensor((128, 128), "float16"), B: T.Tensor((128, 128), "float16"), compute: T.Tensor((128, 128), "float32")) -> None:
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- # with Ts.sblock("root"):
- C_reindex_shared = Ts.sblock_alloc_buffer((4, 8, 2, 1, 16, 16), scope=shared_scope)
- C_reindex_shared_wmma_accumulator = Ts.sblock_alloc_buffer((4, 8, 2, 1, 16, 16), scope="wmma.accumulator")
- A_reindex_shared = Ts.sblock_alloc_buffer((128, 128), "float16", scope=shared_scope)
- B_reindex_shared = Ts.sblock_alloc_buffer((128, 128), "float16", scope=shared_scope)
- A_reindex_shared_wmma_matrix_a = Ts.sblock_alloc_buffer((128, 128), "float16", scope="wmma.matrix_a")
- B_reindex_shared_wmma_matrix_b = Ts.sblock_alloc_buffer((128, 128), "float16", scope="wmma.matrix_b")
- for ax0_0_0_ax1_0_0_fused in T.thread_binding(8, 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(1):
- for ax0_ax1_fused in range(4096):
- with Ts.sblock("A_reindex_shared"):
- v0 = Ts.axis.spatial(128, ax0_0_0_ax1_0_0_fused // 2 * 32 + ax0_ax1_fused // 128)
- v1 = Ts.axis.spatial(128, ax0_ax1_fused % 128)
- Ts.reads(A[v0, v1])
- Ts.writes(A_reindex_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 8})
- A_reindex_shared[v0, v1] = A[v0, v1]
- for ax0_ax1_fused in range(4096):
- with Ts.sblock("B_reindex_shared"):
- v0 = Ts.axis.spatial(128, ax0_ax1_fused // 32)
- v1 = Ts.axis.spatial(128, ax0_0_0_ax1_0_0_fused % 2 * 64 + ax0_0_1_ax1_0_1_fused * 32 + ax0_ax1_fused % 32)
- Ts.reads(B[v0, v1])
- Ts.writes(B_reindex_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 1})
- B_reindex_shared[v0, v1] = B[v0, v1]
- for ax2_0_1 in range(4):
- for ax0_0, ax1_0 in T.grid(2, 2):
- with Ts.sblock("A_reindex_shared_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused // 2 * 2 + ax0_0)
- v1_o = Ts.axis.spatial(8, ax2_0_1 * 2 + ax1_0)
- Ts.reads(A_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(A_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": f"wmma_load_16x16x16_f16_a_{intrin_suffix}"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("A_reindex_shared_wmma.matrix_a"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(A_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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(2, 1):
- with Ts.sblock("B_reindex_shared_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(8, ax2_0_1 * 2 + ax0_0)
- v1_o = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused % 2 * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_0_2_ax1_0_2_fused + ax1_0)
- Ts.reads(B_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(B_reindex_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": f"wmma_load_16x16x16_f16_b_{intrin_suffix}"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("B_reindex_shared_wmma.matrix_b"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(B_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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, 2, 2, 1):
- with Ts.sblock("C_o"):
- v0_o = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused // 2 * 2 + ax0_0_3 * 2 + ax0_0_4)
- v1_o = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused % 2 * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_0_2_ax1_0_2_fused + ax1_0_3 + ax1_0_4)
- v2_o = Ts.axis.reduce(8, ax2_0_0 * 8 + ax2_0_1 * 2 + ax2_0_2)
- Ts.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])
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init": "wmma_fill_16x16x16_f32", "warp_execution": 1})
- with Ts.init():
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("C_init"):
- v0_i_init, v1_i_init = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads()
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, v0_i_init, v1_i_init])
- C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, v0_i_init, v1_i_init] = T.float32(0)
- for ax0_1, ax1_1, ax2_1 in T.grid(16, 16, 16):
- with Ts.sblock("C"):
- v0_i, v1_i, v2_i = Ts.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
- Ts.reads(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, 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])
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, v0_i, v1_i])
- Ts.sblock_attr({"meta_schedule.tiling_structure": "SSSRRSRS"})
- C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, v0_i, v1_i] = C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, 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, 1):
- with Ts.sblock("C_reindex_shared_wmma.accumulator_o"):
- v0 = Ts.axis.spatial(4, ax0_0_0_ax1_0_0_fused // 2)
- v1 = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused % 2 * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_ax1_fused)
- v2 = Ts.axis.spatial(2, ax2 + ax2_1)
- v3 = Ts.axis.spatial(1, ax3)
- v4_o = Ts.axis.spatial(1, 0)
- v5_o = Ts.axis.spatial(1, 0)
- Ts.reads(C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, 0:16, 0:16])
- Ts.writes(C_reindex_shared[v0, v1, v2, v3, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": f"wmma_store_16x16x16_f32_{intrin_suffix}"})
- for ax4, ax5 in T.grid(16, 16):
- with Ts.sblock("C_reindex_shared_wmma.accumulator"):
- v4_i, v5_i = Ts.axis.remap("SS", [ax4, ax5])
- Ts.reads(C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, v4_i, v5_i])
- Ts.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(512):
- with Ts.sblock("C_reindex_shared"):
- v0 = Ts.axis.spatial(4, ax0_0_0_ax1_0_0_fused // 2)
- v1 = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused % 2 * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_ax1_ax3_ax4_ax5_fused // 256)
- v2 = Ts.axis.spatial(2, ax2)
- v3 = Ts.axis.spatial(1, 0)
- v4 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 256 // 16)
- v5 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
- Ts.reads(C_reindex_shared[v0, v1, v2, v3, v4, v5])
- Ts.writes(compute[v4 + v2 * 16 + v0 * 32, v5 + v1 * 16])
- Ts.sblock_attr({"meta_schedule.cooperative_fetch": 4})
- compute[v4 + v2 * 16 + v0 * 32, v5 + v1 * 16] = T.max(C_reindex_shared[v0, v1, v2, v3, v4, v5], T.float32(0))
- # fmt: on
- decision_0 = [
- ("SamplePerfectTile", [4, 1, 1, 1, 2]),
- ("SamplePerfectTile", [2, 2, 2, 1, 1]),
- ("SamplePerfectTile", [1, 4, 2]),
- ("SampleCategorical", 3),
- ("SampleCategorical", 3),
- ("SampleCategorical", 0),
- ]
-
- mod = te.create_function(
- 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({"kind": "cuda", "arch": "sm_70"}),
- types=None,
- sch_rules=[
- multi_level_tiling_tensor_core(
- read_reuse_scope=shared_scope, write_reuse_scope=shared_scope
- ),
- ]
- + get_rules(kind="cuda", types=ms.schedule_rule.AutoInline),
- )
- check_sketches(
- mod,
- sketches=actual,
- expected_mods=[matmul_relu_0],
- expected_decisions=[decision_0],
- )
-
-
-def test_matmul_relu_with_fallback():
- # fmt: off
- @Ts.function
- def matmul_relu_fallback_0(A: T.Tensor((128, 128), "float16"), B: T.Tensor((128, 128), "float16"), compute: T.Tensor((128, 128), "float32")) -> None:
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- # with Ts.sblock("root"):
- C_reindex_shared = Ts.sblock_alloc_buffer((4, 2, 2, 4, 16, 16), scope="shared")
- C_reindex_shared_wmma_accumulator = Ts.sblock_alloc_buffer((4, 2, 2, 4, 16, 16), scope="wmma.accumulator")
- A_reindex_shared = Ts.sblock_alloc_buffer((128, 128), "float16", scope="shared")
- B_reindex_shared = Ts.sblock_alloc_buffer((128, 128), "float16", scope="shared")
- A_reindex_shared_wmma_matrix_a = Ts.sblock_alloc_buffer((128, 128), "float16", scope="wmma.matrix_a")
- B_reindex_shared_wmma_matrix_b = Ts.sblock_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 Ts.sblock("A_reindex_shared"):
- v0 = Ts.axis.spatial(128, ax0_0_0_ax1_0_0_fused * 64 + ax0_0_1_ax1_0_1_fused * 32 + ax0_ax1_fused // 64)
- v1 = Ts.axis.spatial(128, ax2_0_0 * 64 + ax0_ax1_fused % 64)
- Ts.reads(A[v0, v1])
- Ts.writes(A_reindex_shared[v0, v1])
- Ts.sblock_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 Ts.sblock("B_reindex_shared"):
- v0 = Ts.axis.spatial(128, ax2_0_0 * 64 + ax0_ax1_fused // 128)
- v1 = Ts.axis.spatial(128, ax0_ax1_fused % 128)
- Ts.reads(B[v0, v1])
- Ts.writes(B_reindex_shared[v0, v1])
- Ts.sblock_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 Ts.sblock("A_reindex_shared_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_0)
- v1_o = Ts.axis.spatial(8, ax2_0_0 * 4 + ax1_0)
- Ts.reads(A_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(A_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_a_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("A_reindex_shared_wmma.matrix_a"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(A_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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 Ts.sblock("B_reindex_shared_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(8, ax2_0_0 * 4 + ax0_0)
- v1_o = Ts.axis.spatial(8, ax0_0_2_ax1_0_2_fused * 4 + ax1_0)
- Ts.reads(B_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(B_reindex_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_b_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("B_reindex_shared_wmma.matrix_b"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(B_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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 Ts.sblock("C_o"):
- v0_o = Ts.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 = Ts.axis.spatial(8, ax0_0_2_ax1_0_2_fused * 4 + ax1_0_3 * 4 + ax1_0_4)
- v2_o = Ts.axis.reduce(8, ax2_0_0 * 4 + ax2_0_1 * 4 + ax2_0_2)
- Ts.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])
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init": "wmma_fill_16x16x16_f32", "warp_execution": 1})
- with Ts.init():
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("C_init"):
- v0_i_init, v1_i_init = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads()
- Ts.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 Ts.sblock("C"):
- v0_i, v1_i, v2_i = Ts.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
- Ts.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])
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, v0_i, v1_i])
- Ts.sblock_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 Ts.sblock("C_reindex_shared_wmma.accumulator_o"):
- v0 = Ts.axis.spatial(4, ax0_0_0_ax1_0_0_fused * 2 + ax0_0_1_ax1_0_1_fused)
- v1 = Ts.axis.spatial(2, ax0_ax1_fused)
- v2 = Ts.axis.spatial(2, ax2 + ax2_1)
- v3 = Ts.axis.spatial(4, ax3)
- v4_o = Ts.axis.spatial(1, 0)
- v5_o = Ts.axis.spatial(1, 0)
- Ts.reads(C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, 0:16, 0:16])
- Ts.writes(C_reindex_shared[v0, v1, v2, v3, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_store_16x16x16_f32_shared"})
- for ax4, ax5 in T.grid(16, 16):
- with Ts.sblock("C_reindex_shared_wmma.accumulator"):
- v4_i, v5_i = Ts.axis.remap("SS", [ax4, ax5])
- Ts.reads(C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, v4_i, v5_i])
- Ts.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 Ts.sblock("C_reindex_shared"):
- v0 = Ts.axis.spatial(4, ax0_0_0_ax1_0_0_fused * 2 + ax0_0_1_ax1_0_1_fused)
- v1 = Ts.axis.spatial(2, ax0_ax1_ax3_ax4_ax5_fused // 1024)
- v2 = Ts.axis.spatial(2, ax2)
- v3 = Ts.axis.spatial(4, ax0_ax1_ax3_ax4_ax5_fused % 1024 // 256)
- v4 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 256 // 16)
- v5 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
- Ts.reads(C_reindex_shared[v0, v1, v2, v3, v4, v5])
- Ts.writes(compute[v4 + v2 * 16 + v0 * 32, v5 + v3 * 16 + v1 * 64])
- Ts.sblock_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_function(
- 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({"kind": "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],
- )
-
-
-@pytest.mark.parametrize("shared_scope", ["shared", "shared.dyn"])
-def test_conv2d(shared_scope):
- intrin_suffix = shared_scope.replace(".", "_")
-
- # fmt: off
- @Ts.function
- def conv2d_0(inputs: T.Tensor((1, 16, 16, 32), "float16"), weight: T.Tensor((3, 3, 32, 32), "float16"), conv2d_nhwc: T.Tensor((1, 16, 16, 32), "float32")):
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- # with Ts.sblock("root"):
- PadInput = Ts.sblock_alloc_buffer((1, 18, 18, 32), "float16")
- conv2d_nhwc_reindex_shared_dyn = Ts.sblock_alloc_buffer((16, 2, 1, 1, 16, 16), scope=shared_scope)
- conv2d_nhwc_reindex_shared_dyn_wmma_accumulator = Ts.sblock_alloc_buffer((16, 2, 1, 1, 16, 16), scope="wmma.accumulator")
- PadInput_reindex_shared_dyn = Ts.sblock_alloc_buffer((256, 288), "float16", scope=shared_scope)
- weight_reindex_shared_dyn = Ts.sblock_alloc_buffer((288, 32), "float16", scope=shared_scope)
- PadInput_reindex_shared_dyn_wmma_matrix_a = Ts.sblock_alloc_buffer((256, 288), "float16", scope="wmma.matrix_a")
- weight_reindex_shared_dyn_wmma_matrix_b = Ts.sblock_alloc_buffer((288, 32), "float16", scope="wmma.matrix_b")
- for i0, i1, i2, i3 in T.grid(1, 18, 18, 32):
- with Ts.sblock("PadInput"):
- v_i0, v_i1, v_i2, v_i3 = Ts.axis.remap("SSSS", [i0, i1, i2, i3])
- Ts.reads(inputs[v_i0, v_i1 - 1, v_i2 - 1, v_i3])
- Ts.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 Ts.sblock("PadInput_reindex_shared.dyn"):
- v0 = Ts.axis.spatial(256, ax0_0_1_ax1_0_1_fused * 16 + ax0_ax1_fused // 288)
- v1 = Ts.axis.spatial(288, ax0_ax1_fused % 288)
- Ts.reads(PadInput[0, v0 // 16 + v1 // 96, v0 % 16 + v1 % 96 // 32, v1 % 32])
- Ts.writes(PadInput_reindex_shared_dyn[v0, v1])
- Ts.sblock_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 Ts.sblock("weight_reindex_shared.dyn"):
- v0 = Ts.axis.spatial(288, ax0_ax1_fused // 16)
- v1 = Ts.axis.spatial(32, ax0_0_0_ax1_0_0_fused * 16 + ax0_ax1_fused % 16)
- Ts.reads(weight[v0 // 96, v0 % 96 // 32, v0 % 32, v1])
- Ts.writes(weight_reindex_shared_dyn[v0, v1])
- Ts.sblock_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 Ts.sblock("PadInput_reindex_shared.dyn_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(16, ax0_0_1_ax1_0_1_fused + ax0_0)
- v1_o = Ts.axis.spatial(18, ax2_0_1 + ax1_0)
- Ts.reads(PadInput_reindex_shared_dyn[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(PadInput_reindex_shared_dyn_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": f"wmma_load_16x16x16_f16_a_{intrin_suffix}"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("PadInput_reindex_shared.dyn_wmma.matrix_a"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(PadInput_reindex_shared_dyn[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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 Ts.sblock("weight_reindex_shared.dyn_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(18, ax2_0_1 + ax0_0)
- v1_o = Ts.axis.spatial(2, ax0_0_0_ax1_0_0_fused + ax1_0)
- Ts.reads(weight_reindex_shared_dyn[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(weight_reindex_shared_dyn_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": f"wmma_load_16x16x16_f16_b_{intrin_suffix}"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("weight_reindex_shared.dyn_wmma.matrix_b"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(weight_reindex_shared_dyn[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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 Ts.sblock("conv2d_nhwc_o"):
- v0_o = Ts.axis.spatial(16, ax0_0_1_ax1_0_1_fused + ax0_0_3 + ax0_0_4)
- v1_o = Ts.axis.spatial(2, ax0_0_0_ax1_0_0_fused + ax1_0_3 + ax1_0_4)
- v2_o = Ts.axis.reduce(18, ax2_0_0 * 18 + ax2_0_1 + ax2_0_2)
- Ts.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])
- Ts.writes(conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o, 0, 0, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init": "wmma_fill_16x16x16_f32", "warp_execution": 1})
- with Ts.init():
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("conv2d_nhwc_init"):
- v0_i_init, v1_i_init = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads()
- Ts.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 Ts.sblock("conv2d_nhwc"):
- v0_i, v1_i, v2_i = Ts.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
- Ts.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])
- Ts.writes(conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o, 0, 0, v0_i, v1_i])
- Ts.sblock_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 Ts.sblock("conv2d_nhwc_reindex_shared.dyn_wmma.accumulator_o"):
- v0, v1, v2, v3 = Ts.axis.remap("SSSS", [ax0_0_1_ax1_0_1_fused, ax0_0_0_ax1_0_0_fused, ax2_1, ax3])
- v4_o = Ts.axis.spatial(1, 0)
- v5_o = Ts.axis.spatial(1, 0)
- Ts.reads(conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0, v1, v2, v3, 0:16, 0:16])
- Ts.writes(conv2d_nhwc_reindex_shared_dyn[v0, v1, v2, v3, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": f"wmma_store_16x16x16_f32_{intrin_suffix}"})
- for ax4, ax5 in T.grid(16, 16):
- with Ts.sblock("conv2d_nhwc_reindex_shared.dyn_wmma.accumulator"):
- v4_i, v5_i = Ts.axis.remap("SS", [ax4, ax5])
- Ts.reads(conv2d_nhwc_reindex_shared_dyn_wmma_accumulator[v0, v1, v2, v3, v4_i, v5_i])
- Ts.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 Ts.sblock("conv2d_nhwc_reindex_shared.dyn"):
- v0, v1, v2 = Ts.axis.remap("SSS", [ax0_0_1_ax1_0_1_fused, ax0_0_0_ax1_0_0_fused, ax2])
- v3 = Ts.axis.spatial(1, 0)
- v4 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused // 16)
- v5 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
- Ts.reads(conv2d_nhwc_reindex_shared_dyn[v0, v1, v2, v3, v4, v5])
- Ts.writes(conv2d_nhwc[0, (v4 + v0 * 16) // 16, (v4 + v0 * 16) % 16, v5 + v1 * 16])
- Ts.sblock_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_function(
- 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({"kind": "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({"kind": "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],
- )
-
-
-@pytest.mark.parametrize("shared_scope", ["shared", "shared.dyn"])
-def test_matmul_relu_pipeline(shared_scope):
- intrin_suffix = shared_scope.replace(".", "_")
-
- # fmt: off
- @Ts.function
- def matmul_relu_pipeline_0(A: T.Tensor((128, 128), "float16"), B: T.Tensor((128, 128), "float16"), compute: T.Tensor((128, 128), "float32")) -> None:
- # function attr dict
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- # body
- # with Ts.sblock("root")
- C = Ts.sblock_alloc_buffer((128, 128))
- C_reindex_shared = Ts.sblock_alloc_buffer((4, 4, 2, 2, 16, 16), scope=shared_scope)
- C_reindex_shared_wmma_accumulator = Ts.sblock_alloc_buffer((4, 4, 2, 2, 16, 16), scope="wmma.accumulator")
- A_reindex_shared = Ts.sblock_alloc_buffer((128, 128), "float16", scope=shared_scope)
- B_reindex_shared = Ts.sblock_alloc_buffer((128, 128), "float16", scope=shared_scope)
- A_reindex_shared_wmma_matrix_a = Ts.sblock_alloc_buffer((128, 128), "float16", scope="wmma.matrix_a")
- B_reindex_shared_wmma_matrix_b = Ts.sblock_alloc_buffer((128, 128), "float16", scope="wmma.matrix_b")
- for ax0_0_0_ax1_0_0_fused in T.thread_binding(1, 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 T.serial(4, annotations={"software_pipeline_order": [0, 3, 1, 4, 5, 2, 6], "software_pipeline_stage": [0, 0, 0, 0, 0, 1, 1]}):
- for ax0_ax1_fused in range(1024):
- with Ts.sblock("A_reindex_shared"):
- v0 = Ts.axis.spatial(128, ax0_0_1_ax1_0_1_fused // 4 * 32 + ax0_ax1_fused // 32)
- v1 = Ts.axis.spatial(128, ax2_0_0 * 32 + ax0_ax1_fused % 32)
- Ts.reads(A[v0, v1])
- Ts.writes(A_reindex_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "double_buffer_scope": 0, "meta_schedule.cooperative_fetch": 4, "tirx.manifest_shared_memory_local_stage": 1})
- A_reindex_shared[v0, v1] = A[v0, v1]
- for ax0_ax1_fused in range(1024):
- with Ts.sblock("B_reindex_shared"):
- v0 = Ts.axis.spatial(128, ax2_0_0 * 32 + ax0_ax1_fused // 32)
- v1 = Ts.axis.spatial(128, ax0_0_1_ax1_0_1_fused % 4 * 32 + ax0_ax1_fused % 32)
- Ts.reads(B[v0, v1])
- Ts.writes(B_reindex_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "double_buffer_scope": 0, "meta_schedule.cooperative_fetch": 2, "tirx.manifest_shared_memory_local_stage": 1})
- B_reindex_shared[v0, v1] = B[v0, v1]
- for ax2_0_1 in T.serial(2, annotations={"software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, 1]}):
- for ax0_0, ax1_0 in T.grid(2, 1):
- with Ts.sblock("A_reindex_shared_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(8, ax0_0_1_ax1_0_1_fused // 4 * 2 + ax0_0)
- v1_o = Ts.axis.spatial(8, ax2_0_0 * 2 + ax2_0_1 + ax1_0)
- Ts.reads(A_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(A_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": f"wmma_load_16x16x16_f16_a_{intrin_suffix}"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("A_reindex_shared_wmma.matrix_a"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(A_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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(1, 2):
- with Ts.sblock("B_reindex_shared_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(8, ax2_0_0 * 2 + ax2_0_1 + ax0_0)
- v1_o = Ts.axis.spatial(8, ax0_0_1_ax1_0_1_fused % 4 * 2 + ax1_0)
- Ts.reads(B_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(B_reindex_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": f"wmma_load_16x16x16_f16_b_{intrin_suffix}"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("B_reindex_shared_wmma.matrix_b"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(B_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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, 1, 2, 2):
- with Ts.sblock("C_o"):
- v0_o = Ts.axis.spatial(8, ax0_0_1_ax1_0_1_fused // 4 * 2 + ax0_0_3 * 2 + ax0_0_4)
- v1_o = Ts.axis.spatial(8, ax0_0_1_ax1_0_1_fused % 4 * 2 + ax1_0_3 * 2 + ax1_0_4)
- v2_o = Ts.axis.reduce(8, ax2_0_0 * 2 + ax2_0_1 + ax2_0_2)
- Ts.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])
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 2, v0_o % 2, v1_o % 2, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init": "wmma_fill_16x16x16_f32", "warp_execution": 1})
- with Ts.init():
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("C_init"):
- v0_i_init, v1_i_init = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads()
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 2, v0_o % 2, v1_o % 2, v0_i_init, v1_i_init])
- C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 2, v0_o % 2, v1_o % 2, v0_i_init, v1_i_init] = T.float32(0)
- for ax0_1, ax1_1, ax2_1 in T.grid(16, 16, 16):
- with Ts.sblock("C"):
- v0_i, v1_i, v2_i = Ts.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
- Ts.reads(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 2, v0_o % 2, v1_o % 2, 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])
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 2, v0_o % 2, v1_o % 2, v0_i, v1_i])
- Ts.sblock_attr({"meta_schedule.tiling_structure": "SSSRRSRS"})
- C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 2, v0_o % 2, v1_o % 2, v0_i, v1_i] = C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 2, v0_o % 2, v1_o % 2, 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(1, thread="threadIdx.y"):
- for ax2_1, ax3 in T.grid(1, 2):
- with Ts.sblock("C_reindex_shared_wmma.accumulator_o"):
- v0 = Ts.axis.spatial(4, ax0_0_1_ax1_0_1_fused // 4)
- v1 = Ts.axis.spatial(4, ax0_0_1_ax1_0_1_fused % 4)
- v2 = Ts.axis.spatial(2, ax2 + ax2_1)
- v3 = Ts.axis.spatial(2, ax3)
- v4_o = Ts.axis.spatial(1, 0)
- v5_o = Ts.axis.spatial(1, 0)
- Ts.reads(C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, 0:16, 0:16])
- Ts.writes(C_reindex_shared[v0, v1, v2, v3, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": f"wmma_store_16x16x16_f32_{intrin_suffix}"})
- for ax4, ax5 in T.grid(16, 16):
- with Ts.sblock("C_reindex_shared_wmma.accumulator"):
- v4_i, v5_i = Ts.axis.remap("SS", [ax4, ax5])
- Ts.reads(C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, v4_i, v5_i])
- Ts.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(512):
- with Ts.sblock("C_reindex_shared"):
- v0 = Ts.axis.spatial(4, ax0_0_1_ax1_0_1_fused // 4)
- v1 = Ts.axis.spatial(4, ax0_0_1_ax1_0_1_fused % 4)
- v2 = Ts.axis.spatial(2, ax2)
- v3 = Ts.axis.spatial(2, ax0_ax1_ax3_ax4_ax5_fused // 256)
- v4 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 256 // 16)
- v5 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
- Ts.reads(C_reindex_shared[v0, v1, v2, v3, v4, v5])
- Ts.writes(C[v4 + v2 * 16 + v0 * 32, v5 + v3 * 16 + v1 * 32])
- Ts.sblock_attr({"meta_schedule.cooperative_fetch": 3})
- C[v4 + v2 * 16 + v0 * 32, v5 + v3 * 16 + v1 * 32] = C_reindex_shared[v0, v1, v2, v3, v4, v5]
- for i0, i1 in T.grid(128, 128):
- with Ts.sblock("compute"):
- v_i0, v_i1 = Ts.axis.remap("SS", [i0, i1])
- Ts.reads(C[v_i0, v_i1])
- Ts.writes(compute[v_i0, v_i1])
- compute[v_i0, v_i1] = T.max(C[v_i0, v_i1], T.float32(0))
-
- # fmt: on
- decision_0 = [
- ("SamplePerfectTile", [1, 4, 1, 1, 2]),
- ("SamplePerfectTile", [1, 4, 1, 1, 2]),
- ("SamplePerfectTile", [4, 2, 1]),
- ("SampleCategorical", 2),
- ("SampleCategorical", 2),
- ("SampleCategorical", 1),
- ]
- mod = te.create_function(
- 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({"kind": "cuda", "arch": "sm_70"}),
- types=None,
- sch_rules=[
- multi_level_tiling_tensor_core(
- read_reuse_scope=shared_scope,
- write_reuse_scope=shared_scope,
- use_software_pipeline=True,
- ),
- ],
- )
- check_sketches(
- mod,
- sketches=actual,
- expected_mods=[matmul_relu_pipeline_0],
- expected_decisions=[decision_0],
- )
-
-
-def test_matmul_relu_non_tensorizable():
- # expected to do nothing on non-tensorizable workloads
- mod = te.create_function(
- te_workload.matmul_relu( # dtype doesn't match tensor intrin
- n=128,
- m=128,
- k=128,
- )
- )
- (sch,) = generate_design_space(
- kind="cuda",
- mod=mod,
- target=tvm.target.Target({"kind": "cuda", "arch": "sm_70"}),
- types=None,
- sch_rules=[multi_level_tiling_tensor_core(write_reuse_scope="shared")]
- + get_rules("cuda", ms.schedule_rule.AutoInline),
- )
- tvm.ir.assert_structural_equal(mod, sch.mod["main"])
-
-
-def test_padded_matmul_relu():
- # fmt: off
- @Ts.function
- def padded_matmul_relu_0(A: T.Tensor((127, 127), "float16"), B: T.Tensor((127, 127), "float16"), compute: T.Tensor((127, 127), "float32")) -> None:
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- C_reindex_shared = Ts.sblock_alloc_buffer((4, 8, 2, 1, 16, 16), scope="shared")
- C_reindex_shared_wmma_accumulator = Ts.sblock_alloc_buffer((4, 8, 2, 1, 16, 16), scope="wmma.accumulator")
- A_reindex_shared = Ts.sblock_alloc_buffer((128, 128), "float16", scope="shared")
- B_reindex_shared = Ts.sblock_alloc_buffer((128, 128), "float16", scope="shared")
- A_reindex_shared_wmma_matrix_a = Ts.sblock_alloc_buffer((128, 128), "float16", scope="wmma.matrix_a")
- B_reindex_shared_wmma_matrix_b = Ts.sblock_alloc_buffer((128, 128), "float16", scope="wmma.matrix_b")
- for ax0_0_0_ax1_0_0_fused in T.thread_binding(8, 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(1):
- for ax0_ax1_fused in range(4096):
- with Ts.sblock("A_reindex_shared"):
- v0 = Ts.axis.spatial(128, ax0_0_0_ax1_0_0_fused // 2 * 32 + ax0_ax1_fused // 128)
- v1 = Ts.axis.spatial(128, ax0_ax1_fused % 128)
- Ts.reads(A[v0, v1])
- Ts.writes(A_reindex_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 8})
- A_reindex_shared[v0, v1] = T.if_then_else(v0 < 127 and v1 < 127, A[v0, v1], T.float16(0))
- for ax0_ax1_fused in range(4096):
- with Ts.sblock("B_reindex_shared"):
- v0 = Ts.axis.spatial(128, ax0_ax1_fused // 32)
- v1 = Ts.axis.spatial(128, ax0_0_0_ax1_0_0_fused % 2 * 64 + ax0_0_1_ax1_0_1_fused * 32 + ax0_ax1_fused % 32)
- Ts.reads(B[v0, v1])
- Ts.writes(B_reindex_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 1})
- B_reindex_shared[v0, v1] = T.if_then_else(v0 < 127 and v1 < 127, B[v0, v1], T.float16(0))
- for ax2_0_1 in range(4):
- for ax0_0, ax1_0 in T.grid(2, 2):
- with Ts.sblock("A_reindex_shared_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused // 2 * 2 + ax0_0)
- v1_o = Ts.axis.spatial(8, ax2_0_1 * 2 + ax1_0)
- Ts.reads(A_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(A_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_a_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("A_reindex_shared_wmma.matrix_a"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(A_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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(2, 1):
- with Ts.sblock("B_reindex_shared_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(8, ax2_0_1 * 2 + ax0_0)
- v1_o = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused % 2 * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_0_2_ax1_0_2_fused + ax1_0)
- Ts.reads(B_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(B_reindex_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_b_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("B_reindex_shared_wmma.matrix_b"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(B_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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, 2, 2, 1):
- with Ts.sblock("C_o"):
- v0_o = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused // 2 * 2 + ax0_0_3 * 2 + ax0_0_4)
- v1_o = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused % 2 * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_0_2_ax1_0_2_fused + ax1_0_3 + ax1_0_4)
- v2_o = Ts.axis.reduce(8, ax2_0_0 * 8 + ax2_0_1 * 2 + ax2_0_2)
- Ts.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])
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init": "wmma_fill_16x16x16_f32", "warp_execution": 1})
- with Ts.init():
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("C_init"):
- v0_i_init, v1_i_init = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads()
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, v0_i_init, v1_i_init])
- C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, v0_i_init, v1_i_init] = T.float32(0)
- for ax0_1, ax1_1, ax2_1 in T.grid(16, 16, 16):
- with Ts.sblock("C"):
- v0_i, v1_i, v2_i = Ts.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
- Ts.reads(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, 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])
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, v0_i, v1_i])
- Ts.sblock_attr({"meta_schedule.tiling_structure": "SSSRRSRS"})
- C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, v0_i, v1_i] = C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o, v0_o % 2, 0, 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, 1):
- with Ts.sblock("C_reindex_shared_wmma.accumulator_o"):
- v0 = Ts.axis.spatial(4, ax0_0_0_ax1_0_0_fused // 2)
- v1 = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused % 2 * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_ax1_fused)
- v2 = Ts.axis.spatial(2, ax2 + ax2_1)
- v3 = Ts.axis.spatial(1, ax3)
- v4_o = Ts.axis.spatial(1, 0)
- v5_o = Ts.axis.spatial(1, 0)
- Ts.reads(C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, 0:16, 0:16])
- Ts.writes(C_reindex_shared[v0, v1, v2, v3, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_store_16x16x16_f32_shared"})
- for ax4, ax5 in T.grid(16, 16):
- with Ts.sblock("C_reindex_shared_wmma.accumulator"):
- v4_i, v5_i = Ts.axis.remap("SS", [ax4, ax5])
- Ts.reads(C_reindex_shared_wmma_accumulator[v0, v1, v2, v3, v4_i, v5_i])
- Ts.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(512):
- with Ts.sblock("C_reindex_shared"):
- v0 = Ts.axis.spatial(4, ax0_0_0_ax1_0_0_fused // 2)
- v1 = Ts.axis.spatial(8, ax0_0_0_ax1_0_0_fused % 2 * 4 + ax0_0_1_ax1_0_1_fused * 2 + ax0_ax1_ax3_ax4_ax5_fused // 256)
- v2 = Ts.axis.spatial(2, ax2)
- v3 = Ts.axis.spatial(1, 0)
- v4 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 256 // 16)
- v5 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
- Ts.where(ax0_0_0_ax1_0_0_fused // 2 * 32 + ax2 * 16 + ax0_ax1_ax3_ax4_ax5_fused % 256 // 16 < 127 and ax0_0_0_ax1_0_0_fused % 2 * 64 + ax0_0_1_ax1_0_1_fused * 32 + ax0_ax1_ax3_ax4_ax5_fused // 256 * 16 + ax0_ax1_ax3_ax4_ax5_fused % 16 < 127)
- Ts.reads(C_reindex_shared[v0, v1, v2, v3, v4, v5])
- Ts.writes(compute[v4 + v2 * 16 + v0 * 32, v5 + v1 * 16])
- Ts.sblock_attr({"meta_schedule.cooperative_fetch": 4})
- compute[v4 + v2 * 16 + v0 * 32, v5 + v1 * 16] = T.max(C_reindex_shared[v0, v1, v2, v3, v4, v5], T.float32(0))
- # fmt: on
-
- decision_0 = [
- ("SamplePerfectTile", [4, 1, 1, 1, 2]),
- ("SamplePerfectTile", [2, 2, 2, 1, 1]),
- ("SamplePerfectTile", [1, 4, 2]),
- ("SampleCategorical", 3),
- ("SampleCategorical", 3),
- ("SampleCategorical", 0),
- ]
-
- mod = te.create_function(
- te_workload.matmul_relu(
- n=127,
- m=127,
- k=127,
- in_dtype="float16",
- out_dtype="float32",
- )
- )
- actual = generate_design_space(
- kind="cuda",
- mod=mod,
- target=tvm.target.Target({"kind": "cuda", "arch": "sm_70"}),
- types=None,
- sch_rules=[multi_level_tiling_tensor_core(write_reuse_scope="shared")]
- + get_rules("cuda", ms.schedule_rule.AutoInline),
- )
- check_sketches(
- mod,
- sketches=actual,
- expected_mods=[padded_matmul_relu_0],
- expected_decisions=[decision_0],
- )
-
-
-def test_conv_1x1():
- # fmt: off
- @Ts.function
- def conv2d_1x1_0(inputs: T.Tensor((1, 16, 16, 64), "float16"), weight: T.Tensor((1, 1, 64, 64), "float16"), conv2d_nhwc: T.Tensor((1, 16, 16, 64), "float32")):
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- # with Ts.sblock("root"):
- conv2d_nhwc_reindex_shared = Ts.sblock_alloc_buffer((2, 1, 8, 4, 16, 16), scope="shared")
- conv2d_nhwc_reindex_shared_wmma_accumulator = Ts.sblock_alloc_buffer((2, 1, 8, 4, 16, 16), scope="wmma.accumulator")
- PadInput_reindex_shared = Ts.sblock_alloc_buffer((256, 64), "float16", scope="shared")
- weight_reindex_shared = Ts.sblock_alloc_buffer((1, 1, 64, 64), "float16", scope="shared")
- PadInput_reindex_shared_wmma_matrix_a = Ts.sblock_alloc_buffer((256, 64), "float16", scope="wmma.matrix_a")
- weight_reindex_shared_wmma_matrix_b = Ts.sblock_alloc_buffer((1, 1, 64, 64), "float16", scope="wmma.matrix_b")
- for ax0_ax1_ax2_0_0_ax3_0_0_fused in T.thread_binding(1, thread="blockIdx.y"):
- for ax2_0_1_ax3_0_1_fused in T.thread_binding(1, thread="blockIdx.x"):
- for ax2_0_2_ax3_0_2_fused in T.thread_binding(2, thread="threadIdx.y"):
- for ax4_0_0 in range(2):
- for ax0_ax1_fused in range(8192):
- with Ts.sblock("PadInput_reindex_shared"):
- v0 = Ts.axis.spatial(256, ax0_ax1_fused // 32)
- v1 = Ts.axis.spatial(64, ax4_0_0 * 32 + ax0_ax1_fused % 32)
- Ts.reads(inputs[0, v0 // 16, v0 % 16, v1])
- Ts.writes(PadInput_reindex_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 8})
- PadInput_reindex_shared[v0, v1] = inputs[0, v0 // 16, v0 % 16, v1]
- for ax0_ax1_ax2_ax3_fused in range(2048):
- with Ts.sblock("weight_reindex_shared"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial(1, 0)
- v2 = Ts.axis.spatial(64, ax4_0_0 * 32 + ax0_ax1_ax2_ax3_fused // 64)
- v3 = Ts.axis.spatial(64, ax0_ax1_ax2_ax3_fused % 64)
- Ts.reads(weight[v0, v1, v2, v3])
- Ts.writes(weight_reindex_shared[v0, v1, v2, v3])
- Ts.sblock_attr({"buffer_dim_align": [[0, 2, 32, 8]], "meta_schedule.cooperative_fetch": 4})
- weight_reindex_shared[v0, v1, v2, v3] = weight[v0, v1, v2, v3]
- for ax4_0_1 in range(1):
- for ax0_0, ax1_0 in T.grid(8, 2):
- with Ts.sblock("PadInput_reindex_shared_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(16, ax2_0_2_ax3_0_2_fused * 8 + ax0_0)
- v1_o = Ts.axis.spatial(4, ax4_0_0 * 2 + ax1_0)
- Ts.reads(PadInput_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(PadInput_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_a_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("PadInput_reindex_shared_wmma.matrix_a"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(PadInput_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.writes(PadInput_reindex_shared_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- PadInput_reindex_shared_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = PadInput_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i]
- for ax0, ax1, ax2_0, ax3_0 in T.grid(1, 1, 2, 4):
- with Ts.sblock("weight_reindex_shared_wmma.matrix_b_o"):
- v0_o, v1_o = Ts.axis.remap("SS", [ax0, ax1])
- v2_o = Ts.axis.spatial(4, ax4_0_0 * 2 + ax2_0)
- v3_o = Ts.axis.spatial(4, ax3_0)
- Ts.reads(weight_reindex_shared[v0_o, v1_o, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(weight_reindex_shared_wmma_matrix_b[v0_o, v1_o, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_b_shared"})
- for ax2_1, ax3_1 in T.grid(16, 16):
- with Ts.sblock("weight_reindex_shared_wmma.matrix_b"):
- v2_i, v3_i = Ts.axis.remap("SS", [ax2_1, ax3_1])
- Ts.reads(weight_reindex_shared[v0_o, v1_o, v2_o * 16 + v2_i, v3_o * 16 + v3_i])
- Ts.writes(weight_reindex_shared_wmma_matrix_b[v0_o, v1_o, v2_o * 16 + v2_i, v3_o * 16 + v3_i])
- weight_reindex_shared_wmma_matrix_b[v0_o, v1_o, v2_o * 16 + v2_i, v3_o * 16 + v3_i] = weight_reindex_shared[v0_o, v1_o, v2_o * 16 + v2_i, v3_o * 16 + v3_i]
- for ax2_0_3, ax3_0_3, ax4_0_2, ax2_0_4, ax3_0_4 in T.grid(8, 1, 2, 1, 4):
- with Ts.sblock("conv2d_nhwc_o"):
- v0_o = Ts.axis.spatial(1, 0)
- v1_o = Ts.axis.spatial(1, 0)
- v2_o = Ts.axis.spatial(16, ax2_0_2_ax3_0_2_fused * 8 + ax2_0_3 + ax2_0_4)
- v3_o = Ts.axis.spatial(4, ax3_0_3 * 4 + ax3_0_4)
- v4_o = Ts.axis.reduce(4, ax4_0_0 * 2 + ax4_0_1 * 2 + ax4_0_2)
- Ts.reads(PadInput_reindex_shared_wmma_matrix_a[v2_o * 16:v2_o * 16 + 16, v4_o * 16:v4_o * 16 + 16], weight_reindex_shared_wmma_matrix_b[v0_o, v1_o, v4_o * 16:v4_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o // 8, 0, v2_o % 8, v3_o, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init": "wmma_fill_16x16x16_f32", "warp_execution": 1})
- with Ts.init():
- for ax2_1, ax3_1 in T.grid(16, 16):
- with Ts.sblock("conv2d_nhwc_init"):
- v2_i_init, v3_i_init = Ts.axis.remap("SS", [ax2_1, ax3_1])
- Ts.reads()
- Ts.writes(conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o // 8, 0, v2_o % 8, v3_o, v2_i_init, v3_i_init])
- conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o // 8, 0, v2_o % 8, v3_o, v2_i_init, v3_i_init] = T.float32(0)
- for ax2_1, ax3_1, ax4_1 in T.grid(16, 16, 16):
- with Ts.sblock("conv2d_nhwc"):
- v2_i, v3_i, v4_i = Ts.axis.remap("SSR", [ax2_1, ax3_1, ax4_1])
- Ts.reads(conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o // 8, 0, v2_o % 8, v3_o, v2_i, v3_i], PadInput_reindex_shared_wmma_matrix_a[v2_o * 16 + v2_i, v4_o * 16 + v4_i], weight_reindex_shared_wmma_matrix_b[v0_o, v1_o, v4_o * 16 + v4_i, v3_o * 16 + v3_i])
- Ts.writes(conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o // 8, 0, v2_o % 8, v3_o, v2_i, v3_i])
- Ts.sblock_attr({"meta_schedule.tiling_structure": "SSSRRSRS"})
- conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o // 8, 0, v2_o % 8, v3_o, v2_i, v3_i] = conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o // 8, 0, v2_o % 8, v3_o, v2_i, v3_i] + T.Cast("float32", PadInput_reindex_shared_wmma_matrix_a[v2_o * 16 + v2_i, v4_o * 16 + v4_i]) * T.Cast("float32", weight_reindex_shared_wmma_matrix_b[v0_o, v1_o, v4_o * 16 + v4_i, v3_o * 16 + v3_i])
- for ax2 in range(8):
- for ax0_ax1_fused in T.thread_binding(2, thread="threadIdx.y"):
- for ax2_1, ax3 in T.grid(1, 4):
- with Ts.sblock("conv2d_nhwc_reindex_shared_wmma.accumulator_o"):
- v0_o = Ts.axis.spatial(2, ax0_ax1_fused)
- v1_o = Ts.axis.spatial(1, 0)
- v2_o = Ts.axis.spatial(8, ax2 + ax2_1)
- v3_o = Ts.axis.spatial(4, ax3)
- v4_o = Ts.axis.spatial(1, 0)
- v5_o = Ts.axis.spatial(1, 0)
- Ts.reads(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, 0:16, 0:16])
- Ts.writes(conv2d_nhwc_reindex_shared[v0_o, v1_o, v2_o, v3_o, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_store_16x16x16_f32_shared"})
- for ax4, ax5 in T.grid(16, 16):
- with Ts.sblock("conv2d_nhwc_reindex_shared_wmma.accumulator"):
- v4_i, v5_i = Ts.axis.remap("SS", [ax4, ax5])
- Ts.reads(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i])
- Ts.writes(conv2d_nhwc_reindex_shared[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i])
- conv2d_nhwc_reindex_shared[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i] = conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i]
- for ax0_ax1_ax3_ax4_ax5_fused in range(2048):
- with Ts.sblock("conv2d_nhwc_reindex_shared"):
- v0 = Ts.axis.spatial(2, ax0_ax1_ax3_ax4_ax5_fused // 1024)
- v1 = Ts.axis.spatial(1, 0)
- v2 = Ts.axis.spatial(8, ax2)
- v3 = Ts.axis.spatial(4, ax0_ax1_ax3_ax4_ax5_fused % 1024 // 256)
- v4 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 256 // 16)
- v5 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
- Ts.reads(conv2d_nhwc_reindex_shared[v0, v1, v2, v3, v4, v5])
- Ts.writes(conv2d_nhwc[0, (v4 + v2 * 16 + v0 * 128) // 16, (v4 + v2 * 16 + v0 * 128) % 16, v5 + v3 * 16])
- Ts.sblock_attr({"meta_schedule.cooperative_fetch": 1})
- conv2d_nhwc[0, (v4 + v2 * 16 + v0 * 128) // 16, (v4 + v2 * 16 + v0 * 128) % 16, v5 + v3 * 16] = conv2d_nhwc_reindex_shared[v0, v1, v2, v3, v4, v5]
- # fmt: on
-
- decision_0 = [
- ("SamplePerfectTile", [1, 1, 2, 8, 1]),
- ("SamplePerfectTile", [1, 1, 1, 1, 4]),
- ("SamplePerfectTile", [2, 1, 2]),
- ("SampleCategorical", 0),
- ("SampleCategorical", 3),
- ("SampleCategorical", 2),
- ]
-
- mod = te.create_function(
- te_workload.conv2d_nhwc(
- 1,
- 16,
- 16,
- 64,
- 64,
- 1,
- 1,
- 0,
- in_dtype="float16",
- out_dtype="float32",
- )
- )
- actual = generate_design_space(
- kind="cuda",
- mod=mod,
- target=tvm.target.Target({"kind": "cuda", "arch": "sm_70"}),
- types=None,
- sch_rules=[multi_level_tiling_tensor_core(write_reuse_scope="shared")]
- + get_rules("cuda", ms.schedule_rule.AutoInline),
- )
- check_sketches(
- mod,
- sketches=actual,
- expected_mods=[conv2d_1x1_0],
- expected_decisions=[decision_0],
- )
-
-
-def test_padded_conv():
- # fmt: off
- @Ts.function
- def padded_conv2d_0(inputs: T.Tensor((1, 224, 224, 3), "float16"), weight: T.Tensor((7, 7, 3, 64), "float16"), conv2d_nhwc: T.Tensor((1, 112, 112, 64), "float32")):
- T.func_attr({"tirx.noalias": True})
- # with Ts.sblock("root"):
- conv2d_nhwc_reindex_shared = Ts.sblock_alloc_buffer((56, 2, 14, 2, 16, 16), scope="shared")
- conv2d_nhwc_reindex_shared_wmma_accumulator = Ts.sblock_alloc_buffer((56, 2, 14, 2, 16, 16), scope="wmma.accumulator")
- PadInput_reindex_pad_shared = Ts.sblock_alloc_buffer((12544, 160), "float16", scope="shared")
- weight_reindex_pad_shared = Ts.sblock_alloc_buffer((160, 64), "float16", scope="shared")
- PadInput_reindex_pad_shared_wmma_matrix_a = Ts.sblock_alloc_buffer((12544, 160), "float16", scope="wmma.matrix_a")
- weight_reindex_pad_shared_wmma_matrix_b = Ts.sblock_alloc_buffer((160, 64), "float16", scope="wmma.matrix_b")
- for ax0_0_0_ax1_0_0_fused in T.thread_binding(14, thread="blockIdx.y"):
- for ax0_0_1_ax1_0_1_fused in T.thread_binding(1, thread="blockIdx.x"):
- for ax0_0_2_ax1_0_2_fused in T.thread_binding(8, thread="threadIdx.y"):
- for ax2_0_0 in range(10):
- for ax0_ax1_fused in range(28672):
- with Ts.sblock("PadInput_reindex_pad_shared"):
- v0 = Ts.axis.spatial(12544, ax0_0_0_ax1_0_0_fused // 2 * 1792 + ax0_ax1_fused // 16)
- v1 = Ts.axis.spatial(160, ax2_0_0 * 16 + ax0_ax1_fused % 16)
- Ts.reads(inputs[0, v0 // 112 * 2 + v1 // 21 - 3, v0 % 112 * 2 + v1 % 21 // 3 - 3, v1 % 3])
- Ts.writes(PadInput_reindex_pad_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 4})
- PadInput_reindex_pad_shared[v0, v1] = T.if_then_else(v1 < 147, T.if_then_else(3 <= v0 // 112 * 2 + v1 // 21 and v0 // 112 * 2 + v1 // 21 < 227 and 3 <= v0 % 112 * 2 + v1 % 21 // 3 and v0 % 112 * 2 + v1 % 21 // 3 < 227, inputs[0, v0 // 112 * 2 + v1 // 21 - 3, v0 % 112 * 2 + v1 % 21 // 3 - 3, v1 % 3], T.float16(0)), T.float16(0))
- for ax0_ax1_fused in range(512):
- with Ts.sblock("weight_reindex_pad_shared"):
- v0 = Ts.axis.spatial(160, ax2_0_0 * 16 + ax0_ax1_fused // 32)
- v1 = Ts.axis.spatial(64, ax0_0_0_ax1_0_0_fused % 2 * 32 + ax0_ax1_fused % 32)
- Ts.reads(weight[v0 // 21, v0 % 21 // 3, v0 % 3, v1])
- Ts.writes(weight_reindex_pad_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 2})
- weight_reindex_pad_shared[v0, v1] = T.if_then_else(v0 < 147, weight[v0 // 21, v0 % 21 // 3, v0 % 3, v1], T.float16(0))
- for ax2_0_1 in range(1):
- for ax0_0, ax1_0 in T.grid(14, 1):
- with Ts.sblock("PadInput_reindex_pad_shared_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(784, ax0_0_0_ax1_0_0_fused // 2 * 112 + ax0_0_2_ax1_0_2_fused * 14 + ax0_0)
- v1_o = Ts.axis.spatial(10, ax2_0_0 + ax1_0)
- Ts.reads(PadInput_reindex_pad_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(PadInput_reindex_pad_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_a_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("PadInput_reindex_pad_shared_wmma.matrix_a"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(PadInput_reindex_pad_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.writes(PadInput_reindex_pad_shared_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- PadInput_reindex_pad_shared_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = PadInput_reindex_pad_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i]
- for ax0_0, ax1_0 in T.grid(1, 2):
- with Ts.sblock("weight_reindex_pad_shared_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(10, ax2_0_0 + ax0_0)
- v1_o = Ts.axis.spatial(4, ax0_0_0_ax1_0_0_fused % 2 * 2 + ax1_0)
- Ts.reads(weight_reindex_pad_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(weight_reindex_pad_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_b_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("weight_reindex_pad_shared_wmma.matrix_b"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(weight_reindex_pad_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.writes(weight_reindex_pad_shared_wmma_matrix_b[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- weight_reindex_pad_shared_wmma_matrix_b[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = weight_reindex_pad_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(7, 2, 1, 2, 1):
- with Ts.sblock("conv2d_nhwc_o"):
- v0_o = Ts.axis.spatial(784, ax0_0_0_ax1_0_0_fused // 2 * 112 + ax0_0_2_ax1_0_2_fused * 14 + ax0_0_3 * 2 + ax0_0_4)
- v1_o = Ts.axis.spatial(4, ax0_0_0_ax1_0_0_fused % 2 * 2 + ax1_0_3 + ax1_0_4)
- v2_o = Ts.axis.reduce(10, ax2_0_0 + ax2_0_1 + ax2_0_2)
- Ts.reads(PadInput_reindex_pad_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], weight_reindex_pad_shared_wmma_matrix_b[v2_o * 16:v2_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o // 14, v1_o // 2, v0_o % 14, v1_o % 2, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init": "wmma_fill_16x16x16_f32", "warp_execution": 1})
- with Ts.init():
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("conv2d_nhwc_init"):
- v0_i_init, v1_i_init = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads()
- Ts.writes(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o // 14, v1_o // 2, v0_o % 14, v1_o % 2, v0_i_init, v1_i_init])
- conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o // 14, v1_o // 2, v0_o % 14, v1_o % 2, v0_i_init, v1_i_init] = T.float32(0)
- for ax0_1, ax1_1, ax2_1 in T.grid(16, 16, 16):
- with Ts.sblock("conv2d_nhwc"):
- v0_i, v1_i, v2_i = Ts.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
- Ts.reads(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o // 14, v1_o // 2, v0_o % 14, v1_o % 2, v0_i, v1_i], PadInput_reindex_pad_shared_wmma_matrix_a[v0_o * 16 + v0_i, v2_o * 16 + v2_i], weight_reindex_pad_shared_wmma_matrix_b[v2_o * 16 + v2_i, v1_o * 16 + v1_i])
- Ts.writes(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o // 14, v1_o // 2, v0_o % 14, v1_o % 2, v0_i, v1_i])
- Ts.sblock_attr({"meta_schedule.tiling_structure": "SSSRRSRS"})
- conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o // 14, v1_o // 2, v0_o % 14, v1_o % 2, v0_i, v1_i] = conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o // 14, v1_o // 2, v0_o % 14, v1_o % 2, v0_i, v1_i] + T.Cast("float32", PadInput_reindex_pad_shared_wmma_matrix_a[v0_o * 16 + v0_i, v2_o * 16 + v2_i]) * T.Cast("float32", weight_reindex_pad_shared_wmma_matrix_b[v2_o * 16 + v2_i, v1_o * 16 + v1_i])
- for ax2 in range(14):
- for ax0_ax1_fused in T.thread_binding(8, thread="threadIdx.y"):
- for ax2_1, ax3 in T.grid(1, 2):
- with Ts.sblock("conv2d_nhwc_reindex_shared_wmma.accumulator_o"):
- v0_o = Ts.axis.spatial(56, ax0_0_0_ax1_0_0_fused // 2 * 8 + ax0_ax1_fused)
- v1_o = Ts.axis.spatial(2, ax0_0_0_ax1_0_0_fused % 2)
- v2_o = Ts.axis.spatial(14, ax2 + ax2_1)
- v3_o = Ts.axis.spatial(2, ax3)
- v4_o = Ts.axis.spatial(1, 0)
- v5_o = Ts.axis.spatial(1, 0)
- Ts.reads(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, 0:16, 0:16])
- Ts.writes(conv2d_nhwc_reindex_shared[v0_o, v1_o, v2_o, v3_o, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_store_16x16x16_f32_shared"})
- for ax4, ax5 in T.grid(16, 16):
- with Ts.sblock("conv2d_nhwc_reindex_shared_wmma.accumulator"):
- v4_i, v5_i = Ts.axis.remap("SS", [ax4, ax5])
- Ts.reads(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i])
- Ts.writes(conv2d_nhwc_reindex_shared[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i])
- conv2d_nhwc_reindex_shared[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i] = conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i]
- for ax0_ax1_ax3_ax4_ax5_fused in range(4096):
- with Ts.sblock("conv2d_nhwc_reindex_shared"):
- v0 = Ts.axis.spatial(56, ax0_0_0_ax1_0_0_fused // 2 * 8 + ax0_ax1_ax3_ax4_ax5_fused // 512)
- v1 = Ts.axis.spatial(2, ax0_0_0_ax1_0_0_fused % 2)
- v2 = Ts.axis.spatial(14, ax2)
- v3 = Ts.axis.spatial(2, ax0_ax1_ax3_ax4_ax5_fused % 512 // 256)
- v4 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 256 // 16)
- v5 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
- Ts.reads(conv2d_nhwc_reindex_shared[v0, v1, v2, v3, v4, v5])
- Ts.writes(conv2d_nhwc[0, (v4 + v2 * 16 + v0 * 224) // 112, (v4 + v2 * 16 + v0 * 224) % 112, v5 + v3 * 16 + v1 * 32])
- Ts.sblock_attr({"meta_schedule.cooperative_fetch": 3})
- conv2d_nhwc[0, (v4 + v2 * 16 + v0 * 224) // 112, (v4 + v2 * 16 + v0 * 224) % 112, v5 + v3 * 16 + v1 * 32] = conv2d_nhwc_reindex_shared[v0, v1, v2, v3, v4, v5]
- # fmt: on
-
- decision_0 = [
- ("SamplePerfectTile", [7, 1, 8, 7, 2]),
- ("SamplePerfectTile", [2, 1, 1, 2, 1]),
- ("SamplePerfectTile", [10, 1, 1]),
- ("SampleCategorical", 2),
- ("SampleCategorical", 2),
- ("SampleCategorical", 1),
- ]
- mod = te.create_function(
- te_workload.conv2d_nhwc(
- 1,
- 224,
- 224,
- 3,
- 64,
- 7,
- 2,
- 3,
- in_dtype="float16",
- out_dtype="float32",
- )
- )
- actual = generate_design_space(
- kind="cuda",
- mod=mod,
- target=tvm.target.Target({"kind": "cuda", "arch": "sm_70"}),
- types=None,
- sch_rules=[multi_level_tiling_tensor_core(write_reuse_scope="shared")]
- + get_rules("cuda", ms.schedule_rule.AutoInline),
- )
- check_sketches(
- mod,
- sketches=actual,
- expected_mods=[padded_conv2d_0],
- expected_decisions=[decision_0],
- )
-
-
-def test_padded_matmul_single_padded_input():
- # fmt: off
- @Ts.function
- def padded_matmul_single_padded_input_0(A: T.Tensor((1023, 4096), "float16"), B: T.Tensor((4096, 1024), "float16"), C: T.Tensor((1023, 1024), "float32")):
- T.func_attr({"tirx.noalias": True})
- # with Ts.sblock("root"):
- C_reindex_pad_shared = Ts.sblock_alloc_buffer((8, 32, 8, 2, 16, 16), scope="shared")
- C_reindex_pad_shared_wmma_accumulator = Ts.sblock_alloc_buffer((8, 32, 8, 2, 16, 16), scope="wmma.accumulator")
- A_reindex_pad_shared = Ts.sblock_alloc_buffer((1024, 4096), "float16", scope="shared")
- B_reindex_shared = Ts.sblock_alloc_buffer((4096, 1024), "float16", scope="shared")
- A_reindex_pad_shared_wmma_matrix_a = Ts.sblock_alloc_buffer((1024, 4096), "float16", scope="wmma.matrix_a")
- B_reindex_shared_wmma_matrix_b = Ts.sblock_alloc_buffer((4096, 1024), "float16", scope="wmma.matrix_b")
- for ax0_0_0_ax1_0_0_fused in T.thread_binding(1, thread="blockIdx.y"):
- for ax0_0_1_ax1_0_1_fused in T.thread_binding(32, thread="blockIdx.x"):
- for ax0_0_2_ax1_0_2_fused in T.thread_binding(8, thread="threadIdx.y"):
- for ax2_0_0 in range(32):
- for ax0_ax1_fused in range(65536):
- with Ts.sblock("A_reindex_pad_shared"):
- v0 = Ts.axis.spatial(1024, ax0_0_1_ax1_0_1_fused // 16 * 512 + ax0_ax1_fused // 128)
- v1 = Ts.axis.spatial(4096, ax2_0_0 * 128 + ax0_ax1_fused % 128)
- Ts.reads(A[v0, v1])
- Ts.writes(A_reindex_pad_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 2})
- A_reindex_pad_shared[v0, v1] = T.if_then_else(v0 < 1023, A[v0, v1], T.float16(0.0))
- for ax0_ax1_fused in range(8192):
- with Ts.sblock("B_reindex_shared"):
- v0 = Ts.axis.spatial(4096, ax2_0_0 * 128 + ax0_ax1_fused // 64)
- v1 = Ts.axis.spatial(1024, ax0_0_1_ax1_0_1_fused % 16 * 64 + ax0_ax1_fused % 64)
- Ts.reads(B[v0, v1])
- Ts.writes(B_reindex_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 1})
- B_reindex_shared[v0, v1] = B[v0, v1]
- for ax2_0_1 in range(8):
- for ax0_0, ax1_0 in T.grid(8, 1):
- with Ts.sblock("A_reindex_pad_shared_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(64, ax0_0_1_ax1_0_1_fused // 16 * 32 + ax0_0_2_ax1_0_2_fused // 2 * 8 + ax0_0)
- v1_o = Ts.axis.spatial(256, ax2_0_0 * 8 + ax2_0_1 + ax1_0)
- Ts.reads(A_reindex_pad_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(A_reindex_pad_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_a_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("A_reindex_pad_shared_wmma.matrix_a"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(A_reindex_pad_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.writes(A_reindex_pad_shared_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- A_reindex_pad_shared_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = A_reindex_pad_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i]
- for ax0_0, ax1_0 in T.grid(1, 2):
- with Ts.sblock("B_reindex_shared_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(256, ax2_0_0 * 8 + ax2_0_1 + ax0_0)
- v1_o = Ts.axis.spatial(64, ax0_0_1_ax1_0_1_fused % 16 * 4 + ax0_0_2_ax1_0_2_fused % 2 * 2 + ax1_0)
- Ts.reads(B_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(B_reindex_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_b_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("B_reindex_shared_wmma.matrix_b"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(B_reindex_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.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(2, 1, 1, 4, 2):
- with Ts.sblock("C_o"):
- v0_o = Ts.axis.spatial(64, ax0_0_1_ax1_0_1_fused // 16 * 32 + ax0_0_2_ax1_0_2_fused // 2 * 8 + ax0_0_3 * 4 + ax0_0_4)
- v1_o = Ts.axis.spatial(64, ax0_0_1_ax1_0_1_fused % 16 * 4 + ax0_0_2_ax1_0_2_fused % 2 * 2 + ax1_0_3 * 2 + ax1_0_4)
- v2_o = Ts.axis.reduce(256, ax2_0_0 * 8 + ax2_0_1 + ax2_0_2)
- Ts.reads(A_reindex_pad_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])
- Ts.writes(C_reindex_pad_shared_wmma_accumulator[v0_o // 8, v1_o // 2, v0_o % 8, v1_o % 2, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init": "wmma_fill_16x16x16_f32", "warp_execution": 1})
- with Ts.init():
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("C_init"):
- v0_i_init, v1_i_init = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads()
- Ts.writes(C_reindex_pad_shared_wmma_accumulator[v0_o // 8, v1_o // 2, v0_o % 8, v1_o % 2, v0_i_init, v1_i_init])
- C_reindex_pad_shared_wmma_accumulator[v0_o // 8, v1_o // 2, v0_o % 8, v1_o % 2, v0_i_init, v1_i_init] = T.float32(0.0)
- for ax0_1, ax1_1, ax2_1 in T.grid(16, 16, 16):
- with Ts.sblock("C"):
- v0_i, v1_i, v2_i = Ts.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
- Ts.reads(C_reindex_pad_shared_wmma_accumulator[v0_o // 8, v1_o // 2, v0_o % 8, v1_o % 2, v0_i, v1_i], A_reindex_pad_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])
- Ts.writes(C_reindex_pad_shared_wmma_accumulator[v0_o // 8, v1_o // 2, v0_o % 8, v1_o % 2, v0_i, v1_i])
- Ts.sblock_attr({"meta_schedule.tiling_structure": "SSSRRSRS"})
- C_reindex_pad_shared_wmma_accumulator[v0_o // 8, v1_o // 2, v0_o % 8, v1_o % 2, v0_i, v1_i] = C_reindex_pad_shared_wmma_accumulator[v0_o // 8, v1_o // 2, v0_o % 8, v1_o % 2, v0_i, v1_i] + T.Cast("float32", A_reindex_pad_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(8):
- for ax0_ax1_fused in T.thread_binding(8, thread="threadIdx.y"):
- for ax2_1, ax3 in T.grid(1, 2):
- with Ts.sblock("C_reindex_pad_shared_wmma.accumulator_o"):
- v0_o = Ts.axis.spatial(8, ax0_0_1_ax1_0_1_fused // 16 * 4 + ax0_ax1_fused // 2)
- v1_o = Ts.axis.spatial(32, ax0_0_1_ax1_0_1_fused % 16 * 2 + ax0_ax1_fused % 2)
- v2_o = Ts.axis.spatial(8, ax2 + ax2_1)
- v3_o = Ts.axis.spatial(2, ax3)
- v4_o = Ts.axis.spatial(1, 0)
- v5_o = Ts.axis.spatial(1, 0)
- Ts.reads(C_reindex_pad_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, 0:16, 0:16])
- Ts.writes(C_reindex_pad_shared[v0_o, v1_o, v2_o, v3_o, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_store_16x16x16_f32_shared"})
- for ax4, ax5 in T.grid(16, 16):
- with Ts.sblock("C_reindex_pad_shared_wmma.accumulator"):
- v4_i, v5_i = Ts.axis.remap("SS", [ax4, ax5])
- Ts.reads(C_reindex_pad_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i])
- Ts.writes(C_reindex_pad_shared[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i])
- C_reindex_pad_shared[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i] = C_reindex_pad_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i]
- for ax0_ax1_ax3_ax4_ax5_fused in range(4096):
- with Ts.sblock("C_reindex_pad_shared"):
- v0 = Ts.axis.spatial(8, ax0_0_1_ax1_0_1_fused // 16 * 4 + ax0_ax1_ax3_ax4_ax5_fused // 1024)
- v1 = Ts.axis.spatial(32, ax0_0_1_ax1_0_1_fused % 16 * 2 + ax0_ax1_ax3_ax4_ax5_fused % 1024 // 512)
- v2 = Ts.axis.spatial(8, ax2)
- v3 = Ts.axis.spatial(2, ax0_ax1_ax3_ax4_ax5_fused % 512 // 256)
- v4 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 256 // 16)
- v5 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
- Ts.where(ax0_0_1_ax1_0_1_fused // 16 * 512 + ax0_ax1_ax3_ax4_ax5_fused // 1024 * 128 + ax2 * 16 + ax0_ax1_ax3_ax4_ax5_fused % 256 // 16 < 1023)
- Ts.reads(C_reindex_pad_shared[v0, v1, v2, v3, v4, v5])
- Ts.writes(C[v4 + v2 * 16 + v0 * 128, v5 + v3 * 16 + v1 * 32])
- Ts.sblock_attr({"meta_schedule.cooperative_fetch": 4})
- C[v4 + v2 * 16 + v0 * 128, v5 + v3 * 16 + v1 * 32] = C_reindex_pad_shared[v0, v1, v2, v3, v4, v5]
- # fmt: on
-
- decision_0 = [
- ("SamplePerfectTile", [1, 2, 4, 2, 4]),
- ("SamplePerfectTile", [1, 16, 2, 1, 2]),
- ("SamplePerfectTile", [32, 8, 1]),
- ("SampleCategorical", 3),
- ("SampleCategorical", 1),
- ("SampleCategorical", 0),
- ]
- mod = te.create_function(
- te_workload.matmul(
- n=1023,
- m=1024,
- k=4096,
- in_dtype="float16",
- out_dtype="float32",
- )
- )
- actual = generate_design_space(
- kind="cuda",
- mod=mod,
- target=tvm.target.Target({"kind": "cuda", "arch": "sm_70"}),
- types=None,
- sch_rules=[multi_level_tiling_tensor_core()]
- + get_rules("cuda", ms.schedule_rule.AutoInline),
- )
- check_sketches(
- mod,
- sketches=actual,
- expected_mods=[padded_matmul_single_padded_input_0],
- expected_decisions=[decision_0],
- )
-
-
-def test_padded_matmul_no_padded_output():
- # fmt: off
- @Ts.function
- def padded_matmul_no_padded_output_0(A: T.Tensor((1024, 4095), "float16"), B: T.Tensor((4095, 1024), "float16"), C: T.Tensor((1024, 1024), "float32")):
- T.func_attr({"tirx.noalias": True})
- # with Ts.sblock("root"):
- C_reindex_shared = Ts.sblock_alloc_buffer((32, 16, 2, 4, 16, 16), scope="shared")
- C_reindex_shared_wmma_accumulator = Ts.sblock_alloc_buffer((32, 16, 2, 4, 16, 16), scope="wmma.accumulator")
- A_reindex_pad_shared = Ts.sblock_alloc_buffer((1024, 4096), "float16", scope="shared")
- B_reindex_pad_shared = Ts.sblock_alloc_buffer((4096, 1024), "float16", scope="shared")
- A_reindex_pad_shared_wmma_matrix_a = Ts.sblock_alloc_buffer((1024, 4096), "float16", scope="wmma.matrix_a")
- B_reindex_pad_shared_wmma_matrix_b = Ts.sblock_alloc_buffer((4096, 1024), "float16", scope="wmma.matrix_b")
- for ax0_0_0_ax1_0_0_fused in T.thread_binding(64, 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(4, thread="threadIdx.y"):
- for ax2_0_0 in range(128):
- for ax0_ax1_fused in range(4096):
- with Ts.sblock("A_reindex_pad_shared"):
- v0 = Ts.axis.spatial(1024, ax0_0_0_ax1_0_0_fused // 16 * 256 + ax0_0_1_ax1_0_1_fused * 128 + ax0_ax1_fused // 32)
- v1 = Ts.axis.spatial(4096, ax2_0_0 * 32 + ax0_ax1_fused % 32)
- Ts.reads(A[v0, v1])
- Ts.writes(A_reindex_pad_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 8})
- A_reindex_pad_shared[v0, v1] = T.if_then_else(v1 < 4095, A[v0, v1], T.float16(0.0))
- for ax0_ax1_fused in range(2048):
- with Ts.sblock("B_reindex_pad_shared"):
- v0 = Ts.axis.spatial(4096, ax2_0_0 * 32 + ax0_ax1_fused // 64)
- v1 = Ts.axis.spatial(1024, ax0_0_0_ax1_0_0_fused % 16 * 64 + ax0_ax1_fused % 64)
- Ts.reads(B[v0, v1])
- Ts.writes(B_reindex_pad_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]], "meta_schedule.cooperative_fetch": 1})
- B_reindex_pad_shared[v0, v1] = T.if_then_else(v0 < 4095, B[v0, v1], T.float16(0.0))
- for ax2_0_1 in range(2):
- for ax0_0, ax1_0 in T.grid(2, 1):
- with Ts.sblock("A_reindex_pad_shared_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(64, ax0_0_0_ax1_0_0_fused // 16 * 16 + ax0_0_1_ax1_0_1_fused * 8 + ax0_0_2_ax1_0_2_fused * 2 + ax0_0)
- v1_o = Ts.axis.spatial(256, ax2_0_0 * 2 + ax2_0_1 + ax1_0)
- Ts.reads(A_reindex_pad_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(A_reindex_pad_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_a_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("A_reindex_pad_shared_wmma.matrix_a"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(A_reindex_pad_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.writes(A_reindex_pad_shared_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- A_reindex_pad_shared_wmma_matrix_a[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = A_reindex_pad_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i]
- for ax0_0, ax1_0 in T.grid(1, 4):
- with Ts.sblock("B_reindex_pad_shared_wmma.matrix_b_o"):
- v0_o = Ts.axis.spatial(256, ax2_0_0 * 2 + ax2_0_1 + ax0_0)
- v1_o = Ts.axis.spatial(64, ax0_0_0_ax1_0_0_fused % 16 * 4 + ax1_0)
- Ts.reads(B_reindex_pad_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(B_reindex_pad_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_load_16x16x16_f16_b_shared"})
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("B_reindex_pad_shared_wmma.matrix_b"):
- v0_i, v1_i = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads(B_reindex_pad_shared[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- Ts.writes(B_reindex_pad_shared_wmma_matrix_b[v0_o * 16 + v0_i, v1_o * 16 + v1_i])
- B_reindex_pad_shared_wmma_matrix_b[v0_o * 16 + v0_i, v1_o * 16 + v1_i] = B_reindex_pad_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(2, 1, 1, 1, 4):
- with Ts.sblock("C_o"):
- v0_o = Ts.axis.spatial(64, ax0_0_0_ax1_0_0_fused // 16 * 16 + ax0_0_1_ax1_0_1_fused * 8 + ax0_0_2_ax1_0_2_fused * 2 + ax0_0_3 + ax0_0_4)
- v1_o = Ts.axis.spatial(64, ax0_0_0_ax1_0_0_fused % 16 * 4 + ax1_0_3 * 4 + ax1_0_4)
- v2_o = Ts.axis.reduce(256, ax2_0_0 * 2 + ax2_0_1 + ax2_0_2)
- Ts.reads(A_reindex_pad_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], B_reindex_pad_shared_wmma_matrix_b[v2_o * 16:v2_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_sync_16x16x16_f16f16f32", "meta_schedule.auto_tensorize_init": "wmma_fill_16x16x16_f32", "warp_execution": 1})
- with Ts.init():
- for ax0_1, ax1_1 in T.grid(16, 16):
- with Ts.sblock("C_init"):
- v0_i_init, v1_i_init = Ts.axis.remap("SS", [ax0_1, ax1_1])
- Ts.reads()
- Ts.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.0)
- for ax0_1, ax1_1, ax2_1 in T.grid(16, 16, 16):
- with Ts.sblock("C"):
- v0_i, v1_i, v2_i = Ts.axis.remap("SSR", [ax0_1, ax1_1, ax2_1])
- Ts.reads(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, v0_i, v1_i], A_reindex_pad_shared_wmma_matrix_a[v0_o * 16 + v0_i, v2_o * 16 + v2_i], B_reindex_pad_shared_wmma_matrix_b[v2_o * 16 + v2_i, v1_o * 16 + v1_i])
- Ts.writes(C_reindex_shared_wmma_accumulator[v0_o // 2, v1_o // 4, v0_o % 2, v1_o % 4, v0_i, v1_i])
- Ts.sblock_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_pad_shared_wmma_matrix_a[v0_o * 16 + v0_i, v2_o * 16 + v2_i]) * T.Cast("float32", B_reindex_pad_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(4, thread="threadIdx.y"):
- for ax2_1, ax3 in T.grid(1, 4):
- with Ts.sblock("C_reindex_shared_wmma.accumulator_o"):
- v0_o = Ts.axis.spatial(32, ax0_0_0_ax1_0_0_fused // 16 * 8 + ax0_0_1_ax1_0_1_fused * 4 + ax0_ax1_fused)
- v1_o = Ts.axis.spatial(16, ax0_0_0_ax1_0_0_fused % 16)
- v2_o = Ts.axis.spatial(2, ax2 + ax2_1)
- v3_o = Ts.axis.spatial(4, ax3)
- v4_o = Ts.axis.spatial(1, 0)
- v5_o = Ts.axis.spatial(1, 0)
- Ts.reads(C_reindex_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, 0:16, 0:16])
- Ts.writes(C_reindex_shared[v0_o, v1_o, v2_o, v3_o, 0:16, 0:16])
- Ts.sblock_attr({"meta_schedule.auto_tensorize": "wmma_store_16x16x16_f32_shared"})
- for ax4, ax5 in T.grid(16, 16):
- with Ts.sblock("C_reindex_shared_wmma.accumulator"):
- v4_i, v5_i = Ts.axis.remap("SS", [ax4, ax5])
- Ts.reads(C_reindex_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i])
- Ts.writes(C_reindex_shared[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i])
- C_reindex_shared[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i] = C_reindex_shared_wmma_accumulator[v0_o, v1_o, v2_o, v3_o, v4_i, v5_i]
- for ax0_ax1_ax3_ax4_ax5_fused in range(4096):
- with Ts.sblock("C_reindex_shared"):
- v0 = Ts.axis.spatial(32, ax0_0_0_ax1_0_0_fused // 16 * 8 + ax0_0_1_ax1_0_1_fused * 4 + ax0_ax1_ax3_ax4_ax5_fused // 1024)
- v1 = Ts.axis.spatial(16, ax0_0_0_ax1_0_0_fused % 16)
- v2 = Ts.axis.spatial(2, ax2)
- v3 = Ts.axis.spatial(4, ax0_ax1_ax3_ax4_ax5_fused % 1024 // 256)
- v4 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 256 // 16)
- v5 = Ts.axis.spatial(16, ax0_ax1_ax3_ax4_ax5_fused % 16)
- Ts.reads(C_reindex_shared[v0, v1, v2, v3, v4, v5])
- Ts.writes(C[v4 + v2 * 16 + v0 * 32, v5 + v3 * 16 + v1 * 64])
- Ts.sblock_attr({"meta_schedule.cooperative_fetch": 3})
- C[v4 + v2 * 16 + v0 * 32, v5 + v3 * 16 + v1 * 64] = C_reindex_shared[v0, v1, v2, v3, v4, v5]
- # fmt: on
-
- decision_0 = [
- ("SamplePerfectTile", [4, 2, 4, 2, 1]),
- ("SamplePerfectTile", [16, 1, 1, 1, 4]),
- ("SamplePerfectTile", [128, 2, 1]),
- ("SampleCategorical", 2),
- ("SampleCategorical", 3),
- ("SampleCategorical", 0),
- ]
- mod = te.create_function(
- te_workload.matmul(
- n=1024,
- m=1024,
- k=4095,
- in_dtype="float16",
- out_dtype="float32",
- )
- )
- actual = generate_design_space(
- kind="cuda",
- mod=mod,
- target=tvm.target.Target({"kind": "cuda", "arch": "sm_70"}),
- types=None,
- sch_rules=[multi_level_tiling_tensor_core()]
- + get_rules("cuda", ms.schedule_rule.AutoInline),
- )
- check_sketches(
- mod,
- sketches=actual,
- expected_mods=[padded_matmul_no_padded_output_0],
- expected_decisions=[decision_0],
- )
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
index 7402e7ef87..17f18bf450 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
@@ -647,106 +647,6 @@ B_s1 = T.dynamic("B_s1", "int32")
C_3_s0 = T.dynamic("C_3_s0", "int32")
C_3_s1 = T.dynamic("C_3_s1", "int32")
-@tvm.script.ir_module
-class Conv2dInt8_tensorcore_scheduled:
- @Ts.function
- def main(p0: T.Tensor((16, 56, 56, 64), "int8"), p1: T.Tensor((256, 1, 1, 64), "int8"), p2: T.Tensor((1, 1, 1, 256), "int32"), p3: T.Tensor((1, 1, 1, 256), "int32"), p4: T.Tensor((1, 1, 1, 256), "int64"), p5: T.Tensor((1, 1, 1, 256), "int64"), p6: T.Tensor((1, 1, 1, 256), "int64"), p7: T.Tensor((), "int32"), p8: T.Tensor((1,), "int32"), p9: T.Tensor((16, 56, 56, 256), "int32"), compute: T.Tensor((16, 56, 56, 256), "uint8")):
- T.func_attr({"tirx.noalias": True})
- # with Ts.sblock("root"):
- conv2d_nhwc_reindex_shared = Ts.sblock_alloc_buffer((50176, 256), "int32", scope="shared")
- conv2d_nhwc_reindex_shared_wmma_accumulator = Ts.sblock_alloc_buffer((50176, 256), "int32", scope="wmma.accumulator")
- pad_temp_reindex_shared = Ts.sblock_alloc_buffer((50176, 64), "int8", scope="shared")
- p1_reindex_shared = Ts.sblock_alloc_buffer((1, 1, 256, 64), "int8", scope="shared")
- pad_temp_reindex_shared_wmma_matrix_a = Ts.sblock_alloc_buffer((50176, 64), "int8", scope="wmma.matrix_a")
- p1_reindex_shared_wmma_matrix_b = Ts.sblock_alloc_buffer((1, 1, 256, 64), "int8", scope="wmma.matrix_b")
- for ax2_0_0_ax3_0_0_fused in T.thread_binding(3136, thread="blockIdx.x", annotations={"auto_unroll_max_step": 512, "unroll_explicit": 1}):
- for ax2_0_1_ax3_0_1_fused in T.thread_binding(1, thread="vthread.x"):
- for ax2_0_2_ax3_0_2_fused in T.thread_binding(16, thread="threadIdx.x"):
- for ax0_0_init, ax1_0_init, ax0_1_init, ax1_1_init, ax2_0_3_init, ax3_0_3_init, ax0_2_init, ax1_2_init, ax2_0_4_init, ax3_0_4_init in T.grid(1, 1, 1, 1, 1, 1, 1, 1, 1, 1):
- with Ts.sblock("conv2d_nhwc_o_init"):
- v0_o = Ts.axis.spatial(1, ax0_0_init + ax0_1_init + ax0_2_init)
- v1_o = Ts.axis.spatial(1, ax1_0_init + ax1_1_init + ax1_2_init)
- v2_o = Ts.axis.spatial(3136, ax2_0_0_ax3_0_0_fused // 8 * 8 + ax2_0_2_ax3_0_2_fused // 2 + ax2_0_3_init + ax2_0_4_init)
- v3_o = Ts.axis.spatial(16, ax2_0_0_ax3_0_0_fused % 8 * 2 + ax2_0_2_ax3_0_2_fused % 2 + ax3_0_3_init + ax3_0_4_init)
- Ts.reads()
- Ts.writes(conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.thread_extent_high_inclusive": 1024, "meta_schedule.thread_extent_low_inclusive": 32, "warp_execution": 1})
- C = Ts.match_buffer(conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int32", strides=(C_s0, C_s1), scope="wmma.accumulator", offset_factor=16)
- T.gpu_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0))
- for ax0_0, ax1_0, ax4_0_0 in T.grid(1, 1, 2):
- for ax0_ax1_fused_0 in range(16):
- for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
- for ax0_ax1_fused_2 in T.vectorized(16):
- with Ts.sblock("pad_temp_reindex_shared"):
- v0 = Ts.axis.spatial(50176, ax2_0_0_ax3_0_0_fused // 8 * 128 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 16 + ax0_ax1_fused_2) // 32)
- v1 = Ts.axis.spatial(64, ax4_0_0 * 32 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 16 + ax0_ax1_fused_2) % 32)
- Ts.reads(p0[v0 // 3136, v0 % 3136 // 56, v0 % 56, v1])
- Ts.writes(pad_temp_reindex_shared[v0, v1])
- Ts.sblock_attr({"buffer_dim_align": [[0, 0, 32, 16]]})
- pad_temp_reindex_shared[v0, v1] = p0[v0 // 3136, v0 % 3136 // 56, v0 % 56, v1]
- for ax0_ax1_ax2_ax3_fused_0 in range(8):
- for ax0_ax1_ax2_ax3_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
- for ax0_ax1_ax2_ax3_fused_2 in T.vectorized(8):
- with Ts.sblock("p1_reindex_shared"):
- v0 = Ts.axis.spatial(1, 0)
- v1 = Ts.axis.spatial(1, 0)
- v2 = Ts.axis.spatial(256, ax2_0_0_ax3_0_0_fused % 8 * 32 + (ax0_ax1_ax2_ax3_fused_0 * 128 + ax0_ax1_ax2_ax3_fused_1 * 8 + ax0_ax1_ax2_ax3_fused_2) // 32)
- v3 = Ts.axis.spatial(64, ax4_0_0 * 32 + (ax0_ax1_ax2_ax3_fused_0 * 128 + ax0_ax1_ax2_ax3_fused_1 * 8 + ax0_ax1_ax2_ax3_fused_2) % 32)
- Ts.reads(p1[v2, v0, v1, v3])
- Ts.writes(p1_reindex_shared[v0, v1, v2, v3])
- Ts.sblock_attr({"buffer_dim_align": [[0, 2, 32, 16]]})
- p1_reindex_shared[v0, v1, v2, v3] = p1[v2, v0, v1, v3]
- for ax0_1, ax1_1, ax4_0_1 in T.grid(1, 1, 1):
- for ax0_0_1, ax1_0_1 in T.grid(1, 2):
- with Ts.sblock("pad_temp_reindex_shared_wmma.matrix_a_o"):
- v0_o = Ts.axis.spatial(3136, ax2_0_0_ax3_0_0_fused // 8 * 8 + ax2_0_2_ax3_0_2_fused // 2 + ax0_0_1)
- v1_o = Ts.axis.spatial(4, ax4_0_0 * 2 + ax1_0_1)
- Ts.reads(pad_temp_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(pad_temp_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- A = Ts.match_buffer(pad_temp_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16], (16, 16), "int8", strides=(A_s0, A_s1), scope="shared", offset_factor=16)
- C = Ts.match_buffer(pad_temp_reindex_shared_wmma_matrix_a[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16], (16, 16), "int8", strides=(C_1_s0, C_1_s1), scope="wmma.matrix_a", offset_factor=16)
- T.gpu_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.access_ptr("int8", A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
- for ax0, ax1, ax2_0, ax3_0 in T.grid(1, 1, 1, 2):
- with Ts.sblock("p1_reindex_shared_wmma.matrix_b_o"):
- v0_o, v1_o = Ts.axis.remap("SS", [ax0, ax1])
- v2_o = Ts.axis.spatial(16, ax2_0_0_ax3_0_0_fused % 8 * 2 + ax2_0_2_ax3_0_2_fused % 2 + ax2_0)
- v3_o = Ts.axis.spatial(4, ax4_0_0 * 2 + ax3_0)
- Ts.reads(p1_reindex_shared[v0_o, v1_o, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.writes(p1_reindex_shared_wmma_matrix_b[v0_o, v1_o, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- A = Ts.match_buffer(p1_reindex_shared[v0_o, v1_o, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=(A_1_s0, A_1_s1), scope="shared", offset_factor=16)
- C = Ts.match_buffer(p1_reindex_shared_wmma_matrix_b[v0_o, v1_o, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=(C_2_s0, C_2_s1), scope="wmma.matrix_b", offset_factor=16)
- T.gpu_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.access_ptr("int8", A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
- for ax2_0_3, ax3_0_3, ax0_2, ax1_2, ax4_0_2, ax2_0_4, ax3_0_4 in T.grid(1, 1, 1, 1, 2, 1, 1):
- with Ts.sblock("conv2d_nhwc_o_update"):
- v0_o = Ts.axis.spatial(1, ax0_0 + ax0_1 + ax0_2)
- v1_o = Ts.axis.spatial(1, ax1_0 + ax1_1 + ax1_2)
- v2_o = Ts.axis.spatial(3136, ax2_0_0_ax3_0_0_fused // 8 * 8 + ax2_0_2_ax3_0_2_fused // 2 + ax2_0_3 + ax2_0_4)
- v3_o = Ts.axis.spatial(16, ax2_0_0_ax3_0_0_fused % 8 * 2 + ax2_0_2_ax3_0_2_fused % 2 + ax3_0_3 + ax3_0_4)
- v4_o = Ts.axis.reduce(4, ax4_0_0 * 2 + ax4_0_1 * 2 + ax4_0_2)
- Ts.reads(conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], pad_temp_reindex_shared_wmma_matrix_a[v2_o * 16:v2_o * 16 + 16, v4_o * 16:v4_o * 16 + 16], p1_reindex_shared_wmma_matrix_b[v0_o, v1_o, v3_o * 16:v3_o * 16 + 16, v4_o * 16:v4_o * 16 + 16])
- Ts.writes(conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
- Ts.sblock_attr({"meta_schedule.thread_extent_high_inclusive": 1024, "meta_schedule.thread_extent_low_inclusive": 32, "warp_execution": 1})
- A = Ts.match_buffer(pad_temp_reindex_shared_wmma_matrix_a[v2_o * 16:v2_o * 16 + 16, v4_o * 16:v4_o * 16 + 16], (16, 16), "int8", strides=(A_2_s0, A_2_s1), scope="wmma.matrix_a", offset_factor=16)
- B = Ts.match_buffer(p1_reindex_shared_wmma_matrix_b[v0_o, v1_o, v3_o * 16:v3_o * 16 + 16, v4_o * 16:v4_o * 16 + 16], (16, 16), "int8", strides=(B_s0, B_s1), scope="wmma.matrix_b", offset_factor=16)
- C = Ts.match_buffer(conv2d_nhwc_reindex_shared_wmma_accumulator[v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int32", strides=(C_3_s0, C_3_s1), scope="wmma.accumulator", offset_factor=16)
- T.gpu_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
- for ax0_0, ax1_0 in T.grid(1, 1):
- with Ts.sblock("conv2d_nhwc_reindex_shared_wmma.accumulator_o"):
- v0_o = Ts.axis.spatial(3136, ax2_0_0_ax3_0_0_fused // 8 * 8 + ax2_0_2_ax3_0_2_fused // 2 + ax0_0)
- v1_o = Ts.axis.spatial(16, ax2_0_0_ax3_0_0_fused % 8 * 2 + ax2_0_2_ax3_0_2_fused % 2 + ax1_0)
- Ts.reads(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- Ts.writes(conv2d_nhwc_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16])
- A = Ts.match_buffer(conv2d_nhwc_reindex_shared_wmma_accumulator[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16], (16, 16), "int32", strides=(A_3_s0, A_3_s1), scope="wmma.accumulator", offset_factor=16)
- C = Ts.match_buffer(conv2d_nhwc_reindex_shared[v0_o * 16:v0_o * 16 + 16, v1_o * 16:v1_o * 16 + 16], (16, 16), "int32", strides=(C_4_s0, C_4_s1), scope="shared", offset_factor=16)
- T.gpu_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.access_ptr("int32", C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
- for ax0, ax1_0 in T.grid(128, 2):
- for ax1_1 in T.thread_binding(16, thread="threadIdx.x"):
- with Ts.sblock("conv2d_nhwc_reindex_shared"):
- v0 = Ts.axis.spatial(50176, ax2_0_0_ax3_0_0_fused // 8 * 128 + ax0)
- v1 = Ts.axis.spatial(256, ax2_0_0_ax3_0_0_fused % 8 * 32 + ax1_0 * 16 + ax1_1)
- Ts.reads(p7[()], conv2d_nhwc_reindex_shared[v0, v1], p2[0, 0, 0, v1], p3[0, 0, 0, v1], p4[0, 0, 0, v1], p5[0, 0, 0, v1], p6[0, 0, 0, v1], p8[0], p9[v0 // 3136, v0 % 3136 // 56, v0 % 56, v1])
- Ts.writes(compute[v0 // 3136, v0 % 3136 // 56, v0 % 56, v1])
- compute[v0 // 3136, v0 % 3136 // 56, v0 % 56, v1] = T.max(T.min(T.Cast("uint8", T.max(T.min(T.Cast("int32", T.shift_right(T.Cast("int64", T.Cast("int32", T.Cast("uint8", T.max(T.min(p7[()] + T.Cast("int32", T.shift_right(T.Cast("int64", conv2d_nhwc_reindex_shared[v0, v1] - p2[0, 0, 0, v1] + p3[0, 0, 0, v1]) * p4[0, 0, 0, v1] + p5[0, 0, 0, v1], p6[0, 0, 0, v1])), 255), 0))) - p8[0]) * T.Cast("int64", 1098990753), (31) - (1))) + p9[v0 // 3136, v0 % 3136 // 56, v [...]
@tvm.script.ir_module
class Conv2dInt8_NCHWc:
@@ -2050,450 +1950,6 @@ def test_dense_add_gpu():
verify(Dense, apply_anchor_trace, DenseAdd, "cuda", DenseAdd_scheduled_gpu)
-def test_conv2d_int8_tensorcore():
- def apply_trace(sch):
- b0 = sch.get_sblock(name="pad_temp", func_name="main")
- b1 = sch.get_sblock(name="conv2d_nhwc", func_name="main")
- b2 = sch.get_sblock(name="T_subtract", func_name="main")
- b3 = sch.get_sblock(name="T_add", func_name="main")
- b4 = sch.get_sblock(name="T_cast", func_name="main")
- b5 = sch.get_sblock(name="T_multiply", func_name="main")
- b6 = sch.get_sblock(name="T_add_1", func_name="main")
- b7 = sch.get_sblock(name="T_right_shift", func_name="main")
- b8 = sch.get_sblock(name="T_cast_1", func_name="main")
- b9 = sch.get_sblock(name="T_add_2", func_name="main")
- b10 = sch.get_sblock(name="compute", func_name="main")
- b11 = sch.get_sblock(name="T_cast_2", func_name="main")
- b12 = sch.get_sblock(name="T_cast_3", func_name="main")
- b13 = sch.get_sblock(name="T_subtract_1", func_name="main")
- b14 = sch.get_sblock(name="compute_1", func_name="main")
- b15 = sch.get_sblock(name="root", func_name="main")
- sch.annotate(block_or_loop=b1, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")
- b16 = sch.reindex(block=b1, buffer=("write", 0))
- b17 = sch.reindex(block=b1, buffer=("read", 0))
- b18 = sch.reindex(block=b1, buffer=("read", 1))
- sch.transform_layout(
- block=b1,
- buffer=("read", 0),
- index_map=lambda nn, yy, xx, rc: (
- (((nn * 3136) + (yy * 56)) + xx),
- rc,
- ),
- pad_value=None,
- )
- sch.transform_layout(
- block=b1,
- buffer=("read", 1),
- index_map=lambda ff, ry, rx, rc: (
- ry,
- rx,
- ff,
- rc,
- ),
- pad_value=None,
- )
- sch.transform_layout(
- block=b1,
- buffer=("write", 0),
- index_map=lambda nn, yy, xx, ff: (
- (((nn * 3136) + (yy * 56)) + xx),
- ff,
- ),
- pad_value=None,
- )
- sch.transform_block_layout(
- block=b16,
- index_map=lambda nn, yy, xx, ff: (
- (((nn * 3136) + (yy * 56)) + xx),
- ff,
- ),
- )
- sch.transform_block_layout(
- block=b17,
- index_map=lambda nn, yy, xx, rc: (
- (((nn * 3136) + (yy * 56)) + xx),
- rc,
- ),
- )
- sch.transform_block_layout(
- block=b18,
- index_map=lambda ff, ry, rx, rc: (
- ry,
- rx,
- ff,
- rc,
- ),
- )
- sch.transform_block_layout(
- block=b1,
- index_map=lambda nn, yy, xx, ff, ry, rx, rc: (
- ry,
- rx,
- (((nn * 3136) + (yy * 56)) + xx),
- ff,
- rc,
- ),
- )
- l19, l20, l21, l22, l23 = sch.get_loops(block=b1)
- l24, l25 = sch.split(loop=l23, factors=[None, 16], preserve_unit_iters=True)
- l26, l27 = sch.split(loop=l22, factors=[None, 16], preserve_unit_iters=True)
- l28, l29 = sch.split(loop=l21, factors=[None, 16], preserve_unit_iters=True)
- l30, l31, l32, l33, l34, l35, l36, l37 = sch.get_loops(block=b1)
- sch.reorder(l34, l36, l29, l27, l25)
- b38 = sch.blockize(target=l29)
- sch.annotate(
- block_or_loop=b38,
- ann_key="meta_schedule.auto_tensorize",
- ann_val="wmma_sync_16x16x16_s8s8s32_trans",
- )
- sch.annotate(
- block_or_loop=b38,
- ann_key="meta_schedule.auto_tensorize_init",
- ann_val="wmma_fill_16x16x16_s32",
- )
- sch.annotate(block_or_loop=b38, ann_key="warp_execution", ann_val=1)
- l39, l40, l41, l42, l43 = sch.get_loops(block=b38)
- v44, v45, v46 = sch.sample_perfect_tile(
- loop=l39, n=3, max_innermost_factor=4, decision=[1, 1, 1]
- )
- l47, l48, l49 = sch.split(loop=l39, factors=[v44, v45, v46], preserve_unit_iters=True)
- v50, v51, v52 = sch.sample_perfect_tile(
- loop=l40, n=3, max_innermost_factor=4, decision=[1, 1, 1]
- )
- l53, l54, l55 = sch.split(loop=l40, factors=[v50, v51, v52], preserve_unit_iters=True)
- v56, v57, v58, v59, v60 = sch.sample_perfect_tile(
- loop=l41, n=5, max_innermost_factor=4, decision=[392, 1, 8, 1, 1]
- )
- l61, l62, l63, l64, l65 = sch.split(
- loop=l41, factors=[v56, v57, v58, v59, v60], preserve_unit_iters=True
- )
- v66, v67, v68, v69, v70 = sch.sample_perfect_tile(
- loop=l42, n=5, max_innermost_factor=4, decision=[8, 1, 2, 1, 1]
- )
- l71, l72, l73, l74, l75 = sch.split(
- loop=l42, factors=[v66, v67, v68, v69, v70], preserve_unit_iters=True
- )
- v76, v77, v78 = sch.sample_perfect_tile(
- loop=l43, n=3, max_innermost_factor=4, decision=[2, 1, 2]
- )
- l79, l80, l81 = sch.split(loop=l43, factors=[v76, v77, v78], preserve_unit_iters=True)
- sch.reorder(
- l61,
- l71,
- l62,
- l72,
- l63,
- l73,
- l47,
- l53,
- l79,
- l48,
- l54,
- l80,
- l64,
- l74,
- l49,
- l55,
- l81,
- l65,
- l75,
- )
- l82 = sch.fuse(l61, l71, preserve_unit_iters=True)
- sch.bind(loop=l82, thread_axis="blockIdx.x")
- l83 = sch.fuse(l62, l72, preserve_unit_iters=True)
- sch.bind(loop=l83, thread_axis="vthread.x")
- l84 = sch.fuse(l63, l73, preserve_unit_iters=True)
- sch.bind(loop=l84, thread_axis="threadIdx.x")
- sch.annotate(
- block_or_loop=b38, ann_key="meta_schedule.thread_extent_low_inclusive", ann_val=32
- )
- sch.annotate(
- block_or_loop=b38, ann_key="meta_schedule.thread_extent_high_inclusive", ann_val=1024
- )
- b85 = sch.cache_write(block=b38, write_buffer_index=0, storage_scope="shared")
- sch.reverse_compute_at(block=b85, loop=l83, preserve_unit_loops=True, index=-1)
- b86 = sch.cache_write(block=b38, write_buffer_index=0, storage_scope="wmma.accumulator")
- sch.reverse_compute_at(block=b86, loop=l84, preserve_unit_loops=True, index=-1)
- v87 = sch.sample_categorical(
- candidates=[1, 2, 3, 4, 8, 16],
- probs=[
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- ],
- decision=0,
- )
- sch.annotate(block_or_loop=b85, ann_key="meta_schedule.cooperative_fetch", ann_val=v87)
- sch.reverse_compute_inline(block=b16)
- l88, l89, l90, l91, l92 = sch.get_loops(block=b86)
- l93, l94 = sch.split(loop=l92, factors=[None, 16], preserve_unit_iters=True)
- l95, l96 = sch.split(loop=l91, factors=[None, 16], preserve_unit_iters=True)
- l97, l98, l99, l100, l101, l102, l103 = sch.get_loops(block=b86)
- sch.reorder(l102, l96, l94)
- b104 = sch.blockize(target=l96)
- sch.annotate(
- block_or_loop=b104,
- ann_key="meta_schedule.auto_tensorize",
- ann_val="wmma_store_16x16x16_s32_shared",
- )
- b105 = sch.cache_read(
- block=b38, read_buffer_index=0, storage_scope="shared", consumer_blocks=[b38]
- )
- sch.compute_at(block=b105, loop=l79, preserve_unit_loops=True, index=-1)
- l106, l107, l108, l109, l110, l111, l112, l113 = sch.get_loops(block=b105)
- l114 = sch.fuse(l112, l113, preserve_unit_iters=True)
- v115 = sch.sample_categorical(
- candidates=[1, 2, 3, 4, 8, 16],
- probs=[
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- ],
- decision=5,
- )
- sch.annotate(block_or_loop=b105, ann_key="meta_schedule.cooperative_fetch", ann_val=v115)
- b116 = sch.cache_read(
- block=b38, read_buffer_index=1, storage_scope="shared", consumer_blocks=[b38]
- )
- sch.compute_at(block=b116, loop=l79, preserve_unit_loops=True, index=-1)
- l117, l118, l119, l120, l121, l122, l123, l124, l125, l126 = sch.get_loops(block=b116)
- l127 = sch.fuse(l123, l124, l125, l126, preserve_unit_iters=True)
- v128 = sch.sample_categorical(
- candidates=[1, 2, 3, 4, 8, 16],
- probs=[
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- 0.16666666666666666,
- ],
- decision=4,
- )
- sch.annotate(block_or_loop=b116, ann_key="meta_schedule.cooperative_fetch", ann_val=v128)
- b129 = sch.cache_read(block=b38, read_buffer_index=0, storage_scope="wmma.matrix_a")
- sch.compute_at(block=b129, loop=l80, preserve_unit_loops=True, index=-1)
- l130, l131, l132, l133, l134, l135, l136, l137, l138, l139, l140 = sch.get_loops(block=b129)
- l141, l142 = sch.split(loop=l140, factors=[None, 16], preserve_unit_iters=True)
- l143, l144 = sch.split(loop=l139, factors=[None, 16], preserve_unit_iters=True)
- (
- l145,
- l146,
- l147,
- l148,
- l149,
- l150,
- l151,
- l152,
- l153,
- l154,
- l155,
- l156,
- l157,
- ) = sch.get_loops(block=b129)
- sch.reorder(l156, l144, l142)
- b158 = sch.blockize(target=l144)
- sch.annotate(
- block_or_loop=b158,
- ann_key="meta_schedule.auto_tensorize",
- ann_val="wmma_load_16x16x16_s8_a_shared",
- )
- b159 = sch.cache_read(block=b38, read_buffer_index=1, storage_scope="wmma.matrix_b")
- sch.compute_at(block=b159, loop=l80, preserve_unit_loops=True, index=-1)
- (
- l160,
- l161,
- l162,
- l163,
- l164,
- l165,
- l166,
- l167,
- l168,
- l169,
- l170,
- l171,
- l172,
- ) = sch.get_loops(block=b159)
- l173, l174 = sch.split(loop=l172, factors=[None, 16], preserve_unit_iters=True)
- l175, l176 = sch.split(loop=l171, factors=[None, 16], preserve_unit_iters=True)
- (
- l177,
- l178,
- l179,
- l180,
- l181,
- l182,
- l183,
- l184,
- l185,
- l186,
- l187,
- l188,
- l189,
- l190,
- l191,
- ) = sch.get_loops(block=b159)
- sch.reorder(l190, l176, l174)
- b192 = sch.blockize(target=l176)
- sch.annotate(
- block_or_loop=b192,
- ann_key="meta_schedule.auto_tensorize",
- ann_val="wmma_load_16x16x16_s8_b_trans_shared",
- )
- sch.compute_inline(block=b17)
- sch.compute_inline(block=b18)
- sch.storage_align(block=b105, buffer_index=0, axis=-2, factor=32, offset=16)
- sch.storage_align(block=b116, buffer_index=0, axis=-2, factor=32, offset=16)
- sch.reverse_compute_inline(block=b14)
- sch.reverse_compute_inline(block=b13)
- sch.reverse_compute_inline(block=b12)
- sch.reverse_compute_inline(block=b11)
- sch.reverse_compute_inline(block=b10)
- sch.reverse_compute_inline(block=b9)
- sch.reverse_compute_inline(block=b8)
- sch.reverse_compute_inline(block=b7)
- sch.reverse_compute_inline(block=b6)
- sch.reverse_compute_inline(block=b5)
- sch.reverse_compute_inline(block=b4)
- sch.reverse_compute_inline(block=b3)
- sch.reverse_compute_inline(block=b2)
- sch.compute_inline(block=b0)
- v193 = sch.sample_categorical(
- candidates=[0, 16, 64, 512, 1024],
- probs=[
- 0.20000000000000001,
- 0.20000000000000001,
- 0.20000000000000001,
- 0.20000000000000001,
- 0.20000000000000001,
- ],
- decision=3,
- )
- sch.annotate(block_or_loop=b15, ann_key="meta_schedule.unroll_explicit", ann_val=v193)
- sch.enter_postproc()
- sch.unannotate(block_or_loop=b85, ann_key="meta_schedule.cooperative_fetch")
- l194, l195, l196, l197 = sch.get_loops(block=b85)
- l198, l199 = sch.split(loop=l197, factors=[None, 16], preserve_unit_iters=True)
- sch.bind(loop=l199, thread_axis="threadIdx.x")
- sch.unannotate(block_or_loop=b105, ann_key="meta_schedule.cooperative_fetch")
- l200, l201, l202, l203, l204, l205, l206 = sch.get_loops(block=b105)
- l207, l208, l209 = sch.split(loop=l206, factors=[None, 16, 16], preserve_unit_iters=True)
- sch.vectorize(loop=l209)
- sch.bind(loop=l208, thread_axis="threadIdx.x")
- sch.unannotate(block_or_loop=b116, ann_key="meta_schedule.cooperative_fetch")
- l210, l211, l212, l213, l214, l215, l216 = sch.get_loops(block=b116)
- l217, l218, l219 = sch.split(loop=l216, factors=[None, 16, 8], preserve_unit_iters=True)
- sch.vectorize(loop=l219)
- sch.bind(loop=l218, thread_axis="threadIdx.x")
- b220 = sch.get_sblock(name="root", func_name="main")
- sch.unannotate(block_or_loop=b220, ann_key="meta_schedule.unroll_explicit")
- b221, b222, b223, b224, b225, b226, b227 = sch.get_child_blocks(b220)
- l228, l229, l230, l231, l232, l233, l234, l235, l236 = sch.get_loops(block=b221)
- sch.annotate(block_or_loop=l228, ann_key="auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l228, ann_key="unroll_explicit", ann_val=1)
- l237, l238, l239, l240, l241, l242, l243, l244, l245 = sch.get_loops(block=b222)
- sch.annotate(block_or_loop=l237, ann_key="auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l237, ann_key="unroll_explicit", ann_val=1)
- l246, l247, l248, l249, l250, l251, l252, l253, l254, l255, l256 = sch.get_loops(block=b223)
- sch.annotate(block_or_loop=l246, ann_key="auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l246, ann_key="unroll_explicit", ann_val=1)
- (
- l257,
- l258,
- l259,
- l260,
- l261,
- l262,
- l263,
- l264,
- l265,
- l266,
- l267,
- l268,
- l269,
- ) = sch.get_loops(block=b224)
- sch.annotate(block_or_loop=l257, ann_key="auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l257, ann_key="unroll_explicit", ann_val=1)
- (
- l270,
- l271,
- l272,
- l273,
- l274,
- l275,
- l276,
- l277,
- l278,
- l279,
- l280,
- l281,
- l282,
- l283,
- l284,
- l285,
- ) = sch.get_loops(block=b225)
- sch.annotate(block_or_loop=l270, ann_key="auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l270, ann_key="unroll_explicit", ann_val=1)
- l286, l287, l288, l289, l290 = sch.get_loops(block=b226)
- sch.annotate(block_or_loop=l286, ann_key="auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l286, ann_key="unroll_explicit", ann_val=1)
- l291, l292, l293, l294, l295 = sch.get_loops(block=b227)
- sch.annotate(block_or_loop=l291, ann_key="auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l291, ann_key="unroll_explicit", ann_val=1)
- b296 = sch.get_sblock(name="conv2d_nhwc_o", func_name="main")
- (
- l297,
- l298,
- l299,
- l300,
- l301,
- l302,
- l303,
- l304,
- l305,
- l306,
- l307,
- l308,
- l309,
- l310,
- l311,
- l312,
- ) = sch.get_loops(block=b296)
- b313 = sch.decompose_reduction(block=b296, loop=l300)
- sch.unannotate(block_or_loop=b313, ann_key="meta_schedule.auto_tensorize")
- sch.annotate(
- block_or_loop=b313,
- ann_key="meta_schedule.auto_tensorize",
- ann_val="wmma_fill_16x16x16_s32",
- )
- sch.unannotate(block_or_loop=b296, ann_key="meta_schedule.auto_tensorize_init")
- sch.unannotate(block_or_loop=b313, ann_key="meta_schedule.auto_tensorize_init")
- b314 = sch.get_sblock(name="conv2d_nhwc_o_init", func_name="main")
- sch.unannotate(block_or_loop=b314, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b314, tensor_intrin="wmma_fill_16x16x16_s32")
- b315 = sch.get_sblock(name=sch.get(b158).name_hint, func_name="main")
- sch.unannotate(block_or_loop=b315, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b315, tensor_intrin="wmma_load_16x16x16_s8_a_shared")
- b316 = sch.get_sblock(name="p1_reindex_shared_wmma.matrix_b_o", func_name="main")
- sch.unannotate(block_or_loop=b316, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b316, tensor_intrin="wmma_load_16x16x16_s8_b_trans_shared")
- b317 = sch.get_sblock(name="conv2d_nhwc_o_update", func_name="main")
- sch.unannotate(block_or_loop=b317, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b317, tensor_intrin="wmma_sync_16x16x16_s8s8s32_trans")
- b318 = sch.get_sblock(name=sch.get(b104).name_hint, func_name="main")
- sch.unannotate(block_or_loop=b318, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b318, tensor_intrin="wmma_store_16x16x16_s32_shared")
-
- verify(Conv2dInt8, apply_trace, Conv2dInt8_target, "cuda", Conv2dInt8_tensorcore_scheduled)
-
-
def test_conv2d_int8_vnni():
def apply_trace(sch):
b0 = sch.get_sblock(name="compile_engine_const", func_name="main")
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
index cdabb56efd..5d5034b8b3 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
@@ -33,10 +33,6 @@ from tvm.s_tir.schedule.analysis import (
is_output_block,
suggest_index_map,
)
-from tvm.s_tir.tensor_intrin.cuda import (
- WMMA_SYNC_16x16x16_f16f16f16_INTRIN,
- WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
-)
from tvm.s_tir.tensor_intrin.x86 import dot_product_16x4_u8i8i32_desc
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
@@ -315,10 +311,10 @@ def test_get_tensorize_loop_mapping_padding_matmul():
s = Schedule(matmul)
block = s.get_sblock("C")
- desc = TensorIntrin.get(WMMA_SYNC_16x16x16_f16f16f16_INTRIN).desc
+ desc = TensorIntrin.get("mma_sync_m16n8k8_f16f16f16").desc
info = get_tensorize_loop_mapping(s, block, desc, allow_padding=True)
assert info is not None
- expected_padding = [16, 1, 16]
+ expected_padding = [16, 1, 8]
actual_padding = info.block_iter_paddings
assert actual_padding is not None
assert len(actual_padding) == len(expected_padding)
@@ -345,7 +341,7 @@ def test_get_auto_tensorize_mapping_info_conv2d():
check_index_map(
conv2d,
"conv2d_nhwc",
- WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
+ "mma_sync_m16n8k8_f16f16f32",
lambda n, h, w, c, rh, rw, rc: (n * 256 + h * 16 + w, c, rh * 192 + rw * 64 + rc),
)
@@ -357,7 +353,7 @@ def test_get_auto_tensorize_mapping_info_conv2d_unit_batch():
check_index_map(
conv2d,
"conv2d_nhwc",
- WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
+ "mma_sync_m16n8k8_f16f16f32",
lambda n, h, w, c, rh, rw, rc: (n * 256 + h * 16 + w, c, rh * 192 + rw * 64 + rc),
)
@@ -367,9 +363,7 @@ def test_get_auto_tensorize_mapping_info_batch_matmul(b, m, n, k):
matmul = create_function(
te_workload.batch_matmul_nkkm(b, m, n, k, in_dtype="float16", out_dtype="float32")
)
- check_index_map(
- matmul, "Z", WMMA_SYNC_16x16x16_f16f16f32_INTRIN, lambda b, m, n, k: (b, m, n, k)
- )
+ check_index_map(matmul, "Z", "mma_sync_m16n8k8_f16f16f32", lambda b, m, n, k: (b, m, n, k))
@pytest.mark.parametrize(
@@ -390,7 +384,7 @@ def test_get_auto_tensorize_mapping_info_batch_matmul(b, m, n, k):
)
def test_get_auto_tensorize_mapping_info_matmul(n, m, k, expected):
matmul = create_function(te_workload.matmul(n, m, k, in_dtype="float16", out_dtype="float32"))
- check_index_map(matmul, "C", WMMA_SYNC_16x16x16_f16f16f32_INTRIN, expected)
+ check_index_map(matmul, "C", "mma_sync_m16n8k8_f16f16f32", expected)
def test_is_output_block():
diff --git a/tests/python/tirx/script/test_tirx_script_source_locations.py b/tests/python/tirx/script/test_tirx_script_source_locations.py
index 6570ca3892..da84c0b877 100644
--- a/tests/python/tirx/script/test_tirx_script_source_locations.py
+++ b/tests/python/tirx/script/test_tirx_script_source_locations.py
@@ -28,7 +28,7 @@ from tvm_ffi import structural_walk
import tvm
import tvm.testing
-from tvm import ir
+from tvm import ir, tirx
from tvm.ir import Call, SequentialSpan, TensorLoad, assert_structural_equal, prim
from tvm.script import ir as I
from tvm.script import tirx as T
@@ -298,7 +298,7 @@ def test_native_view_keeps_producer_identity_name_and_span(monkeypatch):
def test_native_binding_preserves_metadata_but_binds_buffer_expressions():
# Inert metadata must retain resource identity; a non-Var buffer expression still needs a
# located binding.
- from tvm import ir, tirx
+ from tvm import ir
from tvm.script.ir_builder import base
producer_span = ir.Span(ir.SourceName("producer.py"), 7, 7, 2, 19)