mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
[Refactor] Remove GEMM v1 and promote gemm_py to be the canonical gemm op (#2033)
* [Refactor] Remove GEMM v1 and merge gemm_py into the canonical gemm op
GEMM v2 has been the default for a while and no tests or examples still
depend on v1. This removes the legacy C++ GemmNode path entirely and
promotes the former GemmPyNode (Python-delegated lowering) to be the one
and only GEMM tile op.
C++:
- Merge src/op/gemm_py.{h,cc} into src/op/gemm.{h,cc}. GemmPyNode becomes
GemmNode, GemmPy becomes Gemm, FFI type key "tl.GemmPy" becomes "tl.Gemm",
and the op names "tl.tileop.{gemm_py,wgmma_gemm_py,tcgen05_gemm_py}"
collapse to "tl.tileop.{gemm,wgmma_gemm,tcgen05_gemm}". The shared
infrastructure (GemmWarpPolicyNode, GemmInst enum,
GemmWarpPolicyNode::computeWarpPartition) is preserved in place so
GemmSP/GemmSPPy keep compiling unchanged.
- Delete the tl_gemm builtin, its codegen dispatch in
codegen_{cuda,hip,cutedsl}.cc, and its matching branches in
inject_fence_proxy.cc and merge_shared_memory_allocations.cc. tl_gemm_sp
is kept as the remaining async-proxy gemm builtin.
- Drop #include "gemm_py.h" and GemmPyNode/GemmPy references from
inject_pipeline.cc, producer_consumer_ws.cc, and lower_blackwell_2sm.cc.
Python:
- tileop/gemm: rename GemmPy -> Gemm, register as "tl.Gemm", and rename
the FFI callbacks tl.gemm_py.{lower,infer_layout} -> tl.gemm.{lower,
infer_layout}. _ffi_api.GemmPyGemmInst becomes GemmGetGemmInst.
- language/gemm_op.py: delete gemm_v1/gemm_v2 wrappers. gemm() now emits
tl.tileop.gemm directly; wgmma_gemm/tcgen05_gemm drop the _py suffix.
- env.py: remove TILELANG_USE_GEMM_V1 and the use_gemm_v1() helper.
- ir.py: drop the empty v1-era `class Gemm(Node, Scriptable): ...`
placeholder that was colliding with the real tileop.gemm.Gemm wrapper
under the "tl.Gemm" FFI type key.
- contrib/cutedsl: delete cutedsl/gemm_v1.py (runtime helper for the v1
cutedsl lowering path) and its re-export in cutedsl/__init__.py.
Tests / maint:
- testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py switches
from T.gemm_v2 to T.gemm.
- testing/python/transform/test_tilelang_transform_inject_fence_proxy.py::
test_lower_fence_proxy is rewritten to exercise tl.tl_gemm_sp instead
of the removed tl.tl_gemm, preserving the fence-proxy assertion.
- Delete maint/gemm_v2/ (v1-vs-v2 benchmarking scripts) and drop the
v1-specific comment from maint/scripts/run_local_ci_test.sh.
C++ headers under src/tl_templates/{cuda,hip,cpp}/gemm*.h are left in
place as reference implementations; they are no longer reachable from
the generated code path since the only consumer (tl_gemm builtin) is
gone, but keeping them avoids a large, hard-to-reverse deletion.
* [Fix] Restore cutedsl/gemm_v1.py (Gemm_SM80 / Gemm_SM90 helpers)
The previous commit deleted cutedsl/gemm_v1.py on the assumption that it
was the CuteDSL lowering helper for the removed T.gemm_v1 path. It is
actually a runtime helper module that defines SM80/SM90 CUTE DSL
primitives (make_smem_layout_AB, get_tma_atom, etc.) which are emitted
into generated cutedsl kernel code via
`import tilelang.contrib.cutedsl as tl; tl.Gemm_SM90.get_tma_atom(...)`
(see jit/adapter/cutedsl/wrapper.py CUBIN_TMA_ATOM_INIT_TEMPLATE).
The "v1" in the filename refers to the CuteDSL helper version, not to
the deleted T.gemm_v1 op, so removing it broke the cutedsl code path for
attention_sink, blocksparse_attention, blocksparse_gemm, and deepgemm
examples with `module 'tilelang.contrib.cutedsl' has no attribute
'Gemm_SM90'`. Restore the file and its `from .gemm_v1 import *` wildcard
re-export.
* [Refactor] Restore tl_gemm builtin as a call_intrin path to tl:: templates
The previous commit in this PR deleted the `tl_gemm` builtin entirely on
the argument that no code emitted it anymore. That was too aggressive:
keeping the builtin preserves a deliberate extension point for any
Python-side gemm lowering backend (for example a future
GemmLegacyExtern) to delegate into the existing C++ templates in
src/tl_templates/<target>/gemm*.h via
T.call_intrin("handle", tir.op.Op.get("tl.tl_gemm"),
op_instance_str, A_ptr, B_ptr, C_ptr)
without any extra C++ plumbing. It is also the only thing that keeps
those tl_templates headers reachable rather than dead code.
Restore:
- tl_gemm builtin declaration in src/op/builtin.h (with an updated
docstring describing the call_intrin entry point) and its
TIR_DEFINE_TL_BUILTIN registration in src/op/builtin.cc.
- Codegen dispatch branches in codegen_cuda.cc, codegen_hip.cc, and
codegen_cutedsl.cc that turn a tl_gemm Call into a PrintCallExtern
emission of `tl::gemm_xx<...>(A, B, C)`.
- SMEM alignment and async-proxy fence handling for tl_gemm in
merge_shared_memory_allocations.cc and inject_fence_proxy.cc. These
passes are tied to the builtin's semantics (not to the removed v1
GemmNode) and are still correct the moment anything emits tl_gemm.
- test_lower_fence_proxy in test_tilelang_transform_inject_fence_proxy
reverts to exercising `tl.tl_gemm` instead of the temporary
`tl.tl_gemm_sp` substitution introduced when the builtin was gone.
Note that the current sole gemm lowering path (the merged GemmNode that
delegates to Python via tl.gemm.lower) still emits inline mma/wgmma/
tcgen5mma via the Python macro generators, not tl_gemm. So restoring
the builtin is a no-op for the default path - it only re-enables the
"call into v1 templates from Python" option for anyone who wants it.
* [Fix] Default HIP execution backend to tvm_ffi, not cython
resolve_execution_backend() only promoted cuda and metal to tvm_ffi for
the "auto" path, so hip fell into the else branch and picked cython. But
allowed_backends_for_target() already lists tvm_ffi as an allowed
backend for hip, and the cuda default is tvm_ffi for good reasons:
- The cython SourceWrapper path in src/tilelang/jit/adapter/wrapper.py
has known gaps (for example T.alloc_global is not handled in
create_dispatch_func, which breaks any test that uses an internally
allocated global scratch buffer).
- libgen.py appends user-supplied compile_flags verbatim to the
hipcc/nvcc command. nvcc accepts --use_fast_math, clang++/hipcc does
not, so a CUDA-style flag leaks into HIP compilation.
Neither of these leak through the tvm_ffi backend. Aligning hip's
default with cuda makes the two paths symmetric and unbreaks
test_alloc_global and test_all_attrs_together_lazy on ROCm without
having to mark the tests skip or rewrite the cython wrapper.
cython is still available on hip via execution_backend='cython' for
anyone who explicitly wants it.
* [Test] Gate CUDA-only language tests with requires_cuda
Three language tests hardcode target='cuda' (or rely on CUDA-only
codegen behavior) but lacked a CUDA availability marker, so they failed
on the ROCm CI runner with "Cannot find global function
target.build.tilelang_cuda" once the earlier pre-existing failures got
past --maxfail=3. Add @tilelang.testing.requires_cuda to:
- test_tilelang_if_range
- test_annotate_min_blocks_per_sm_launch_bounds
- test_tilelang_transpose and test_tilelang_transpose_square
Also replace the hand-rolled __main__ block in test_tilelang_transpose
with tilelang.testing.main() so the skip markers are honored when the
file is run directly.
* [Test] Delete tile library gemm test
Remove testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py.
This commit is contained in:
Lei Wang
@@ -258,4 +258,4 @@ However, fixing a specific layout introduces several potential issues:
|
||||
|
||||
3. Alignment requirements: `CUTLASS` enforces strict alignment checks, and many hyperparameter configurations can lead to compilation errors. (For reference, sm8x was implemented in `CUTLASS 2`.)
|
||||
|
||||
`T.gemm_sp_v2` was designed to address these limitations, following the approach of `T.gemm_v2`. It lowers directly to PTX, removing the need for a fixed metadata layout.
|
||||
`T.gemm_sp_v2` was designed to address these limitations, following the approach of `T.gemm`. It lowers directly to PTX, removing the need for a fixed metadata layout.
|
||||
|
||||
@@ -1,808 +0,0 @@
|
||||
# pytest correctness_evaluation.py -n 32
|
||||
import pytest
|
||||
from tilelang import tvm as tvm
|
||||
import tilelang.testing
|
||||
from tilelang import language as T
|
||||
import torch
|
||||
|
||||
|
||||
def matmul(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope="shared.dyn")
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope="shared.dyn")
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
if trans_A:
|
||||
T.copy(A[k * block_K, by * block_M], A_shared)
|
||||
else:
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
if trans_B:
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
else:
|
||||
T.copy(B[k * block_K, bx * block_N], B_shared)
|
||||
T.gemm(A_shared, B_shared, C_local, trans_A, trans_B)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def _compile_and_check(
|
||||
program,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
):
|
||||
kernel = tilelang.compile(
|
||||
program,
|
||||
out_idx=[2],
|
||||
pass_configs={
|
||||
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
||||
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
||||
# tilelang.PassConfigKey.TIR_USE_ASYNC_COPY: False,
|
||||
},
|
||||
)
|
||||
|
||||
print(kernel.get_kernel_source())
|
||||
|
||||
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
||||
|
||||
def ref_program(A, B):
|
||||
if trans_A:
|
||||
A = A.T
|
||||
if trans_B:
|
||||
B = B.T
|
||||
if in_dtype == T.float32:
|
||||
A = (A.view(torch.int32) - 0x1000).view(torch.float32)
|
||||
B = (B.view(torch.int32) - 0x1000).view(torch.float32)
|
||||
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
||||
C = C.to(torch.__getattribute__(out_dtype))
|
||||
return C
|
||||
|
||||
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
||||
print("assert_allclose")
|
||||
|
||||
|
||||
def run_gemm(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=2,
|
||||
num_threads=128,
|
||||
):
|
||||
if block_N >= 256 or block_M >= 256 or block_K >= 256:
|
||||
num_stages = 0
|
||||
program = matmul(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
|
||||
_compile_and_check(program, trans_A, trans_B, in_dtype, out_dtype)
|
||||
|
||||
|
||||
def matmul_rs(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
A_frag_shape = A_shared_shape
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope="shared.dyn")
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope="shared.dyn")
|
||||
A_frag = T.alloc_fragment(A_frag_shape, in_dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
if trans_A:
|
||||
T.copy(A[k * block_K, by * block_M], A_shared)
|
||||
else:
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
if trans_B:
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
else:
|
||||
T.copy(B[k * block_K, bx * block_N], B_shared)
|
||||
T.copy(A_shared, A_frag)
|
||||
T.gemm_v2(A_frag, B_shared, C_local, trans_A, trans_B)
|
||||
# T.gemm(A_frag, B_shared, C_local, trans_A, trans_B)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def run_gemm_rs(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=2,
|
||||
num_threads=128,
|
||||
):
|
||||
if block_N >= 256 or block_M >= 256 or block_K >= 256:
|
||||
num_stages = 0
|
||||
program = matmul_rs(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
_compile_and_check(program, trans_A, trans_B, in_dtype, out_dtype)
|
||||
|
||||
|
||||
def matmul_sr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
B_frag_shape = B_shared_shape
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope="shared.dyn")
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope="shared.dyn")
|
||||
B_frag = T.alloc_fragment(B_frag_shape, in_dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
if trans_A:
|
||||
T.copy(A[k * block_K, by * block_M], A_shared)
|
||||
else:
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
if trans_B:
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
else:
|
||||
T.copy(B[k * block_K, bx * block_N], B_shared)
|
||||
T.copy(B_shared, B_frag)
|
||||
T.gemm_v2(A_shared, B_frag, C_local, trans_A, trans_B)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def run_gemm_sr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=2,
|
||||
num_threads=128,
|
||||
):
|
||||
if block_N >= 256 or block_M >= 256 or block_K >= 256:
|
||||
num_stages = 0
|
||||
program = matmul_sr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
|
||||
_compile_and_check(program, trans_A, trans_B, in_dtype, out_dtype)
|
||||
|
||||
|
||||
def matmul_rr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
A_frag_shape = A_shared_shape
|
||||
B_frag_shape = B_shared_shape
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope="shared.dyn")
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope="shared.dyn")
|
||||
A_frag = T.alloc_fragment(A_frag_shape, in_dtype)
|
||||
B_frag = T.alloc_fragment(B_frag_shape, in_dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
if trans_A:
|
||||
T.copy(A[k * block_K, by * block_M], A_shared)
|
||||
else:
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
if trans_B:
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
else:
|
||||
T.copy(B[k * block_K, bx * block_N], B_shared)
|
||||
T.copy(A_shared, A_frag)
|
||||
T.copy(B_shared, B_frag)
|
||||
T.gemm_v2(A_frag, B_frag, C_local, trans_A, trans_B)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def run_gemm_rr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=2,
|
||||
num_threads=128,
|
||||
):
|
||||
if block_N >= 256 or block_M >= 256 or block_K >= 256:
|
||||
num_stages = 0
|
||||
program = matmul_rr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
|
||||
_compile_and_check(program, trans_A, trans_B, in_dtype, out_dtype)
|
||||
|
||||
|
||||
M_VALUES = [64, 128, 256]
|
||||
N_VALUES = [16, 32, 64, 128, 256, 512]
|
||||
K_VALUES = [16, 32, 64, 128]
|
||||
K_VALUES_8Bit = [32, 64, 128]
|
||||
NUM_THREADS_VALUES = [128, 256]
|
||||
|
||||
|
||||
def _generate_dtype_cases(k_values, num_threads):
|
||||
"""Generate dtype test cases for given K values and num_threads."""
|
||||
return (
|
||||
[
|
||||
pytest.param(
|
||||
k,
|
||||
T.float16,
|
||||
T.float16,
|
||||
T.float16,
|
||||
num_threads,
|
||||
id=f"K{k}-float16-float16-float16-threads{num_threads}",
|
||||
)
|
||||
for k in k_values
|
||||
]
|
||||
+ [
|
||||
pytest.param(
|
||||
k,
|
||||
T.int8,
|
||||
T.int32,
|
||||
T.int32,
|
||||
num_threads,
|
||||
id=f"K{k}-int8-int32-int32-threads{num_threads}",
|
||||
)
|
||||
for k in K_VALUES_8Bit
|
||||
]
|
||||
+ [
|
||||
pytest.param(
|
||||
k,
|
||||
T.float8_e5m2,
|
||||
T.float32,
|
||||
T.float32,
|
||||
num_threads,
|
||||
id=f"K{k}-float8_e5m2-float32-float32-threads{num_threads}",
|
||||
)
|
||||
for k in K_VALUES_8Bit
|
||||
]
|
||||
+ [
|
||||
pytest.param(
|
||||
k,
|
||||
T.float8_e4m3fn,
|
||||
T.float32,
|
||||
T.float32,
|
||||
num_threads,
|
||||
id=f"K{k}-float8_e4m3fn-float32-float32-threads{num_threads}",
|
||||
)
|
||||
for k in K_VALUES_8Bit
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
# num_threads=128 can work with any N
|
||||
FALSE_TRUE_CASES_128 = _generate_dtype_cases(K_VALUES, 128)
|
||||
# num_threads=256 requires N >= 32
|
||||
FALSE_TRUE_CASES_256 = _generate_dtype_cases(K_VALUES, 256)
|
||||
FALSE_TRUE_CASES = FALSE_TRUE_CASES_128 + FALSE_TRUE_CASES_256
|
||||
|
||||
|
||||
def _ensure_torch_dtypes(*dtype_names):
|
||||
import torch
|
||||
|
||||
for name in set(dtype_names):
|
||||
if not hasattr(torch, name):
|
||||
pytest.skip(f"Torch does not expose dtype {name}")
|
||||
|
||||
|
||||
def _skip_if_threads_exceed_n(num_threads, n):
|
||||
"""Skip test if num_threads=256 and N < 32."""
|
||||
if num_threads == 256 and n < 32:
|
||||
pytest.skip(f"num_threads=256 requires N >= 32, but N={n}")
|
||||
|
||||
|
||||
def run_gemm_rs_false_true(m, n, k, in_dtype, out_dtype, accum_dtype, num_threads=128):
|
||||
run_gemm_rs(m, n, k * 3, False, True, in_dtype, out_dtype, accum_dtype, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_rs_false_false(m, n, k, num_threads=128):
|
||||
run_gemm_rs(m, n, k * 3, False, False, T.float16, T.float16, T.float16, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_rs_true_false(m, n, k, num_threads=128):
|
||||
run_gemm_rs(m, n, k * 3, True, False, T.float16, T.float16, T.float16, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_rs_true_true(m, n, k, num_threads=128):
|
||||
run_gemm_rs(m, n, k * 3, True, True, T.float16, T.float16, T.float16, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_sr_false_true(m, n, k, in_dtype, out_dtype, accum_dtype, num_threads=128):
|
||||
run_gemm_sr(m, n, k * 3, False, True, in_dtype, out_dtype, accum_dtype, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_sr_false_false(m, n, k, num_threads=128):
|
||||
run_gemm_sr(m, n, k * 3, False, False, T.float16, T.float16, T.float16, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_sr_true_false(m, n, k, num_threads=128):
|
||||
run_gemm_sr(m, n, k * 3, True, False, T.float16, T.float16, T.float16, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_sr_true_true(m, n, k, num_threads=128):
|
||||
run_gemm_sr(m, n, k * 3, True, True, T.float16, T.float16, T.float16, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_rr_false_true(m, n, k, in_dtype, out_dtype, accum_dtype, num_threads=128):
|
||||
run_gemm_rr(m, n, k * 3, False, True, in_dtype, out_dtype, accum_dtype, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_rr_false_false(m, n, k, num_threads=128):
|
||||
run_gemm_rr(m, n, k * 3, False, False, T.float16, T.float16, T.float16, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_rr_true_false(m, n, k, num_threads=128):
|
||||
run_gemm_rr(m, n, k * 3, True, False, T.float16, T.float16, T.float16, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
def run_gemm_rr_true_true(m, n, k, num_threads=128):
|
||||
run_gemm_rr(m, n, k * 3, True, True, T.float16, T.float16, T.float16, m, n, k, num_threads=num_threads)
|
||||
|
||||
|
||||
TRANS_CASES = [
|
||||
pytest.param(False, False, id="nn"),
|
||||
pytest.param(False, True, id="nt"),
|
||||
pytest.param(True, False, id="tn"),
|
||||
pytest.param(True, True, id="tt"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def _setup_tilelang_environment():
|
||||
tilelang.disable_cache()
|
||||
tilelang.testing.set_random_seed(42)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k,in_dtype,out_dtype,accum_dtype,num_threads", FALSE_TRUE_CASES)
|
||||
def test_gemm_false_true(m, n, k, in_dtype, out_dtype, accum_dtype, num_threads):
|
||||
import torch
|
||||
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
|
||||
required_torch_attrs = {
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
}
|
||||
for attr in required_torch_attrs:
|
||||
if not hasattr(torch, attr):
|
||||
pytest.skip(f"Torch does not expose dtype {attr}")
|
||||
run_gemm(
|
||||
m,
|
||||
n,
|
||||
k * 3,
|
||||
False,
|
||||
True,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
num_threads=num_threads,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_false_false(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
run_gemm(
|
||||
m,
|
||||
n,
|
||||
k * 3,
|
||||
False,
|
||||
False,
|
||||
T.float16,
|
||||
T.float16,
|
||||
T.float16,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
num_threads=num_threads,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_true_false(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
run_gemm(
|
||||
m,
|
||||
n,
|
||||
k * 3,
|
||||
True,
|
||||
False,
|
||||
T.float16,
|
||||
T.float16,
|
||||
T.float16,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
num_threads=num_threads,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_true_true(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
run_gemm(
|
||||
m,
|
||||
n,
|
||||
k * 3,
|
||||
True,
|
||||
True,
|
||||
T.float16,
|
||||
T.float16,
|
||||
T.float16,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
num_threads=num_threads,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k,in_dtype,out_dtype,accum_dtype,num_threads", FALSE_TRUE_CASES)
|
||||
def test_gemm_rs_false_true(m, n, k, in_dtype, out_dtype, accum_dtype, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(in_dtype, out_dtype, accum_dtype)
|
||||
run_gemm_rs_false_true(m, n, k, in_dtype, out_dtype, accum_dtype, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_rs_false_false(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(T.float16)
|
||||
run_gemm_rs_false_false(m, n, k, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_rs_true_false(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(T.float16)
|
||||
run_gemm_rs_true_false(m, n, k, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_rs_true_true(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(T.float16)
|
||||
run_gemm_rs_true_true(m, n, k, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k,in_dtype,out_dtype,accum_dtype,num_threads", FALSE_TRUE_CASES)
|
||||
def test_gemm_sr_false_true(m, n, k, in_dtype, out_dtype, accum_dtype, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(in_dtype, out_dtype, accum_dtype)
|
||||
run_gemm_sr_false_true(m, n, k, in_dtype, out_dtype, accum_dtype, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_sr_false_false(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(T.float16)
|
||||
run_gemm_sr_false_false(m, n, k, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_sr_true_false(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(T.float16)
|
||||
run_gemm_sr_true_false(m, n, k, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_sr_true_true(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(T.float16)
|
||||
run_gemm_sr_true_true(m, n, k, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k,in_dtype,out_dtype,accum_dtype,num_threads", FALSE_TRUE_CASES)
|
||||
def test_gemm_rr_false_true(m, n, k, in_dtype, out_dtype, accum_dtype, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(in_dtype, out_dtype, accum_dtype)
|
||||
run_gemm_rr_false_true(m, n, k, in_dtype, out_dtype, accum_dtype, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_rr_false_false(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(T.float16)
|
||||
run_gemm_rr_false_false(m, n, k, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_rr_true_false(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(T.float16)
|
||||
run_gemm_rr_true_false(m, n, k, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
@pytest.mark.parametrize("num_threads", NUM_THREADS_VALUES, ids=lambda v: f"threads{v}")
|
||||
def test_gemm_rr_true_true(m, n, k, num_threads):
|
||||
_skip_if_threads_exceed_n(num_threads, n)
|
||||
_ensure_torch_dtypes(T.float16)
|
||||
run_gemm_rr_true_true(m, n, k, num_threads)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_gemm(
|
||||
M=64,
|
||||
N=192,
|
||||
K=64,
|
||||
trans_A=False,
|
||||
trans_B=False,
|
||||
in_dtype=T.bfloat16,
|
||||
out_dtype=T.bfloat16,
|
||||
dtypeAccum=T.float32,
|
||||
block_M=64,
|
||||
block_N=192,
|
||||
block_K=64,
|
||||
num_stages=0,
|
||||
num_threads=256,
|
||||
)
|
||||
|
||||
# # Test Pass
|
||||
# for m in [64, 128, 256]:
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [16, 32, 64, 128]:
|
||||
# print(f"======================= Test {m} {n} {k} False True =============================")
|
||||
# run_gemm(m, n, k * 3, False, True, T.float16, T.float16, T.float16, m, n, k, 2, 128)
|
||||
# print(f"Test {m} {n} {k} Pass")
|
||||
|
||||
# # Test Pass
|
||||
# for m in [64, 128, 256]:
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [16, 32, 64, 128]:
|
||||
# print(f"======================= Test {m} {n} {k} False False =============================")
|
||||
# run_gemm(m, n, k * 3, False, False, T.float16, T.float16, T.float16, m, n, k, 2, 128)
|
||||
# print(f"Test {m} {n} {k} Pass")
|
||||
|
||||
# # Test Pass
|
||||
# for m in [64, 128, 256]:
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [16, 32, 64, 128]:
|
||||
# print(f"======================= Test {m} {n} {k} True False =============================")
|
||||
# run_gemm(m, n, k * 3, True, False, T.float16, T.float16, T.float16, m, n, k, 2, 128)
|
||||
# print(f"Test {m}, {n} {k} Pass")
|
||||
# print(f"Test {n} Pass")
|
||||
|
||||
# # Test Pass
|
||||
# for m in [64, 128, 256]:
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [16, 32, 64, 128]:
|
||||
# print(f"======================= Test {m} {n} {k} True True =============================")
|
||||
# run_gemm(m, n, k * 3, True, True, T.float16, T.float16, T.float16, m, n, k, 2, 128)
|
||||
# print(f"Test {m}, {n} {k} Pass")
|
||||
# print(f"Test {n} Pass")
|
||||
|
||||
# Test Pass
|
||||
# for m in [64, 128, 256]:
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [16, 32, 64, 128]:
|
||||
# print(f"======================= Test {m} {n} {k} False True =============================")
|
||||
# run_gemm_rs(m, n, k * 3, False, True, T.float16, T.float16, T.float16, m, n, k, 2, 128)
|
||||
# print(f"Test {m} {n} {k} Pass")
|
||||
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [16, 32, 64, 128]:
|
||||
# run_gemm_rs(64, n, k, False, False, T.float16, T.float16, T.float16, 64, n, k, 0, 256)
|
||||
# print(f"Test {64} {n} {k} Pass")
|
||||
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [16, 32, 64, 128]:
|
||||
# run_gemm(64, n, k, False, False, T.float16, T.float16, T.float16, 64, n, k, 0, 256)
|
||||
# print(f"Test {64} {n} {k} Pass")
|
||||
@@ -1,350 +0,0 @@
|
||||
# pytest maint/gemm_v2/correctness_evaluation_sm70.py -n 32
|
||||
import pytest
|
||||
from tilelang import tvm as tvm
|
||||
import tilelang.testing
|
||||
from tilelang import language as T
|
||||
|
||||
|
||||
def matmul(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope="shared.dyn")
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope="shared.dyn")
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
if trans_A:
|
||||
T.copy(A[k * block_K, by * block_M], A_shared)
|
||||
else:
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
if trans_B:
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
else:
|
||||
T.copy(B[k * block_K, bx * block_N], B_shared)
|
||||
# T.gemm(A_shared, B_shared, C_local, trans_A, trans_B)
|
||||
T.gemm_v2(A_shared, B_shared, C_local, trans_A, trans_B)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def _compile_and_check(
|
||||
program,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
):
|
||||
kernel = tilelang.compile(
|
||||
program,
|
||||
out_idx=[2],
|
||||
pass_configs={
|
||||
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
||||
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
||||
# tilelang.PassConfigKey.TIR_USE_ASYNC_COPY: False,
|
||||
},
|
||||
)
|
||||
|
||||
print(kernel.get_kernel_source())
|
||||
|
||||
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
||||
|
||||
def ref_program(A, B):
|
||||
import torch
|
||||
|
||||
if trans_A:
|
||||
A = A.T
|
||||
if trans_B:
|
||||
B = B.T
|
||||
if in_dtype == T.float32:
|
||||
A = (A.view(torch.int32) - 0x1000).view(torch.float32)
|
||||
B = (B.view(torch.int32) - 0x1000).view(torch.float32)
|
||||
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
||||
C = C.to(torch.__getattribute__(out_dtype))
|
||||
return C
|
||||
|
||||
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
||||
print("assert_allclose")
|
||||
|
||||
|
||||
def run_gemm(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=3,
|
||||
num_threads=128,
|
||||
):
|
||||
program = matmul(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
|
||||
_compile_and_check(program, trans_A, trans_B, in_dtype, out_dtype)
|
||||
|
||||
|
||||
def matmul_rs(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
A_frag_shape = A_shared_shape
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope="shared.dyn")
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope="shared.dyn")
|
||||
A_frag = T.alloc_fragment(A_frag_shape, in_dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
if trans_A:
|
||||
T.copy(A[k * block_K, by * block_M], A_shared)
|
||||
else:
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
if trans_B:
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
else:
|
||||
T.copy(B[k * block_K, bx * block_N], B_shared)
|
||||
T.copy(A_shared, A_frag)
|
||||
T.gemm_v2(A_frag, B_shared, C_local, trans_A, trans_B)
|
||||
# T.gemm(A_frag, B_shared, C_local, trans_A, trans_B)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def run_gemm_rs(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=3,
|
||||
num_threads=128,
|
||||
):
|
||||
program = matmul_rs(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
_compile_and_check(program, trans_A, trans_B, in_dtype, out_dtype)
|
||||
|
||||
|
||||
M_VALUES = [64, 128]
|
||||
N_VALUES = [32, 64, 128]
|
||||
K_VALUES = [16, 32, 64]
|
||||
FALSE_TRUE_CASES = [
|
||||
pytest.param(
|
||||
k,
|
||||
T.float16,
|
||||
T.float16,
|
||||
T.float16,
|
||||
id=f"K{k}-float16-float16-float16",
|
||||
)
|
||||
for k in K_VALUES
|
||||
] + [
|
||||
pytest.param(
|
||||
k,
|
||||
T.float16,
|
||||
T.float16,
|
||||
T.float32,
|
||||
id=f"K{k}-float16-float16-float32",
|
||||
)
|
||||
for k in K_VALUES
|
||||
]
|
||||
|
||||
|
||||
def _ensure_torch_dtypes(*dtype_names):
|
||||
import torch
|
||||
|
||||
for name in set(dtype_names):
|
||||
if not hasattr(torch, name):
|
||||
pytest.skip(f"Torch does not expose dtype {name}")
|
||||
|
||||
|
||||
def run_gemm_rs_false_true(m, n, k, in_dtype, out_dtype, accum_dtype):
|
||||
run_gemm_rs(m, n, k * 3, False, True, in_dtype, out_dtype, accum_dtype, m, n, k, 2, 128)
|
||||
|
||||
|
||||
def run_gemm_rs_false_false(m, n, k):
|
||||
run_gemm_rs(m, n, k * 3, False, False, T.float16, T.float16, T.float16, m, n, k, 2, 128)
|
||||
|
||||
|
||||
TRANS_CASES = [
|
||||
pytest.param(False, False, id="nn"),
|
||||
pytest.param(False, True, id="nt"),
|
||||
pytest.param(True, False, id="tn"),
|
||||
pytest.param(True, True, id="tt"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def _setup_tilelang_environment():
|
||||
tilelang.disable_cache()
|
||||
tilelang.testing.set_random_seed(42)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k,in_dtype,out_dtype,accum_dtype", FALSE_TRUE_CASES)
|
||||
def test_gemm_false_true(m, n, k, in_dtype, out_dtype, accum_dtype):
|
||||
import torch
|
||||
|
||||
required_torch_attrs = {
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
}
|
||||
for attr in required_torch_attrs:
|
||||
if not hasattr(torch, attr):
|
||||
pytest.skip(f"Torch does not expose dtype {attr}")
|
||||
run_gemm(
|
||||
m,
|
||||
n,
|
||||
k * 3,
|
||||
False,
|
||||
True,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
2,
|
||||
128,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
def test_gemm_false_false(m, n, k):
|
||||
run_gemm(
|
||||
m,
|
||||
n,
|
||||
k * 3,
|
||||
False,
|
||||
False,
|
||||
T.float16,
|
||||
T.float16,
|
||||
T.float16,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
2,
|
||||
128,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k,in_dtype,out_dtype,accum_dtype", FALSE_TRUE_CASES)
|
||||
def test_gemm_rs_false_true(m, n, k, in_dtype, out_dtype, accum_dtype):
|
||||
_ensure_torch_dtypes(in_dtype, out_dtype, accum_dtype)
|
||||
run_gemm_rs_false_true(m, n, k, in_dtype, out_dtype, accum_dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k", K_VALUES, ids=lambda v: f"K{v}")
|
||||
def test_gemm_rs_false_false(m, n, k):
|
||||
_ensure_torch_dtypes(T.float16)
|
||||
run_gemm_rs_false_false(m, n, k)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tilelang.testing.main()
|
||||
|
||||
# # Test Pass
|
||||
# for m in [64, 128]:
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [16, 32, 64]:
|
||||
# print(f"======================= Test {m} {n} {k} False True =============================")
|
||||
# run_gemm(m, n, k * 3, False, True, T.float16, T.float16, T.float16, m, n, k, 2, 128)
|
||||
# print(f"Test {m} {n} {k} Pass")
|
||||
|
||||
# # Test Pass
|
||||
# for m in [64, 128]:
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [16, 32, 64]:
|
||||
# print(f"======================= Test {m} {n} {k} False False =============================")
|
||||
# run_gemm(m, n, k * 3, False, False, T.float16, T.float16, T.float16, m, n, k, 2, 128)
|
||||
# print(f"Test {m} {n} {k} Pass")
|
||||
@@ -1,231 +0,0 @@
|
||||
# pytest correctness_evaluation.py -n 32
|
||||
import pytest
|
||||
from tilelang import tvm as tvm
|
||||
import tilelang.testing
|
||||
import tilelang.language as T
|
||||
|
||||
|
||||
def matmul(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype)
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
||||
C_tmem = T.alloc_tmem([block_M, block_N], accum_dtype)
|
||||
mbar = T.alloc_barrier(1)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
C_shared = T.alloc_shared((block_M, block_N), out_dtype)
|
||||
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
T.tcgen05_gemm(A_shared, B_shared, C_tmem, trans_A, trans_B, mbar=mbar, clear_accum=k == 0)
|
||||
T.mbarrier_wait_parity(mbar, k % 2)
|
||||
|
||||
T.copy(C_tmem, C_local)
|
||||
T.copy(C_local, C_shared)
|
||||
|
||||
T.copy(C_shared, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def _compile_and_check(
|
||||
program,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
):
|
||||
kernel = tilelang.compile(
|
||||
program,
|
||||
out_idx=[2],
|
||||
pass_configs={
|
||||
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
||||
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
||||
},
|
||||
)
|
||||
|
||||
print(kernel.get_kernel_source())
|
||||
|
||||
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
||||
|
||||
def ref_program(A, B):
|
||||
import torch
|
||||
|
||||
if trans_A:
|
||||
A = A.T
|
||||
if trans_B:
|
||||
B = B.T
|
||||
if in_dtype == T.float32:
|
||||
A = (A.view(torch.int32) - 0x1000).view(torch.float32)
|
||||
B = (B.view(torch.int32) - 0x1000).view(torch.float32)
|
||||
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
||||
C = C.to(torch.__getattribute__(out_dtype))
|
||||
return C
|
||||
|
||||
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
||||
print("assert_allclose")
|
||||
|
||||
|
||||
def run_gemm(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=2,
|
||||
num_threads=128,
|
||||
):
|
||||
if block_N >= 256 or block_M >= 256 or block_K >= 256:
|
||||
num_stages = 0
|
||||
program = matmul(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
|
||||
_compile_and_check(program, trans_A, trans_B, in_dtype, out_dtype)
|
||||
|
||||
|
||||
M_VALUES = [32, 64, 128, 256]
|
||||
N_VALUES = [64, 128, 256, 512]
|
||||
K_VALUES = [16, 32, 64, 128]
|
||||
K_VALUES_8Bit = [32, 64, 128]
|
||||
FALSE_TRUE_CASES = (
|
||||
[
|
||||
pytest.param(
|
||||
k,
|
||||
T.float16,
|
||||
T.float32,
|
||||
T.float32,
|
||||
id=f"K{k}-float16-float-float",
|
||||
)
|
||||
for k in K_VALUES
|
||||
]
|
||||
+ [
|
||||
pytest.param(
|
||||
k,
|
||||
T.float8_e5m2,
|
||||
T.float32,
|
||||
T.float32,
|
||||
id="K32-float8_e5m2-float32-float32",
|
||||
)
|
||||
for k in K_VALUES_8Bit
|
||||
]
|
||||
+ [
|
||||
pytest.param(
|
||||
k,
|
||||
T.int8,
|
||||
T.int32,
|
||||
T.int32,
|
||||
id="K32-int8-int32-int32",
|
||||
)
|
||||
for k in K_VALUES_8Bit
|
||||
]
|
||||
)
|
||||
|
||||
TRANS_CASES = [
|
||||
pytest.param(False, True, id="nt"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("k,in_dtype,out_dtype,accum_dtype", FALSE_TRUE_CASES)
|
||||
def test_gemm_false_true(m, n, k, in_dtype, out_dtype, accum_dtype):
|
||||
import torch
|
||||
|
||||
required_torch_attrs = {
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
}
|
||||
for attr in required_torch_attrs:
|
||||
if not hasattr(torch, attr):
|
||||
pytest.skip(f"Torch does not expose dtype {attr}")
|
||||
run_gemm(
|
||||
m,
|
||||
n,
|
||||
k * 3,
|
||||
False,
|
||||
True,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tilelang.testing.main()
|
||||
|
||||
# # Test Pass
|
||||
# for m in [32, 64, 128, 256]:
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [16, 32, 64, 128]:
|
||||
# if m in [32, 64] and (n not in [64, 128, 256]):
|
||||
# continue
|
||||
# print(f"======================= Test {m} {n} {k} False True =============================")
|
||||
# run_gemm(m, n, k * 3, False, True, T.float16, T.float, T.float, m, n, k, 2, 128)
|
||||
# print(f"Test {m} {n} {k} Pass")
|
||||
|
||||
# # Test Pass
|
||||
# for m in [32, 64, 128, 256]:
|
||||
# for n in [32, 64, 128]:
|
||||
# for k in [16, 32, 64, 128]:
|
||||
# if m in [32, 64] and (n not in [64, 128, 256]):
|
||||
# continue
|
||||
# print(f"======================= Test {m} {n} {k} False True =============================")
|
||||
# run_gemm(m, n, k * 3, False, True, T.float16, T.float, T.float, m, n, k, 2, 256)
|
||||
# print(f"Test {m} {n} {k} Pass")
|
||||
|
||||
# # Test Pass
|
||||
# for m in [32, 64, 128, 256]:
|
||||
# for n in [16, 32, 64, 128]:
|
||||
# for k in [32, 64, 128]:
|
||||
# if m in [32, 64] and (n not in [64, 128, 256]):
|
||||
# continue
|
||||
# print(f"======================= Test {m} {n} {k} False True =============================")
|
||||
# run_gemm(m, n, k * 3, False, True, T.float8_e5m2, T.float, T.float, m, n, k, 2, 128)
|
||||
@@ -1,168 +0,0 @@
|
||||
# pytest correctness_evaluation_tcgen05_2cta.py -n 32
|
||||
import pytest
|
||||
from tilelang import tvm as tvm
|
||||
import tilelang
|
||||
import tilelang.testing
|
||||
import tilelang.language as T
|
||||
|
||||
tilelang.disable_cache()
|
||||
|
||||
|
||||
def matmul_2cta(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
):
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor((M, K), in_dtype),
|
||||
B: T.Tensor((K, N), in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=128, cluster_dims=2) as (bx, by):
|
||||
A_shared = T.alloc_shared((num_stages, block_M, block_K), in_dtype)
|
||||
B_shared = T.alloc_shared((num_stages, block_K, block_N // 2), in_dtype) # each CTA holds half of B
|
||||
C_tmem = T.alloc_tmem([block_M, block_N], accum_dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
C_shared = T.alloc_shared((block_M, block_N), out_dtype)
|
||||
loaded = T.alloc_cluster_barrier([32 * 2] * num_stages)
|
||||
consumed = T.alloc_cluster_barrier([1] * num_stages)
|
||||
tmem_full = T.alloc_barrier([1])
|
||||
|
||||
tx = T.get_thread_binding()
|
||||
cta_id = T.block_rank_in_cluster()
|
||||
T.assume(cta_id < 2)
|
||||
|
||||
T.use_swizzle(16)
|
||||
|
||||
if tx < 32: # warp 0: issue TMA loads
|
||||
for k in T.serial(T.ceildiv(K, block_K)):
|
||||
T.mbarrier_wait_parity(consumed[k % num_stages], ((k // num_stages) & 1) ^ 1)
|
||||
T.tma_copy(
|
||||
A[bx * block_M : (bx + 1) * block_M, k * block_K : (k + 1) * block_K],
|
||||
A_shared[k % num_stages, :, :],
|
||||
barrier=loaded[k % num_stages],
|
||||
)
|
||||
T.tma_copy(
|
||||
B[k * block_K : (k + 1) * block_K, (by * 2 + cta_id) * (block_N // 2) : (by * 2 + cta_id + 1) * (block_N // 2)],
|
||||
B_shared[k % num_stages, :, :],
|
||||
barrier=loaded[k % num_stages],
|
||||
)
|
||||
T.mbarrier_arrive(loaded[k % num_stages], 0) # arrive on leader CTA's barrier
|
||||
elif cta_id == 0 and tx < 64: # warp 1 on leader CTA: issue tcgen5 MMA
|
||||
for k in T.serial(T.ceildiv(K, block_K)):
|
||||
T.mbarrier_wait_parity(loaded[k % num_stages], (k // num_stages) & 1)
|
||||
T.tcgen05_gemm(
|
||||
A_shared[k % num_stages, :, :],
|
||||
B_shared[k % num_stages, :, :],
|
||||
C_tmem,
|
||||
mbar=consumed[k % num_stages],
|
||||
clear_accum=k == 0,
|
||||
use_2cta=True,
|
||||
)
|
||||
T.tcgen05_mma_arrive(tmem_full, arrive_2cta=True)
|
||||
|
||||
T.mbarrier_wait_parity(tmem_full, 0)
|
||||
T.copy(C_tmem, C_local)
|
||||
T.copy(C_local, C_shared)
|
||||
T.copy(C_shared, C[bx * block_M, by * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def _compile_and_check(program, out_dtype):
|
||||
kernel = tilelang.compile(program, out_idx=[2])
|
||||
|
||||
print(kernel.get_kernel_source())
|
||||
|
||||
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
||||
|
||||
def ref_program(A, B):
|
||||
import torch
|
||||
|
||||
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
||||
return C.to(torch.__getattribute__(out_dtype))
|
||||
|
||||
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
||||
print("assert_allclose passed")
|
||||
|
||||
|
||||
def run_gemm(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=4,
|
||||
):
|
||||
program = matmul_2cta(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
)
|
||||
_compile_and_check(program, out_dtype)
|
||||
|
||||
|
||||
M_VALUES = [64, 128, 256]
|
||||
N_VALUES = [64, 128, 256]
|
||||
|
||||
# atom_k=16 for fp16/bf16 (K%16==0), atom_k=32 for fp8/int8 (K%32==0)
|
||||
K_VALUES_16 = [16, 32, 64, 128]
|
||||
K_VALUES_32 = [32, 64, 128]
|
||||
|
||||
# Dtype cases: (block_K, in_dtype, out_dtype, accum_dtype)
|
||||
FP16_CASES = [pytest.param(k, T.float16, T.float32, T.float32, id=f"K{k}-fp16-fp32-fp32") for k in K_VALUES_16]
|
||||
|
||||
FP8_E5M2_CASES = [pytest.param(k, T.float8_e5m2, T.float32, T.float32, id=f"K{k}-fp8e5m2-fp32-fp32") for k in K_VALUES_32]
|
||||
|
||||
INT8_CASES = [pytest.param(k, T.int8, T.int32, T.int32, id=f"K{k}-int8-int32-int32") for k in K_VALUES_32]
|
||||
|
||||
ALL_DTYPE_CASES = FP16_CASES + FP8_E5M2_CASES + INT8_CASES
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", M_VALUES, ids=lambda v: f"M{v}")
|
||||
@pytest.mark.parametrize("n", N_VALUES, ids=lambda v: f"N{v}")
|
||||
@pytest.mark.parametrize("block_k,in_dtype,out_dtype,accum_dtype", ALL_DTYPE_CASES)
|
||||
def test_gemm_2cta(m, n, block_k, in_dtype, out_dtype, accum_dtype):
|
||||
import torch
|
||||
|
||||
for attr in {in_dtype, out_dtype, accum_dtype}:
|
||||
if not hasattr(torch, attr):
|
||||
pytest.skip(f"Torch does not expose dtype {attr}")
|
||||
|
||||
# M = 2 * block_M so ceildiv(M, block_M) = 2 (cluster needs >= 2 tiles in M dim)
|
||||
# K = 3 * block_K to exercise multi-iteration pipelining
|
||||
k = block_k * 3
|
||||
run_gemm(
|
||||
m * 2,
|
||||
n,
|
||||
k,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
block_M=m,
|
||||
block_N=n,
|
||||
block_K=block_k,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tilelang.testing.main()
|
||||
@@ -1,98 +0,0 @@
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--use_v2", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
use_v2 = args.use_v2
|
||||
|
||||
|
||||
# @tilelang.jit(target="cuda")
|
||||
# target currently can be "cuda" or "hip" or "cpu".
|
||||
# if not specified, it will be inferred from the input tensors during compile time
|
||||
@tilelang.jit
|
||||
def matmul(M, N, K, block_M, block_N, block_K, dtype=T.float16, accum_dtype=T.float32):
|
||||
@T.prim_func
|
||||
def matmul_relu_kernel(
|
||||
A: T.Tensor((M, K), dtype),
|
||||
B: T.Tensor((K, N), dtype),
|
||||
C: T.Tensor((M, N), dtype),
|
||||
):
|
||||
# Initialize Kernel Context
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
||||
A_shared = T.alloc_shared((block_M, block_K), dtype)
|
||||
B_shared = T.alloc_shared((block_K, block_N), dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
|
||||
# Enable rasterization for better L2 cache locality (Optional)
|
||||
# T.use_swizzle(panel_size=10, enable=True)
|
||||
|
||||
# Clear local accumulation
|
||||
T.clear(C_local)
|
||||
|
||||
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
|
||||
# Copy tile of A
|
||||
# This is a sugar syntax for parallelized copy
|
||||
T.copy(A[by * block_M, ko * block_K], A_shared)
|
||||
|
||||
# Copy tile of B
|
||||
T.copy(B[ko * block_K, bx * block_N], B_shared)
|
||||
|
||||
# Perform a tile-level GEMM on the shared buffers
|
||||
# Currently we dispatch to the cute/hip on Nvidia/AMD GPUs
|
||||
if use_v2:
|
||||
T.gemm_v2(A_shared, B_shared, C_local)
|
||||
else:
|
||||
T.gemm_v1(A_shared, B_shared, C_local)
|
||||
|
||||
# relu
|
||||
for i, j in T.Parallel(block_M, block_N):
|
||||
C_local[i, j] = T.max(C_local[i, j], 0)
|
||||
|
||||
# Copy result back to global memory
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return matmul_relu_kernel
|
||||
|
||||
|
||||
M = 16384 # M = T.dynamic("m") if you want to use dynamic shape
|
||||
N = 16384
|
||||
K = 16384
|
||||
block_M = 128
|
||||
block_N = 128
|
||||
block_K = 32
|
||||
|
||||
# 1. Define the kernel (matmul) and compile/lower it into an executable module
|
||||
matmul_relu_kernel = matmul(M, N, K, block_M, block_N, block_K)
|
||||
|
||||
# 3. Test the kernel in Python with PyTorch data
|
||||
import torch
|
||||
|
||||
# Create random input tensors on the GPU
|
||||
a = torch.randn(M, K, device="cuda", dtype=torch.float16)
|
||||
b = torch.randn(K, N, device="cuda", dtype=torch.float16)
|
||||
c = torch.empty(M, N, device="cuda", dtype=torch.float16)
|
||||
|
||||
# Run the kernel through the Profiler
|
||||
matmul_relu_kernel(a, b, c)
|
||||
|
||||
print(c)
|
||||
# Reference multiplication using PyTorch
|
||||
ref_c = torch.relu(a @ b)
|
||||
|
||||
# Validate correctness
|
||||
torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
|
||||
print("Kernel output matches PyTorch reference.")
|
||||
|
||||
# 4. Retrieve and inspect the generated CUDA source (optional)
|
||||
# cuda_source = jit_kernel.get_kernel_source()
|
||||
# print("Generated CUDA kernel:\n", cuda_source)
|
||||
|
||||
# 5.Profile latency with kernel
|
||||
profiler = matmul_relu_kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
||||
|
||||
latency = profiler.do_bench()
|
||||
|
||||
print(f"Latency: {latency} ms")
|
||||
@@ -1,98 +0,0 @@
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--use_v2", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
use_v2 = args.use_v2
|
||||
|
||||
|
||||
# @tilelang.jit(target="cuda")
|
||||
# target currently can be "cuda" or "hip" or "cpu".
|
||||
# if not specified, it will be inferred from the input tensors during compile time
|
||||
@tilelang.jit
|
||||
def matmul(M, N, K, block_M, block_N, block_K, dtype=T.float16, accum_dtype=T.float32):
|
||||
@T.prim_func
|
||||
def matmul_relu_kernel(
|
||||
A: T.Tensor((M, K), dtype),
|
||||
B: T.Tensor((K, N), dtype),
|
||||
C: T.Tensor((M, N), dtype),
|
||||
):
|
||||
# Initialize Kernel Context
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
||||
A_shared = T.alloc_shared((block_M, block_K), dtype)
|
||||
B_shared = T.alloc_shared((block_K, block_N), dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
|
||||
# Enable rasterization for better L2 cache locality (Optional)
|
||||
# T.use_swizzle(panel_size=10, enable=True)
|
||||
|
||||
# Clear local accumulation
|
||||
T.clear(C_local)
|
||||
|
||||
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
|
||||
# Copy tile of A
|
||||
# This is a sugar syntax for parallelized copy
|
||||
T.copy(A[by * block_M, ko * block_K], A_shared)
|
||||
|
||||
# Copy tile of B
|
||||
T.copy(B[ko * block_K, bx * block_N], B_shared)
|
||||
|
||||
# Perform a tile-level GEMM on the shared buffers
|
||||
# Currently we dispatch to the cute/hip on Nvidia/AMD GPUs
|
||||
if use_v2:
|
||||
T.gemm_v2(A_shared, B_shared, C_local)
|
||||
else:
|
||||
T.gemm_v1(A_shared, B_shared, C_local)
|
||||
|
||||
# relu
|
||||
for i, j in T.Parallel(block_M, block_N):
|
||||
C_local[i, j] = T.max(C_local[i, j], 0)
|
||||
|
||||
# Copy result back to global memory
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return matmul_relu_kernel
|
||||
|
||||
|
||||
M = 16384 # M = T.dynamic("m") if you want to use dynamic shape
|
||||
N = 16384
|
||||
K = 16384
|
||||
block_M = 128
|
||||
block_N = 128
|
||||
block_K = 64
|
||||
|
||||
# 1. Define the kernel (matmul) and compile/lower it into an executable module
|
||||
matmul_relu_kernel = matmul(M, N, K, block_M, block_N, block_K)
|
||||
|
||||
# 3. Test the kernel in Python with PyTorch data
|
||||
import torch
|
||||
|
||||
# Create random input tensors on the GPU
|
||||
a = torch.randn(M, K, device="cuda", dtype=torch.float16)
|
||||
b = torch.randn(K, N, device="cuda", dtype=torch.float16)
|
||||
c = torch.empty(M, N, device="cuda", dtype=torch.float16)
|
||||
|
||||
# Run the kernel through the Profiler
|
||||
matmul_relu_kernel(a, b, c)
|
||||
|
||||
print(c)
|
||||
# Reference multiplication using PyTorch
|
||||
ref_c = torch.relu(a @ b)
|
||||
|
||||
# Validate correctness
|
||||
torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
|
||||
print("Kernel output matches PyTorch reference.")
|
||||
|
||||
# 4. Retrieve and inspect the generated CUDA source (optional)
|
||||
# cuda_source = jit_kernel.get_kernel_source()
|
||||
# print("Generated CUDA kernel:\n", cuda_source)
|
||||
|
||||
# 5.Profile latency with kernel
|
||||
profiler = matmul_relu_kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
||||
|
||||
latency = profiler.do_bench()
|
||||
|
||||
print(f"Latency: {latency} ms")
|
||||
@@ -1,228 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import tilelang
|
||||
from tilelang.autotuner import *
|
||||
import tilelang.language as T
|
||||
import itertools
|
||||
import argparse
|
||||
from functools import partial
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--batch", type=int, default=128, help="batch size")
|
||||
parser.add_argument("--heads", type=int, default=16, help="heads")
|
||||
parser.add_argument("--seq_q", type=int, default=1024, help="query sequence length")
|
||||
parser.add_argument("--seq_kv", type=int, default=1024, help="key/value sequence length")
|
||||
parser.add_argument("--dim", type=int, default=256, help="dim")
|
||||
parser.add_argument("--is_causal", action="store_true", help="causal")
|
||||
parser.add_argument("--tune", action="store_true", help="tune configs")
|
||||
parser.add_argument("--use_v2", action="store_true")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
use_v2 = args.use_v2
|
||||
|
||||
|
||||
def get_configs():
|
||||
iter_params = dict(block_M=[128], block_N=[128], num_stages=[2], threads=[256])
|
||||
return [dict(zip(iter_params, values)) for values in itertools.product(*iter_params.values())]
|
||||
|
||||
|
||||
@autotune(configs=get_configs(), warmup=10, rep=10)
|
||||
@tilelang.jit(
|
||||
out_idx=[3],
|
||||
pass_configs={
|
||||
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
|
||||
},
|
||||
)
|
||||
def flashattn(batch, heads, seq_q, seq_kv, dim, is_causal, block_M=64, block_N=64, num_stages=0, threads=128):
|
||||
scale = (1.0 / dim) ** 0.5 * 1.44269504 # log2(e)
|
||||
q_shape = [batch, heads, seq_q, dim]
|
||||
kv_shape = [batch, heads, seq_kv, dim]
|
||||
dtype = T.float16
|
||||
accum_dtype = T.float32
|
||||
|
||||
past_len = seq_kv - seq_q
|
||||
assert past_len >= 0, "seq_kv must be greater than or equal to seq_q"
|
||||
|
||||
@T.macro
|
||||
def MMA0(
|
||||
K: T.Tensor(kv_shape, dtype),
|
||||
Q_shared: T.SharedBuffer([block_M, dim], dtype),
|
||||
K_shared: T.SharedBuffer([block_N, dim], dtype),
|
||||
acc_s: T.FragmentBuffer([block_M, block_N], accum_dtype),
|
||||
k: T.int32,
|
||||
bx: T.int32,
|
||||
by: T.int32,
|
||||
bz: T.int32,
|
||||
):
|
||||
T.copy(K[bz, by, k * block_N : (k + 1) * block_N, :], K_shared)
|
||||
if is_causal:
|
||||
for i, j in T.Parallel(block_M, block_N):
|
||||
q_idx = bx * block_M + i + past_len
|
||||
k_idx = k * block_N + j
|
||||
acc_s[i, j] = T.if_then_else(q_idx >= k_idx, 0, -T.infinity(acc_s.dtype))
|
||||
else:
|
||||
T.clear(acc_s)
|
||||
if use_v2:
|
||||
T.gemm_v2(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
|
||||
else:
|
||||
T.gemm_v1(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
|
||||
|
||||
@T.macro
|
||||
def MMA1(
|
||||
V: T.Tensor(kv_shape, dtype),
|
||||
V_shared: T.SharedBuffer([block_N, dim], dtype),
|
||||
acc_s_cast: T.FragmentBuffer([block_M, block_N], dtype),
|
||||
acc_o: T.FragmentBuffer([block_M, dim], accum_dtype),
|
||||
k: T.int32,
|
||||
by: T.int32,
|
||||
bz: T.int32,
|
||||
):
|
||||
T.copy(V[bz, by, k * block_N : (k + 1) * block_N, :], V_shared)
|
||||
# T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
|
||||
if use_v2:
|
||||
T.gemm_v2(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
|
||||
else:
|
||||
T.gemm_v1(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
|
||||
|
||||
@T.macro
|
||||
def Softmax(
|
||||
acc_s: T.FragmentBuffer([block_M, block_N], accum_dtype),
|
||||
acc_s_cast: T.FragmentBuffer([block_M, block_N], dtype),
|
||||
scores_max: T.FragmentBuffer([block_M], accum_dtype),
|
||||
scores_max_prev: T.FragmentBuffer([block_M], accum_dtype),
|
||||
scores_scale: T.FragmentBuffer([block_M], accum_dtype),
|
||||
scores_sum: T.FragmentBuffer([block_M], accum_dtype),
|
||||
logsum: T.FragmentBuffer([block_M], accum_dtype),
|
||||
):
|
||||
T.copy(scores_max, scores_max_prev)
|
||||
T.fill(scores_max, -T.infinity(accum_dtype))
|
||||
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
|
||||
# To do causal softmax, we need to set the scores_max to 0 if it is -inf
|
||||
# This process is called Check_inf in FlashAttention3 code, and it only need to be done
|
||||
# in the first ceil_div(kBlockM, kBlockN) steps.
|
||||
# for i in T.Parallel(block_M):
|
||||
# scores_max[i] = T.if_then_else(scores_max[i] == -T.infinity(accum_dtype), 0, scores_max[i])
|
||||
for i in T.Parallel(block_M):
|
||||
scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
|
||||
|
||||
for i, j in T.Parallel(block_M, block_N):
|
||||
# Instead of computing exp(x - max), we compute exp2(x * log_2(e) -
|
||||
# max * log_2(e)) This allows the compiler to use the ffma
|
||||
# instruction instead of fadd and fmul separately.
|
||||
acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
|
||||
T.reduce_sum(acc_s, scores_sum, dim=1)
|
||||
for i in T.Parallel(block_M):
|
||||
logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
|
||||
T.copy(acc_s, acc_s_cast)
|
||||
|
||||
@T.macro
|
||||
def Rescale(
|
||||
acc_o: T.FragmentBuffer([block_M, dim], accum_dtype),
|
||||
scores_scale: T.FragmentBuffer([block_M], accum_dtype),
|
||||
):
|
||||
for i, j in T.Parallel(block_M, dim):
|
||||
acc_o[i, j] *= scores_scale[i]
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
Q: T.Tensor(q_shape, dtype),
|
||||
K: T.Tensor(kv_shape, dtype),
|
||||
V: T.Tensor(kv_shape, dtype),
|
||||
Output: T.Tensor(q_shape, dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(seq_q, block_M), heads, batch, threads=threads) as (bx, by, bz):
|
||||
Q_shared = T.alloc_shared([block_M, dim], dtype)
|
||||
K_shared = T.alloc_shared([block_N, dim], dtype)
|
||||
V_shared = T.alloc_shared([block_N, dim], dtype)
|
||||
O_shared = T.alloc_shared([block_M, dim], dtype)
|
||||
acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)
|
||||
acc_s_cast = T.alloc_fragment([block_M, block_N], dtype)
|
||||
acc_o = T.alloc_fragment([block_M, dim], accum_dtype)
|
||||
scores_max = T.alloc_fragment([block_M], accum_dtype)
|
||||
scores_max_prev = T.alloc_fragment([block_M], accum_dtype)
|
||||
scores_scale = T.alloc_fragment([block_M], accum_dtype)
|
||||
scores_sum = T.alloc_fragment([block_M], accum_dtype)
|
||||
logsum = T.alloc_fragment([block_M], accum_dtype)
|
||||
|
||||
T.copy(Q[bz, by, bx * block_M : (bx + 1) * block_M, :], Q_shared)
|
||||
T.fill(acc_o, 0)
|
||||
T.fill(logsum, 0)
|
||||
T.fill(scores_max, -T.infinity(accum_dtype))
|
||||
|
||||
loop_range = (
|
||||
T.min(T.ceildiv(seq_kv, block_N), T.ceildiv((bx + 1) * block_M + past_len, block_N))
|
||||
if is_causal
|
||||
else T.ceildiv(seq_kv, block_N)
|
||||
)
|
||||
|
||||
for k in T.Pipelined(loop_range, num_stages=num_stages):
|
||||
MMA0(K, Q_shared, K_shared, acc_s, k, bx, by, bz)
|
||||
Softmax(acc_s, acc_s_cast, scores_max, scores_max_prev, scores_scale, scores_sum, logsum)
|
||||
Rescale(acc_o, scores_scale)
|
||||
MMA1(V, V_shared, acc_s_cast, acc_o, k, by, bz)
|
||||
for i, j in T.Parallel(block_M, dim):
|
||||
acc_o[i, j] /= logsum[i]
|
||||
T.copy(acc_o, O_shared)
|
||||
T.copy(O_shared, Output[bz, by, bx * block_M : (bx + 1) * block_M, :])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def ref_program(Q, K, V, is_causal):
|
||||
dim = Q.size(-1)
|
||||
scores = torch.einsum("bhqd,bhkd->bhqk", Q, K)
|
||||
scores = scores / torch.sqrt(torch.tensor(dim, dtype=scores.dtype))
|
||||
if is_causal:
|
||||
seq_q = Q.size(2)
|
||||
seq_kv = K.size(2)
|
||||
mask = torch.tril(torch.ones(seq_q, seq_kv, device=scores.device), seq_kv - seq_q)
|
||||
mask = mask.unsqueeze(0).unsqueeze(0)
|
||||
scores = scores.masked_fill(mask == 0, float("-inf"))
|
||||
attention_weights = F.softmax(scores, dim=-1)
|
||||
output = torch.einsum("bhqk,bhkd->bhqd", attention_weights, V)
|
||||
return output
|
||||
|
||||
|
||||
def main(
|
||||
batch: int = 1,
|
||||
heads: int = 1,
|
||||
seq_q: int = 256,
|
||||
seq_kv: int = 256,
|
||||
dim: int = 64,
|
||||
is_causal: bool = False,
|
||||
tune: bool = False,
|
||||
):
|
||||
flops_per_matmul = 2.0 * batch * heads * seq_q * seq_kv * dim
|
||||
total_flops = 2 * flops_per_matmul
|
||||
if is_causal:
|
||||
total_flops *= 0.5
|
||||
|
||||
if not tune:
|
||||
kernel = flashattn(batch, heads, seq_q, seq_kv, dim, is_causal, block_M=64, block_N=64, num_stages=0, threads=128)
|
||||
print(kernel.get_kernel_source())
|
||||
ref_program_processed = partial(ref_program, is_causal=is_causal)
|
||||
|
||||
profiler = kernel.get_profiler()
|
||||
profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01)
|
||||
print("All checks pass.")
|
||||
latency = profiler.do_bench(ref_program_processed, warmup=500)
|
||||
print(f"Ref: {latency:.2f} ms")
|
||||
print(f"Ref: {total_flops / latency * 1e-9:.2f} TFlops")
|
||||
latency = profiler.do_bench(warmup=500)
|
||||
print(f"Tile-lang: {latency:.2f} ms")
|
||||
print(f"Tile-lang: {total_flops / latency * 1e-9:.2f} TFlops")
|
||||
else:
|
||||
kernel = flashattn(batch, heads, seq_q, seq_kv, dim, is_causal)
|
||||
best_latency = kernel.latency
|
||||
best_config = kernel.config
|
||||
ref_latency = kernel.ref_latency
|
||||
print(f"Best latency: {best_latency}")
|
||||
print(f"Best TFlops: {total_flops / best_latency * 1e-9}")
|
||||
print(f"Best config: {best_config}")
|
||||
print(f"Ref latency: {ref_latency}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tilelang.disable_cache()
|
||||
main(args.batch, args.heads, args.seq_q, args.seq_kv, args.dim, args.is_causal, args.tune)
|
||||
@@ -190,7 +190,5 @@ python -m pytest "${PYTEST_ARGS_DEVICE[@]}" . "${PYTEST_ARGS_COMMON[@]}"
|
||||
cd .. || exit 1
|
||||
|
||||
# Run pytest in parallel for all tests in the testing/python directory.
|
||||
# CuTeDSL backend now defaults to GEMM v2 (WGMMA descriptor-based).
|
||||
# Set TILELANG_USE_GEMM_V1=1 only if you need to debug V1-specific issues.
|
||||
cd testing/python || exit 1
|
||||
python -m pytest "${PYTEST_ARGS_DEVICE[@]}" . "${PYTEST_ARGS_COMMON[@]}"
|
||||
|
||||
+7
-1
@@ -696,7 +696,13 @@ TVM_DLL const Op &tvm_rdna_wmma_store();
|
||||
/*!
|
||||
* \brief tilelang intrinsic for general matrix multiplication (GEMM).
|
||||
*
|
||||
* This op is used to represent a generic GEMM operation in tilelang.
|
||||
* This op wraps a templated `tl::gemm_*<...>` call into the generated device
|
||||
* code. Python-side lowering backends that want to delegate to the C++
|
||||
* template implementations in `src/tl_templates/<target>/gemm*.h` can emit a
|
||||
* call to this builtin directly via
|
||||
* T.call_intrin("handle", "tl.tl_gemm", op_instance_str, A_ptr, B_ptr,
|
||||
* C_ptr) where `op_instance_str` is the fully-instantiated `tl::gemm_ss<M, N,
|
||||
* K, ...>` template string.
|
||||
*/
|
||||
TVM_DLL const Op &tl_gemm();
|
||||
|
||||
|
||||
+98
-468
@@ -21,35 +21,21 @@ namespace tl {
|
||||
using namespace tir;
|
||||
|
||||
/**
|
||||
* @brief Construct a Gemm operator from serialized TL arguments and a buffer
|
||||
* map.
|
||||
* @brief Construct a Gemm operator from serialized TL arguments.
|
||||
*
|
||||
* This constructor deserializes operator parameters from `args` and resolves
|
||||
* buffer references via `vmap`, populating an internal GemmNode with:
|
||||
* - device pointers for A, B, C and their corresponding Buffer objects,
|
||||
* - transpose flags for A and B,
|
||||
* - matrix dimensions M, N, K,
|
||||
* - warp allocation policy and clear_accum flag,
|
||||
* - strides and memory offsets for A and B,
|
||||
* - optional kPack (must be 1 or 2) and optional internal wg_wait.
|
||||
*
|
||||
* The populated GemmNode is stored into the wrapper's internal `data_`.
|
||||
* Deserializes operator parameters from `args` and resolves buffer references,
|
||||
* populating an internal GemmNode with buffers, transpose flags, M/N/K,
|
||||
* warp policy, clear_accum, strides, offsets, optional kPack/wg_wait, and
|
||||
* optional mbarrier.
|
||||
*
|
||||
* @param args Positional serialized arguments produced by the TL frontend:
|
||||
* expected layout is:
|
||||
* [Aptr, Bptr, Cptr, trans_A (Bool), trans_B (Bool),
|
||||
* M (Int), N (Int), K (Int), policy (Int), clear_accum (Bool),
|
||||
* stride_A (Int), stride_B (Int), offset_A (Int), offset_B (Int),
|
||||
* (optional) kPack (Int), (optional) internal wg_wait (Int)]
|
||||
*
|
||||
* @note If `kPack` is provided it must be 1; otherwise the constructor
|
||||
* fails with an ICHECK (runtime assertion). No other validation is
|
||||
* performed here.
|
||||
* (optional) kPack (Int), (optional) internal wg_wait (Int),
|
||||
* (optional) mbar (BufferLoad), cCoord_y (PrimExpr), cCoord_x (PrimExpr)]
|
||||
*/
|
||||
// NormalizeToBufferRegion moved to src/op/utils.{h,cc}
|
||||
|
||||
// MakeAccessPtrFromRegion moved to src/op/utils.{h,cc}
|
||||
|
||||
Gemm::Gemm(Array<PrimExpr> args, Map<String, ObjectRef> annotations) {
|
||||
ObjectPtr<GemmNode> node = tvm::ffi::make_object<GemmNode>();
|
||||
|
||||
@@ -115,14 +101,6 @@ AccessRegions GemmNode::GetAccessRegions() const {
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Create a copy of this GemmNode as a TileOperator.
|
||||
*
|
||||
* Constructs a new GemmNode by copying the current node state and returns it
|
||||
* wrapped in a Gemm TileOperator.
|
||||
*
|
||||
* @return TileOperator A Gemm operator that owns a copy of this node.
|
||||
*/
|
||||
TileOperator GemmNode::Clone() const {
|
||||
auto op = tvm::ffi::make_object<GemmNode>(*this);
|
||||
return Gemm(op);
|
||||
@@ -133,6 +111,7 @@ bool GemmNode::allowTcgen5Mma(Target target) const {
|
||||
IsSharedBuffer(b_) && c_.scope() == "shared.tmem";
|
||||
if (!TargetIsSm100(target) || !scope_ok)
|
||||
return false;
|
||||
// For TS variant (A from TMEM), use B's dtype as the input dtype
|
||||
DataType ab_dtype = (a_.scope() == "shared.tmem") ? b_->dtype : a_->dtype;
|
||||
return GetTCGEN5MMAMeta(m_, n_, k_, ab_dtype, c_->dtype).first;
|
||||
}
|
||||
@@ -172,9 +151,11 @@ GemmInst GemmNode::getGemmInst(int block_size, Target target) const {
|
||||
}
|
||||
return GemmInst::kTCGEN5MMA;
|
||||
}
|
||||
if (allowTcgen5Mma(target)) {
|
||||
bool allow_tcgen5mma = allowTcgen5Mma(target);
|
||||
bool allow_wgmma = allowWgmma(block_size, target);
|
||||
if (allow_tcgen5mma) {
|
||||
return GemmInst::kTCGEN5MMA;
|
||||
} else if (allowWgmma(block_size, target)) {
|
||||
} else if (allow_wgmma) {
|
||||
return GemmInst::kWGMMA;
|
||||
} else if (TargetIsCDNA(target)) {
|
||||
return GemmInst::kMFMA;
|
||||
@@ -182,8 +163,10 @@ GemmInst GemmNode::getGemmInst(int block_size, Target target) const {
|
||||
return GemmInst::kWMMA;
|
||||
} else if (TargetIsCuda(target)) {
|
||||
return GemmInst::kMMA;
|
||||
} else if (TargetIsCPU(target)) {
|
||||
return GemmInst::kScalar;
|
||||
} else {
|
||||
ICHECK(0) << "Unsupported target for gemm: " << target;
|
||||
ICHECK(0) << "Unsupported target for gemm: " << target->str();
|
||||
return GemmInst::kMMA;
|
||||
}
|
||||
}
|
||||
@@ -379,32 +362,8 @@ std::pair<int, int> GemmWarpPolicyNode::computeWarpPartition(
|
||||
/**
|
||||
* @brief Checks whether WGMMA (warp-group MMA) can be used for this GEMM.
|
||||
*
|
||||
* Evaluates device-memory placement, data-type combinations, transpose flags,
|
||||
* and K divisibility constraints required for the Hopper WGMMA code path.
|
||||
*
|
||||
* The check returns true only when:
|
||||
* - B resides in shared memory ("shared" or "shared.dyn"); and
|
||||
* - (C, A, B) dtypes match one of the supported combinations below and K
|
||||
* satisfies the required alignment; and
|
||||
* - for combinations that require specific orientations, A is not transposed
|
||||
* and B is transposed.
|
||||
*
|
||||
* Supported combinations and constraints:
|
||||
* - C=float16:
|
||||
* - A=float16, B=float16: K % 16 == 0
|
||||
* - Various float8 mixes (e4m3/e5m2): require (!trans_A && trans_B) and K %
|
||||
* 32 == 0
|
||||
* - C=float32:
|
||||
* - A=float16, B=float16: K % 16 == 0
|
||||
* - A=bfloat16, B=bfloat16: K % 16 == 0
|
||||
* - A=float32, B=float32: require (!trans_A && trans_B) and K % 8 == 0
|
||||
* - Various float8 mixes: require (!trans_A && trans_B) and K % 32 == 0
|
||||
* - C=int32:
|
||||
* - 8-bit integer combinations (Int8/UInt8): require (!trans_A && trans_B)
|
||||
* and K % 32 == 0
|
||||
*
|
||||
* @return true if WGMMA is supported for the current buffers, dtypes, and
|
||||
* transpose/shape constraints; false otherwise.
|
||||
* Returns true only when B resides in shared memory and the (C, A, B) dtype
|
||||
* combination plus K alignment matches one of the supported WGMMA variants.
|
||||
*/
|
||||
bool GemmNode::checkWgmma() const {
|
||||
if (b_.scope() != "shared.dyn" && b_.scope() != "shared") {
|
||||
@@ -447,436 +406,71 @@ bool GemmNode::checkWgmma() const {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Parse and return the numeric GPU architecture from a Target's "arch"
|
||||
* attribute.
|
||||
*
|
||||
* Examines the target's "arch" string and, if it matches the pattern
|
||||
* "sm_<num>", returns <num> as an int. If the attribute is present but does not
|
||||
* match that pattern, returns 0.
|
||||
*
|
||||
* Preconditions: the target must have an "arch" attribute (this is checked via
|
||||
* ICHECK).
|
||||
*
|
||||
* @return int The parsed architecture number (e.g., 80 for "sm_80"), or 0 if
|
||||
* the arch string does not match "sm_<num>".
|
||||
*/
|
||||
static int GetArchInt(Target target) {
|
||||
int arch_int = 0;
|
||||
auto s = target->GetAttr<tvm::ffi::String>("arch");
|
||||
ICHECK(s.has_value());
|
||||
std::string arch = s.value();
|
||||
if (arch.rfind("sm_", 0) == 0) {
|
||||
arch_int = std::stoi(arch.substr(3));
|
||||
} else {
|
||||
arch_int = 0;
|
||||
}
|
||||
return arch_int;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Lower the GEMM operator to a TL TIR call expression.
|
||||
*
|
||||
* Constructs a tl::gemm call string parameterized by M, N, K, warp partition,
|
||||
* transpose flags, accumulation clearing, target-specific stride/offset/kPack
|
||||
* and optional workgroup wait value, then returns an Evaluate(call) node
|
||||
* invoking tl::tl_gemm with the composed string and the A/B/C buffer handles.
|
||||
*
|
||||
* @param T Contains lowering context including thread bounds and target.
|
||||
* @param analyzer Optional arithmetic analyzer used by lowering (may be
|
||||
* nullptr).
|
||||
* @return Stmt A TIR statement representing the evaluated TL GEMM call.
|
||||
*/
|
||||
Stmt GemmNode::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const {
|
||||
auto block_size = *as_const_int(T.thread_bounds->extent);
|
||||
GemmInst gemm_inst = getGemmInst(block_size, T.target);
|
||||
auto [warp_m, warp_n] =
|
||||
policy_->computeWarpPartition(m_, n_, block_size, T.target, gemm_inst);
|
||||
|
||||
// Build access pointers from regions locally
|
||||
PrimExpr Aptr =
|
||||
MakeAccessPtrFromRegion(aRegion_, /*r*/ 1, /*require_2d*/ true);
|
||||
PrimExpr Bptr =
|
||||
MakeAccessPtrFromRegion(bRegion_, /*r*/ 1, /*require_2d*/ true);
|
||||
PrimExpr Cptr =
|
||||
MakeAccessPtrFromRegion(cRegion_, /*rw*/ 3, /*require_2d*/ true);
|
||||
|
||||
std::stringstream ss;
|
||||
std::string op_name;
|
||||
|
||||
if (gemm_inst == GemmInst::kTCGEN5MMA) {
|
||||
DataType ab_dtype = (a_.scope() == "shared.tmem") ? b_->dtype : a_->dtype;
|
||||
auto [can_use_tcgen5mma, meta] =
|
||||
GetTCGEN5MMAMeta(m_, n_, k_, ab_dtype, c_->dtype);
|
||||
ICHECK(can_use_tcgen5mma);
|
||||
ICHECK(IsSharedBuffer(b_));
|
||||
ICHECK(c_.scope() == "shared.tmem");
|
||||
ICHECK(mbar_.defined()) << "mbar must be provided for TCGEN5MMA";
|
||||
if (a_.scope() == "shared.tmem") {
|
||||
op_name = "tl::tcgen5mma_gemm_ts";
|
||||
} else if (IsSharedBuffer(a_)) {
|
||||
op_name = "tl::tcgen5mma_gemm_ss";
|
||||
} else {
|
||||
ICHECK(0)
|
||||
<< "Unsupported A scope for TCGEN5MMA: "
|
||||
<< a_.scope(); // If this is triggered, it means Tilelang has bugs.
|
||||
}
|
||||
ICHECK(wgWait_ == 0 || wgWait_ == -1)
|
||||
<< "TCGEN5MMA only accepts internal wg_wait values 0 or -1. "
|
||||
"Public T.gemm() uses 0 and synchronization is still managed "
|
||||
"manually via mbarrier.";
|
||||
|
||||
std::string accum_dtype = "";
|
||||
if (c_->dtype.is_float()) {
|
||||
if (c_->dtype.bits() == 32) {
|
||||
accum_dtype = "float";
|
||||
}
|
||||
}
|
||||
ICHECK(!accum_dtype.empty())
|
||||
<< "Unsupported C dtype for TCGEN5MMA: " << c_->dtype;
|
||||
ss << op_name << "<" << m_ << ", " << n_ << ", " << k_ << ", ";
|
||||
ss << meta.atom_m << ", " << meta.atom_n << ", " << meta.atom_k << ", ";
|
||||
ss << transA_ << ", " << transB_ << ", ";
|
||||
ss << accum_dtype;
|
||||
ss << ">";
|
||||
|
||||
auto C_buffer = T.buffer_remap.count(c_) ? T.buffer_remap[c_] : c_;
|
||||
Array<PrimExpr> new_args;
|
||||
auto mbarPtr = MakeAccessPtrFromBufferLoad(mbar_, /*rw*/ 3);
|
||||
new_args.push_back(StringImm(ss.str()));
|
||||
new_args.push_back(Aptr);
|
||||
new_args.push_back(Bptr);
|
||||
new_args.push_back(BufferLoad(C_buffer, cCoords_));
|
||||
new_args.push_back(mbarPtr);
|
||||
new_args.push_back(clearAccum_);
|
||||
auto new_call = Call(DataType::Handle(), builtin::call_extern(), new_args);
|
||||
|
||||
// Since TCGEN5MMA atoms provided by CUTLASS always have an internal
|
||||
// `elect_one_sync()`, we check if we are calling it using full warps
|
||||
constexpr int warp_size = 32;
|
||||
ICHECK(
|
||||
analyzer->CanProveEqual(FloorMod(T.thread_bounds->min, warp_size), 0) &&
|
||||
analyzer->CanProveEqual(FloorMod(T.thread_bounds->extent, warp_size),
|
||||
0))
|
||||
<< "TCGEN5MMA requires thread bounds to be multiples of warp size (32) "
|
||||
"and aligned to warps.";
|
||||
Stmt tcgen5mma_call;
|
||||
if (analyzer->CanProveEqual(T.thread_bounds->extent, warp_size)) {
|
||||
// If the thread bounds is exactly one warp, we can use the original call
|
||||
tcgen5mma_call = Evaluate(new_call);
|
||||
} else {
|
||||
// Add an if-else clause
|
||||
tcgen5mma_call = IfThenElse(EQ(FloorDiv(T.thread_var, warp_size),
|
||||
FloorDiv(T.thread_bounds->min, warp_size)),
|
||||
Evaluate(new_call));
|
||||
}
|
||||
if (isTcgen05_) {
|
||||
return tcgen5mma_call;
|
||||
}
|
||||
if (const auto f = ffi::Function::GetGlobal("tl.gemm.lower")) {
|
||||
PrimExpr mbar_phase = T.mbar_phase_expr;
|
||||
if (auto explicit_phase = GetAnnotatedMbarPhaseExpr(annotations_)) {
|
||||
mbar_phase = explicit_phase.value();
|
||||
}
|
||||
Stmt wait_stmt = Evaluate(
|
||||
Call(DataType::Handle(), mbarrier_wait_parity(), {mbar_, mbar_phase}));
|
||||
return SeqStmt({tcgen5mma_call, wait_stmt});
|
||||
}
|
||||
|
||||
if (IsFragmentBuffer(a_)) {
|
||||
ICHECK(!IsFragmentBuffer(b_));
|
||||
ICHECK(!transA_)
|
||||
<< "gemm_rs requires the A operand to be in non-transposed layout.";
|
||||
op_name = "tl::gemm_rs";
|
||||
} else if (IsFragmentBuffer(b_)) {
|
||||
op_name = "tl::gemm_sr";
|
||||
} else {
|
||||
op_name = "tl::gemm_ss";
|
||||
}
|
||||
ICHECK(IsFragmentBuffer(c_));
|
||||
|
||||
ss << op_name << "<" << m_ << ", " << n_ << ", " << k_ << ", ";
|
||||
ss << warp_m << ", " << warp_n << ", ";
|
||||
ss << transA_ << ", " << transB_;
|
||||
auto clear_accum_bool = clearAccum_.as<Bool>();
|
||||
ICHECK(clear_accum_bool.has_value())
|
||||
<< "clear_accum must be a constant Bool type, got " << clearAccum_;
|
||||
ss << ", " << bool(clear_accum_bool.value());
|
||||
if (TargetIsCuda(T.target) && (GetArchInt(T.target) >= 75)) {
|
||||
ss << ", " << strideA_ << ", " << strideB_;
|
||||
ss << ", " << offsetA_ << ", " << offsetB_;
|
||||
}
|
||||
if (TargetIsCDNA(T.target)) {
|
||||
// for cdna gemm, we need to specify kPack
|
||||
ss << ", " << kPack_;
|
||||
} else if (TargetIsHopper(T.target)) {
|
||||
ss << ", " << (gemm_inst == GemmInst::kWGMMA ? "true" : "false");
|
||||
}
|
||||
|
||||
// Emit internal wg_wait only for Hopper WGMMA lowering.
|
||||
if (TargetIsHopper(T.target)) {
|
||||
if (wgWait_ != 0) {
|
||||
ss << ", " << wgWait_;
|
||||
// NOTE(wt): Decide GemmInst and compute warp partition on Python side
|
||||
auto prim_func = Downcast<PrimFunc>(
|
||||
(*f)(tvm::ffi::GetRef<Gemm>(this), T.layout_map, T.target,
|
||||
T.thread_bounds, T.thread_var, mbar_phase));
|
||||
ICHECK(prim_func->attrs.defined());
|
||||
auto global_symbol =
|
||||
prim_func->attrs.GetAttr<tvm::ffi::String>("global_symbol");
|
||||
ICHECK(global_symbol.has_value());
|
||||
if (prim_func->body.as<BlockRealizeNode>()) {
|
||||
BlockRealize block_realize = Downcast<BlockRealize>(prim_func->body);
|
||||
auto block = block_realize->block;
|
||||
{
|
||||
BlockNode *n = block.CopyOnWrite();
|
||||
n->name_hint = global_symbol.value();
|
||||
n->annotations.Set(tl::attr::kLexicalAllocScope,
|
||||
IntImm(DataType::Int(32), 1));
|
||||
}
|
||||
return BlockRealize(block_realize->iter_values, block_realize->predicate,
|
||||
block);
|
||||
}
|
||||
} else if (TargetIsSm100(T.target)) {
|
||||
// NOTE On sm100, only the leading thread issues the TCGEN5MMA instruction
|
||||
// but all threads need to wait, so we emit another statement for cases
|
||||
// where wg_wait == 0.
|
||||
ICHECK(wgWait_ == 0 || wgWait_ == -1)
|
||||
<< "wg_wait must be 0 or -1 for Sm100";
|
||||
// wrap with block realize node
|
||||
Map<String, ObjectRef> block_annotations;
|
||||
block_annotations.Set(tl::attr::kLexicalAllocScope,
|
||||
IntImm(DataType::Int(32), 1));
|
||||
return BlockRealize(
|
||||
/*iter_values=*/Array<PrimExpr>(),
|
||||
/*predicate=*/const_true(),
|
||||
/*block=*/
|
||||
Block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{},
|
||||
/*name_hint=*/global_symbol.value(), prim_func->body,
|
||||
/*init=*/Optional<Stmt>(), /*alloc_buffers=*/{},
|
||||
/*match_buffers=*/{}, /*annotations=*/block_annotations));
|
||||
} else {
|
||||
ICHECK(wgWait_ == 0)
|
||||
<< "wg_wait must be 0 for non-Hopper and non-Sm100 targets";
|
||||
LOG(FATAL) << "No lower function found for gemm";
|
||||
return Stmt();
|
||||
}
|
||||
ss << ">";
|
||||
|
||||
auto new_call = Call(DataType::Handle(), tl::tl_gemm(),
|
||||
Array<PrimExpr>{StringImm(ss.str()), Aptr, Bptr, Cptr});
|
||||
return Evaluate(new_call);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Infer and bind target-specific memory/layout mappings for A, B, and C.
|
||||
*
|
||||
* Infers per-buffer layouts (fragment or shared-memory layouts) for this GEMM
|
||||
* operator according to the target architecture, thread bounds, warp
|
||||
* partitioning, data types, and transpose flags, then binds fragment layouts
|
||||
* to the thread range when required.
|
||||
*
|
||||
* Preconditions:
|
||||
* - C.scope() == "local.fragment"
|
||||
*
|
||||
* Side effects:
|
||||
* - Marks layout inference as completed (sets completed_ = true).
|
||||
* - May abort via ICHECK on unsupported targets, invalid buffer scopes, or
|
||||
* incompatible shape constraints.
|
||||
*
|
||||
* @param T Input layout-inference context (provides thread bounds and target).
|
||||
* @return LayoutMap mapping A, B, and C to their inferred layouts.
|
||||
*/
|
||||
LayoutMap GemmNode::InferLayout(const LayoutInferArgs &T,
|
||||
InferLevel level) const {
|
||||
if (completed_)
|
||||
return {};
|
||||
LayoutMap results;
|
||||
auto thread_range = T.thread_bounds;
|
||||
auto block_size = *as_const_int(thread_range->extent);
|
||||
GemmInst gemm_inst = getGemmInst(block_size, T.target);
|
||||
auto [warp_m, warp_n] =
|
||||
policy_->computeWarpPartition(m_, n_, block_size, T.target, gemm_inst);
|
||||
if (TargetIsVolta(T.target)) {
|
||||
ICHECK(IsFragmentBuffer(c_))
|
||||
<< "Volta gemm only supports C in local.fragment scope, got "
|
||||
<< c_.scope();
|
||||
auto fragment = makeGemmVoltaFragmentC(m_, n_, m_ / warp_m, n_ / warp_n,
|
||||
c_->dtype.bits());
|
||||
results.Set(c_, fragment->BindThreadRange(thread_range));
|
||||
if (IsSharedBuffer(a_)) {
|
||||
int dim_A = a_->shape.size();
|
||||
auto layout = makeGemmVoltaABLayout(*as_const_int(a_->shape[dim_A - 2]),
|
||||
*as_const_int(a_->shape[dim_A - 1]),
|
||||
true, !transA_);
|
||||
results.Set(a_, ExpandLayoutToMatchBuffer(layout, a_));
|
||||
} else if (IsFragmentBuffer(a_)) {
|
||||
ICHECK(transA_ == false);
|
||||
auto fragment =
|
||||
makeGemmVoltaFragmentA(m_, n_, k_, m_ / warp_m, n_ / warp_n);
|
||||
results.Set(a_, fragment->BindThreadRange(thread_range));
|
||||
} else {
|
||||
ICHECK(0);
|
||||
}
|
||||
|
||||
ICHECK(IsSharedBuffer(b_));
|
||||
int dim_B = b_->shape.size();
|
||||
auto layout = makeGemmVoltaABLayout(*as_const_int(b_->shape[dim_B - 2]),
|
||||
*as_const_int(b_->shape[dim_B - 1]),
|
||||
false, transB_);
|
||||
results.Set(b_, ExpandLayoutToMatchBuffer(layout, b_));
|
||||
} else if (TargetIsAmpere(T.target) || TargetIsTuring(T.target) ||
|
||||
TargetIsSM120(T.target) ||
|
||||
(TargetIsSm100(T.target) && gemm_inst == GemmInst::kMMA)) {
|
||||
ICHECK(IsFragmentBuffer(c_))
|
||||
<< "MMA only supports C in local.fragment scope, got " << c_.scope();
|
||||
|
||||
auto fragment =
|
||||
makeGemmFragmentC(m_, n_, m_ / warp_m, n_ / warp_n, c_->dtype.bits());
|
||||
results.Set(c_, fragment->BindThreadRange(thread_range));
|
||||
|
||||
if (IsSharedBuffer(a_)) {
|
||||
int dim_A = a_->shape.size();
|
||||
const int64_t mat_stride = *as_const_int(a_->shape[dim_A - 2]);
|
||||
const int64_t mat_continuous = *as_const_int(a_->shape[dim_A - 1]);
|
||||
auto layout = makeGemmABLayout(mat_stride, mat_continuous, mat_continuous,
|
||||
a_->dtype.bits(), !transA_);
|
||||
results.Set(a_, ExpandLayoutToMatchBuffer(layout, a_));
|
||||
} else if (IsFragmentBuffer(a_)) {
|
||||
auto fragment = makeGemmFragmentA(m_, n_, k_, m_ / warp_m, n_ / warp_n,
|
||||
a_->dtype.bits(), transA_);
|
||||
results.Set(a_, fragment->BindThreadRange(thread_range));
|
||||
} else {
|
||||
ICHECK(0);
|
||||
}
|
||||
if (IsSharedBuffer(b_)) {
|
||||
int dim_B = b_->shape.size();
|
||||
const int64_t mat_stride = *as_const_int(b_->shape[dim_B - 2]);
|
||||
const int64_t mat_continuous = *as_const_int(b_->shape[dim_B - 1]);
|
||||
auto layout = makeGemmABLayout(mat_stride, mat_continuous, mat_continuous,
|
||||
b_->dtype.bits(), transB_);
|
||||
results.Set(b_, ExpandLayoutToMatchBuffer(layout, b_));
|
||||
} else if (IsFragmentBuffer(b_)) {
|
||||
auto fragment =
|
||||
makeGemmFragmentB(m_, n_, k_, m_ / warp_m, n_ / warp_n, transB_);
|
||||
results.Set(b_, fragment->BindThreadRange(thread_range));
|
||||
} else {
|
||||
ICHECK(0);
|
||||
}
|
||||
} else if (TargetIsHopper(T.target)) {
|
||||
ICHECK(IsFragmentBuffer(c_))
|
||||
<< (gemm_inst == GemmInst::kWGMMA ? "WGMMA " : "MMA ")
|
||||
<< "only supports C in local.fragment scope, got " << c_.scope();
|
||||
auto fragment = gemm_inst == GemmInst::kWGMMA
|
||||
? makeGemmFragmentCHopper(m_, n_, m_ / warp_m,
|
||||
n_ / warp_n, c_->dtype.bits())
|
||||
: makeGemmFragmentC(m_, n_, m_ / warp_m, n_ / warp_n,
|
||||
c_->dtype.bits());
|
||||
results.Set(c_, fragment->BindThreadRange(thread_range));
|
||||
if (IsSharedBuffer(a_)) {
|
||||
int dim_A = a_->shape.size();
|
||||
const int64_t mat_stride = *as_const_int(a_->shape[dim_A - 2]);
|
||||
const int64_t mat_continuous = *as_const_int(a_->shape[dim_A - 1]);
|
||||
const int64_t continuity =
|
||||
transA_ ? 4 * mat_continuous / warp_m : mat_continuous;
|
||||
auto ABLayout =
|
||||
gemm_inst == GemmInst::kWGMMA
|
||||
? makeGemmABLayoutHopper(mat_stride, mat_continuous, continuity,
|
||||
a_->dtype.bits(), !transA_)
|
||||
: makeGemmABLayout(mat_stride, mat_continuous, mat_continuous,
|
||||
a_->dtype.bits(), !transA_);
|
||||
results.Set(a_, ExpandLayoutToMatchBuffer(ABLayout, a_));
|
||||
} else {
|
||||
auto fragment = makeGemmFragmentA(m_, n_, k_, m_ / warp_m, n_ / warp_n,
|
||||
a_->dtype.bits(), transA_);
|
||||
results.Set(a_, fragment->BindThreadRange(thread_range));
|
||||
}
|
||||
if (IsSharedBuffer(b_)) {
|
||||
int dim_B = b_->shape.size();
|
||||
const int64_t mat_stride = *as_const_int(b_->shape[dim_B - 2]);
|
||||
const int64_t mat_continuous = *as_const_int(b_->shape[dim_B - 1]);
|
||||
const int64_t continuity =
|
||||
transB_ ? mat_continuous : mat_continuous / warp_n;
|
||||
|
||||
auto ABLayout =
|
||||
gemm_inst == GemmInst::kWGMMA
|
||||
? makeGemmABLayoutHopper(mat_stride, mat_continuous, continuity,
|
||||
b_->dtype.bits(), transB_)
|
||||
: makeGemmABLayout(mat_stride, mat_continuous, mat_continuous,
|
||||
b_->dtype.bits(), transB_);
|
||||
results.Set(b_, ExpandLayoutToMatchBuffer(ABLayout, b_));
|
||||
} else {
|
||||
auto fragment =
|
||||
makeGemmFragmentB(m_, n_, k_, m_ / warp_m, n_ / warp_n, transB_);
|
||||
results.Set(b_, fragment->BindThreadRange(thread_range));
|
||||
}
|
||||
} else if (gemm_inst == GemmInst::kTCGEN5MMA) {
|
||||
ICHECK(c_.scope() == "shared.tmem")
|
||||
<< "TCGEN5MMA only supports C in shared.tmem scope, got " << c_.scope();
|
||||
ICHECK(IsSharedBuffer(a_))
|
||||
<< "Current TCGEN5MMA only supports A in shared.dyn scope";
|
||||
DataType ab_dtype = (a_.scope() == "shared.tmem") ? b_->dtype : a_->dtype;
|
||||
auto [can_use_tcgen5mma, meta] =
|
||||
GetTCGEN5MMAMeta(m_, n_, k_, ab_dtype, c_->dtype);
|
||||
ICHECK(can_use_tcgen5mma);
|
||||
{
|
||||
int dim_A = a_->shape.size();
|
||||
const int64_t mat_stride = *as_const_int(a_->shape[dim_A - 2]);
|
||||
const int64_t mat_continuous = *as_const_int(a_->shape[dim_A - 1]);
|
||||
auto layout =
|
||||
makeGemmABLayoutSm100(mat_stride, mat_continuous, mat_continuous,
|
||||
a_->dtype.bits(), transA_ ? 1 : 2);
|
||||
results.Set(a_, ExpandLayoutToMatchBuffer(layout, a_));
|
||||
}
|
||||
{
|
||||
int dim_B = b_->shape.size();
|
||||
const int64_t mat_stride = *as_const_int(b_->shape[dim_B - 2]);
|
||||
const int64_t mat_continuous = *as_const_int(b_->shape[dim_B - 1]);
|
||||
const int64_t continuity = mat_continuous;
|
||||
auto layout =
|
||||
makeGemmABLayoutSm100(mat_stride, mat_continuous, continuity,
|
||||
b_->dtype.bits(), transB_ ? 2 : 1);
|
||||
results.Set(b_, ExpandLayoutToMatchBuffer(layout, b_));
|
||||
}
|
||||
{
|
||||
Layout res;
|
||||
IterVar i = make_itervar("i", m_);
|
||||
IterVar j = make_itervar("j", n_);
|
||||
ICHECK(m_ % meta.atom_m == 0);
|
||||
PrimExpr atom_idx = FloorDiv(i, meta.atom_m) +
|
||||
FloorDiv(j, meta.atom_n) * (m_ / meta.atom_m);
|
||||
PrimExpr ai = FloorMod(i, meta.atom_m); // "ai" means "atom_i"
|
||||
PrimExpr aj = FloorMod(j, meta.atom_n);
|
||||
if (meta.atom_m == 128) {
|
||||
// Layout D
|
||||
// (https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-d)
|
||||
res = Layout(Array{i, j}, {ai, aj + atom_idx * meta.atom_n});
|
||||
} else if (meta.atom_m == 64) {
|
||||
// Layout E
|
||||
// (https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-e)
|
||||
// since .ws variant is used About why we use .ws variant here, please
|
||||
// refer to gemm_sm100.h
|
||||
res = Layout(Array{i, j}, {FloorDiv(ai, 32) * 32 + FloorMod(ai, 32) +
|
||||
FloorDiv(aj, meta.atom_n / 2) * 64,
|
||||
FloorMod(aj, meta.atom_n / 2) +
|
||||
atom_idx * (meta.atom_n / 2)});
|
||||
} else if (meta.atom_m == 32) {
|
||||
// Layout G
|
||||
// (https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-g)
|
||||
res = Layout(
|
||||
Array{i, j},
|
||||
{FloorMod(ai, 32) + FloorDiv(aj, meta.atom_n / 4) * 32,
|
||||
FloorMod(aj, meta.atom_n / 4) + atom_idx * (meta.atom_n / 4)});
|
||||
} else {
|
||||
ICHECK(0);
|
||||
if (const auto f = ffi::Function::GetGlobal("tl.gemm.infer_layout")) {
|
||||
results = Downcast<LayoutMap>(
|
||||
(*f)(tvm::ffi::GetRef<Gemm>(this), T.target, T.thread_bounds));
|
||||
// Bind all fragment layouts with the provided thread range
|
||||
for (auto kv : results) {
|
||||
const Buffer &buf = kv.first;
|
||||
const Layout &layout = kv.second;
|
||||
if (auto frag = layout.as<Fragment>()) {
|
||||
results.Set(buf, frag.value()->BindThreadRange(T.thread_bounds));
|
||||
}
|
||||
results.Set(c_, res);
|
||||
}
|
||||
} else if (TargetIsCDNA(T.target)) {
|
||||
ICHECK(IsFragmentBuffer(c_))
|
||||
<< "CDNA gemm (FMMA) only supports C in local.fragment scope, got "
|
||||
<< c_.scope();
|
||||
auto fragment = makeGemmFragmentCCDNA(m_, n_, m_ / warp_m, n_ / warp_n,
|
||||
c_->dtype.bits());
|
||||
results.Set(c_, fragment->BindThreadRange(thread_range));
|
||||
|
||||
if (IsSharedBuffer(a_)) {
|
||||
int dim_A = a_->shape.size();
|
||||
auto shared_layout = makeGemmABLayoutCDNA(
|
||||
*as_const_int(a_->shape[dim_A - 2]),
|
||||
*as_const_int(a_->shape[dim_A - 1]), a_->dtype.bits(), kPack_);
|
||||
results.Set(a_, ExpandLayoutToMatchBuffer(shared_layout, a_));
|
||||
} else if (IsFragmentBuffer(a_)) {
|
||||
auto fragment =
|
||||
makeGemmFragmentACDNA(m_, n_, k_, m_ / warp_m, n_ / warp_n,
|
||||
a_->dtype.bits(), kPack_, transA_);
|
||||
results.Set(a_, fragment->BindThreadRange(thread_range));
|
||||
} else {
|
||||
ICHECK(0);
|
||||
}
|
||||
if (IsSharedBuffer(b_)) {
|
||||
int dim_B = b_->shape.size();
|
||||
auto shared_layout = makeGemmABLayoutCDNA(
|
||||
*as_const_int(b_->shape[dim_B - 2]),
|
||||
*as_const_int(b_->shape[dim_B - 1]), b_->dtype.bits(), kPack_);
|
||||
|
||||
results.Set(b_, ExpandLayoutToMatchBuffer(shared_layout, b_));
|
||||
} else if (IsFragmentBuffer(b_)) {
|
||||
auto fragment =
|
||||
makeGemmFragmentB(m_, n_, k_, m_ / warp_m, n_ / warp_n, transB_);
|
||||
results.Set(b_, fragment->BindThreadRange(thread_range));
|
||||
} else {
|
||||
ICHECK(0);
|
||||
}
|
||||
} else {
|
||||
ICHECK(0) << "Not supported " << T.target->str();
|
||||
LOG(FATAL) << "No infer layout function found for gemm";
|
||||
}
|
||||
|
||||
completed_ = true;
|
||||
return results;
|
||||
}
|
||||
@@ -920,6 +514,9 @@ TVM_REGISTER_OP("tl.GemmWarpPolicy")
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
GemmNode::RegisterReflection();
|
||||
GemmWarpPolicyNode::RegisterReflection();
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("tl.GemmWarpPolicyComputeWarpPartition",
|
||||
[](GemmWarpPolicy policy, int M, int N, int block_size,
|
||||
@@ -927,6 +524,39 @@ TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
policy->computeWarpPartition(M, N, block_size, target,
|
||||
gemm_inst);
|
||||
});
|
||||
refl::GlobalDef().def("tl.GemmGetGemmInst",
|
||||
[](Gemm gemm, int block_size, Target target) {
|
||||
return gemm->getGemmInst(block_size, target);
|
||||
});
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def(
|
||||
"tl.get_tcgen5_mma_meta", [](int M, int N, int K, DataType ab_dtype,
|
||||
DataType c_dtype, bool disable_2cta) {
|
||||
auto [success, meta] =
|
||||
GetTCGEN5MMAMeta(M, N, K, ab_dtype, c_dtype, disable_2cta);
|
||||
Array<Integer> result;
|
||||
if (success) {
|
||||
result.push_back(Integer(meta.atom_m));
|
||||
result.push_back(Integer(meta.atom_n));
|
||||
result.push_back(Integer(meta.atom_k));
|
||||
result.push_back(Integer(meta.enable_ws));
|
||||
result.push_back(Integer(meta.enable_2cta));
|
||||
}
|
||||
return result;
|
||||
});
|
||||
refl::GlobalDef().def(
|
||||
"tl.get_tcgen5_instr_desc",
|
||||
[](int atom_m, int atom_n, int atom_k, DataType ab_dtype,
|
||||
DataType c_dtype, bool a_is_k_major, bool b_is_k_major, int scale_in_a,
|
||||
int scale_in_b) {
|
||||
uint32_t desc = GetTCGEN5InstrDesc(atom_m, atom_n, atom_k, ab_dtype,
|
||||
c_dtype, a_is_k_major, b_is_k_major,
|
||||
scale_in_a, scale_in_b);
|
||||
return Integer(static_cast<int64_t>(desc));
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace tl
|
||||
|
||||
+8
-7
@@ -129,6 +129,8 @@ public:
|
||||
class GemmNode : public TileOperatorNode {
|
||||
public:
|
||||
bool checkWgmma() const;
|
||||
bool allowTcgen5Mma(Target target) const;
|
||||
bool allowWgmma(int block_size, Target target) const;
|
||||
tir::Buffer a_, b_, c_;
|
||||
// BufferRegion for A, B and C
|
||||
BufferRegion aRegion_, bRegion_, cRegion_;
|
||||
@@ -137,16 +139,17 @@ public:
|
||||
int strideA_, strideB_;
|
||||
int offsetA_, offsetB_;
|
||||
PrimExpr clearAccum_ = const_false();
|
||||
tir::BufferLoad mbar_; // mbar is optional, only used for TCGEN5MMA
|
||||
Array<PrimExpr> cCoords_;
|
||||
// k_pack please ref to bitblas/tl/mfma_macro_generator.py::k_pack
|
||||
// only will be enabled under cdna mfma instructions
|
||||
int kPack_ = 1;
|
||||
int wgWait_ = 0;
|
||||
bool isWgmma_ = false;
|
||||
bool isTcgen05_ = false;
|
||||
tir::BufferLoad mbar_; // mbar is optional, only used for TCGEN5MMA
|
||||
Array<PrimExpr> cCoords_;
|
||||
mutable GemmWarpPolicy policy_;
|
||||
Map<String, ObjectRef> annotations_;
|
||||
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tl.Gemm", GemmNode, TileOperatorNode);
|
||||
|
||||
static void RegisterReflection() {
|
||||
@@ -168,12 +171,12 @@ public:
|
||||
.def_ro("offsetA", &GemmNode::offsetA_)
|
||||
.def_ro("offsetB", &GemmNode::offsetB_)
|
||||
.def_ro("clearAccum", &GemmNode::clearAccum_)
|
||||
.def_ro("mbar", &GemmNode::mbar_)
|
||||
.def_ro("cCoords", &GemmNode::cCoords_)
|
||||
.def_ro("kPack", &GemmNode::kPack_)
|
||||
.def_ro("wgWait", &GemmNode::wgWait_)
|
||||
.def_ro("isWgmma", &GemmNode::isWgmma_)
|
||||
.def_ro("isTcgen05", &GemmNode::isTcgen05_)
|
||||
.def_ro("mbar", &GemmNode::mbar_)
|
||||
.def_ro("cCoords", &GemmNode::cCoords_)
|
||||
.def_ro("policy", &GemmNode::policy_)
|
||||
.def_ro("annotations", &GemmNode::annotations_);
|
||||
}
|
||||
@@ -185,12 +188,10 @@ public:
|
||||
|
||||
TileOperator Clone() const;
|
||||
|
||||
// Target GEMM instruction
|
||||
GemmInst getGemmInst(int block_size, Target target) const;
|
||||
|
||||
private:
|
||||
bool allowTcgen5Mma(Target target) const;
|
||||
bool allowWgmma(int block_size, Target target) const;
|
||||
|
||||
mutable bool completed_ = false;
|
||||
};
|
||||
|
||||
|
||||
@@ -1,438 +0,0 @@
|
||||
/*!
|
||||
* \file tl/op/gemm_py.cc
|
||||
* \brief Implementation of General Matrix Multiplication (GEMM) operators
|
||||
*/
|
||||
|
||||
#include "gemm_py.h"
|
||||
|
||||
#include "builtin.h"
|
||||
#include <tvm/tir/builtin.h>
|
||||
#include <tvm/tir/op.h>
|
||||
#include <tvm/tir/op_attr_types.h>
|
||||
#include <tvm/tir/transform.h>
|
||||
|
||||
#include "../target/utils.h"
|
||||
#include "tcgen5_meta.h"
|
||||
#include "utils.h"
|
||||
|
||||
namespace tvm {
|
||||
namespace tl {
|
||||
|
||||
using namespace tir;
|
||||
|
||||
// NormalizeToBufferRegion moved to src/op/utils.{h,cc}
|
||||
|
||||
// MakeAccessPtrFromRegion moved to src/op/utils.{h,cc}
|
||||
|
||||
/**
|
||||
* @brief Construct a Gemm operator from serialized TL arguments and a buffer
|
||||
* map.
|
||||
*
|
||||
* This constructor deserializes operator parameters from `args` and resolves
|
||||
* buffer references via `vmap`, populating an internal GemmPyNode with:
|
||||
* - device pointers for A, B, C and their corresponding Buffer objects,
|
||||
* - transpose flags for A and B,
|
||||
* - matrix dimensions M, N, K,
|
||||
* - warp allocation policy and clear_accum flag,
|
||||
* - strides and memory offsets for A and B,
|
||||
* - optional kPack (must be 1 or 2) and optional internal wg_wait.
|
||||
*
|
||||
* The populated GemmPyNode is stored into the wrapper's internal `data_`.
|
||||
*
|
||||
* @param args Positional serialized arguments produced by the TL frontend:
|
||||
* expected layout is:
|
||||
* [Aptr, Bptr, Cptr, trans_A (Bool), trans_B (Bool),
|
||||
* M (Int), N (Int), K (Int), policy (Int), clear_accum (Bool),
|
||||
* stride_A (Int), stride_B (Int), offset_A (Int), offset_B (Int),
|
||||
* (optional) kPack (Int), (optional) internal wg_wait (Int)]
|
||||
*
|
||||
* @note If `kPack` is provided it must be 1 or 2; otherwise the constructor
|
||||
* fails with an ICHECK (runtime assertion). No other validation is
|
||||
* performed here.
|
||||
*/
|
||||
GemmPy::GemmPy(Array<PrimExpr> args, Map<String, ObjectRef> annotations) {
|
||||
ObjectPtr<GemmPyNode> node = tvm::ffi::make_object<GemmPyNode>();
|
||||
|
||||
auto a_access = NormalizeToAccessRegion(args[0], kAccessRead);
|
||||
auto b_access = NormalizeToAccessRegion(args[1], kAccessRead);
|
||||
auto c_access = NormalizeToAccessRegion(args[2], kAccessReadWrite);
|
||||
|
||||
node->aRegion_ = a_access.region;
|
||||
node->bRegion_ = b_access.region;
|
||||
node->cRegion_ = c_access.region;
|
||||
node->SetAccessRegions({a_access, b_access, c_access});
|
||||
|
||||
node->a_ = node->aRegion_->buffer;
|
||||
node->b_ = node->bRegion_->buffer;
|
||||
node->c_ = node->cRegion_->buffer;
|
||||
node->transA_ = args[3].as<Bool>().value();
|
||||
node->transB_ = args[4].as<Bool>().value();
|
||||
node->m_ = args[5].as<IntImm>().value()->value;
|
||||
node->n_ = args[6].as<IntImm>().value()->value;
|
||||
node->k_ = args[7].as<IntImm>().value()->value;
|
||||
node->policy_ = GemmWarpPolicy(args[8].as<IntImm>().value()->value);
|
||||
node->clearAccum_ = args[9].as<PrimExpr>().value();
|
||||
node->strideA_ = args[10].as<IntImm>().value()->value;
|
||||
node->strideB_ = args[11].as<IntImm>().value()->value;
|
||||
node->offsetA_ = args[12].as<IntImm>().value()->value;
|
||||
node->offsetB_ = args[13].as<IntImm>().value()->value;
|
||||
if (args.size() > 14) {
|
||||
node->kPack_ = args[14].as<IntImm>().value()->value;
|
||||
if (node->kPack_ != 1 && node->kPack_ != 2) {
|
||||
ICHECK(false) << "kPack must be 1 or 2";
|
||||
}
|
||||
}
|
||||
if (args.size() > 15) {
|
||||
node->wgWait_ = args[15].as<IntImm>().value()->value;
|
||||
}
|
||||
if (auto val = annotations.Get("is_wgmma")) {
|
||||
const auto *int_val = val->as<IntImmNode>();
|
||||
ICHECK(int_val) << "is_wgmma annotation must be IntImmNode";
|
||||
node->isWgmma_ = int_val->value != 0;
|
||||
}
|
||||
if (auto val = annotations.Get("is_tcgen05")) {
|
||||
const auto *int_val = val->as<IntImmNode>();
|
||||
ICHECK(int_val) << "is_tcgen05 annotation must be IntImmNode";
|
||||
node->isTcgen05_ = int_val->value != 0;
|
||||
}
|
||||
if (args.size() > 16 && args[16]->IsInstance<BufferLoadNode>()) {
|
||||
node->mbar_ = Downcast<BufferLoad>(args[16]);
|
||||
}
|
||||
node->cCoords_ = Array<PrimExpr>(
|
||||
{args[17].as<PrimExpr>().value(), args[18].as<PrimExpr>().value()});
|
||||
node->annotations_ = annotations;
|
||||
data_ = std::move(node);
|
||||
}
|
||||
|
||||
AccessRegions GemmPyNode::GetAccessRegions() const {
|
||||
AccessRegions result;
|
||||
result.reads.push_back(aRegion_);
|
||||
result.reads.push_back(bRegion_);
|
||||
if (!is_one(clearAccum_)) {
|
||||
result.reads.push_back(cRegion_);
|
||||
}
|
||||
result.writes.push_back(cRegion_);
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Create a copy of this GemmPyNode as a TileOperator.
|
||||
*
|
||||
* Constructs a new GemmPyNode by copying the current node state and returns it
|
||||
* wrapped in a Gemm TileOperator.
|
||||
*
|
||||
* @return TileOperator A Gemm operator that owns a copy of this node.
|
||||
*/
|
||||
TileOperator GemmPyNode::Clone() const {
|
||||
auto op = tvm::ffi::make_object<GemmPyNode>(*this);
|
||||
return GemmPy(op);
|
||||
}
|
||||
|
||||
bool GemmPyNode::allowTcgen5Mma(Target target) const {
|
||||
bool scope_ok = (IsSharedBuffer(a_) || a_.scope() == "shared.tmem") &&
|
||||
IsSharedBuffer(b_) && c_.scope() == "shared.tmem";
|
||||
if (!TargetIsSm100(target) || !scope_ok)
|
||||
return false;
|
||||
// For TS variant (A from TMEM), use B's dtype as the input dtype
|
||||
DataType ab_dtype = (a_.scope() == "shared.tmem") ? b_->dtype : a_->dtype;
|
||||
return GetTCGEN5MMAMeta(m_, n_, k_, ab_dtype, c_->dtype).first;
|
||||
}
|
||||
|
||||
bool GemmPyNode::allowWgmma(int block_size, Target target) const {
|
||||
tvm::transform::PassContext ctxt = tvm::transform::PassContext::Current();
|
||||
|
||||
int warp_size = TargetGetWarpSize(target);
|
||||
int num_warps = block_size / warp_size;
|
||||
return !ctxt->GetConfig(kDisableWGMMA, Optional<Bool>()).value_or(false) &&
|
||||
TargetIsHopper(target) && (this->m_ >= 64) && (num_warps % 4 == 0) &&
|
||||
checkWgmma();
|
||||
}
|
||||
|
||||
GemmInst GemmPyNode::getGemmInst(int block_size, Target target) const {
|
||||
if (isWgmma_) {
|
||||
if (!allowWgmma(block_size, target)) {
|
||||
LOG(FATAL) << "T.wgmma_gemm() requires Hopper WGMMA lowering, but "
|
||||
"constraints were not satisfied. Got target="
|
||||
<< target << ", A(scope=" << a_.scope()
|
||||
<< ", dtype=" << a_->dtype << "), B(scope=" << b_.scope()
|
||||
<< ", dtype=" << b_->dtype << "), C(scope=" << c_.scope()
|
||||
<< ", dtype=" << c_->dtype << "), M=" << m_ << ", N=" << n_
|
||||
<< ", K=" << k_ << ".";
|
||||
}
|
||||
return GemmInst::kWGMMA;
|
||||
}
|
||||
if (isTcgen05_) {
|
||||
if (!allowTcgen5Mma(target)) {
|
||||
LOG(FATAL) << "T.tcgen05_gemm() requires Blackwell TCGEN5MMA lowering, "
|
||||
"but constraints were not satisfied. Got target="
|
||||
<< target << ", A(scope=" << a_.scope()
|
||||
<< ", dtype=" << a_->dtype << "), B(scope=" << b_.scope()
|
||||
<< ", dtype=" << b_->dtype << "), C(scope=" << c_.scope()
|
||||
<< ", dtype=" << c_->dtype << "), M=" << m_ << ", N=" << n_
|
||||
<< ", K=" << k_ << ".";
|
||||
}
|
||||
return GemmInst::kTCGEN5MMA;
|
||||
}
|
||||
bool allow_tcgen5mma = allowTcgen5Mma(target);
|
||||
bool allow_wgmma = allowWgmma(block_size, target);
|
||||
if (allow_tcgen5mma) {
|
||||
return GemmInst::kTCGEN5MMA;
|
||||
} else if (allow_wgmma) {
|
||||
return GemmInst::kWGMMA;
|
||||
} else if (TargetIsCDNA(target)) {
|
||||
return GemmInst::kMFMA;
|
||||
} else if (TargetIsRDNA(target)) {
|
||||
return GemmInst::kWMMA;
|
||||
} else if (TargetIsCuda(target)) {
|
||||
return GemmInst::kMMA;
|
||||
} else if (TargetIsCPU(target)) {
|
||||
return GemmInst::kScalar;
|
||||
} else {
|
||||
ICHECK(0) << "Unsupported target for gemm: " << target->str();
|
||||
return GemmInst::kMMA; // This line will never be reached due to ICHECK, but
|
||||
// satisfies compiler
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Checks whether WGMMA (warp-group MMA) can be used for this GEMM.
|
||||
*
|
||||
* Evaluates device-memory placement, data-type combinations, transpose flags,
|
||||
* and K divisibility constraints required for the Hopper WGMMA code path.
|
||||
*
|
||||
* The check returns true only when:
|
||||
* - B resides in shared memory ("shared" or "shared.dyn"); and
|
||||
* - (C, A, B) dtypes match one of the supported combinations below and K
|
||||
* satisfies the required alignment; and
|
||||
* - for combinations that require specific orientations, A is not transposed
|
||||
* and B is transposed.
|
||||
*
|
||||
* Supported combinations and constraints:
|
||||
* - C=float16:
|
||||
* - A=float16, B=float16: K % 16 == 0
|
||||
* - Various float8 mixes (e4m3/e5m2): require (!trans_A && trans_B) and K %
|
||||
* 32 == 0
|
||||
* - C=float32:
|
||||
* - A=float16, B=float16: K % 16 == 0
|
||||
* - A=bfloat16, B=bfloat16: K % 16 == 0
|
||||
* - A=float32, B=float32: require (!trans_A && trans_B) and K % 8 == 0
|
||||
* - Various float8 mixes: require (!trans_A && trans_B) and K % 32 == 0
|
||||
* - C=int32:
|
||||
* - 8-bit integer combinations (Int8/UInt8): require (!trans_A && trans_B)
|
||||
* and K % 32 == 0
|
||||
*
|
||||
* @return true if WGMMA is supported for the current buffers, dtypes, and
|
||||
* transpose/shape constraints; false otherwise.
|
||||
*/
|
||||
bool GemmPyNode::checkWgmma() const {
|
||||
if (b_.scope() != "shared.dyn" && b_.scope() != "shared") {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (c_->dtype == DataType::Float(16)) {
|
||||
if (a_->dtype == DataType::Float(16) && b_->dtype == DataType::Float(16))
|
||||
return k_ % 16 == 0;
|
||||
else if (a_->dtype.is_float8() && b_->dtype.is_float8())
|
||||
return (!transA_) && transB_ && k_ % 32 == 0;
|
||||
else
|
||||
return false;
|
||||
} else if (c_->dtype == DataType::Float(32)) {
|
||||
if (a_->dtype == DataType::Float(16) && b_->dtype == DataType::Float(16))
|
||||
return k_ % 16 == 0;
|
||||
else if (a_->dtype == DataType::BFloat(16) &&
|
||||
b_->dtype == DataType::BFloat(16))
|
||||
return k_ % 16 == 0;
|
||||
else if (a_->dtype == DataType::Float(32) &&
|
||||
b_->dtype == DataType::Float(32))
|
||||
return (!transA_) && transB_ && k_ % 8 == 0;
|
||||
else if (a_->dtype.is_float8() && b_->dtype.is_float8())
|
||||
return (!transA_) && transB_ && k_ % 32 == 0;
|
||||
else
|
||||
return false;
|
||||
} else if (c_->dtype == DataType::Int(32)) {
|
||||
if (a_->dtype == DataType::Int(8) && b_->dtype == DataType::Int(8))
|
||||
return (!transA_) && transB_ && k_ % 32 == 0;
|
||||
else if (a_->dtype == DataType::Int(8) && b_->dtype == DataType::UInt(8))
|
||||
return (!transA_) && transB_ && k_ % 32 == 0;
|
||||
else if (a_->dtype == DataType::UInt(8) && b_->dtype == DataType::Int(8))
|
||||
return (!transA_) && transB_ && k_ % 32 == 0;
|
||||
else if (a_->dtype == DataType::UInt(8) && b_->dtype == DataType::UInt(8))
|
||||
return (!transA_) && transB_ && k_ % 32 == 0;
|
||||
else
|
||||
return false;
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Parse and return the numeric GPU architecture from a Target's "arch"
|
||||
* attribute.
|
||||
*
|
||||
* Examines the target's "arch" string and, if it matches the pattern
|
||||
* "sm_<num>", returns <num> as an int. If the attribute is present but does not
|
||||
* match that pattern, returns 0.
|
||||
*
|
||||
* Preconditions: the target must have an "arch" attribute (this is checked via
|
||||
* ICHECK).
|
||||
*
|
||||
* @return int The parsed architecture number (e.g., 80 for "sm_80"), or 0 if
|
||||
* the arch string does not match "sm_<num>".
|
||||
*/
|
||||
static int GetArchInt(Target target) {
|
||||
int arch_int = 0;
|
||||
auto s = target->GetAttr<tvm::ffi::String>("arch");
|
||||
ICHECK(s.has_value());
|
||||
std::string arch = s.value();
|
||||
if (arch.rfind("sm_", 0) == 0) {
|
||||
arch_int = std::stoi(arch.substr(3));
|
||||
} else {
|
||||
arch_int = 0;
|
||||
}
|
||||
return arch_int;
|
||||
}
|
||||
|
||||
Stmt GemmPyNode::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const {
|
||||
if (const auto f = ffi::Function::GetGlobal("tl.gemm_py.lower")) {
|
||||
PrimExpr mbar_phase = T.mbar_phase_expr;
|
||||
if (auto explicit_phase = GetAnnotatedMbarPhaseExpr(annotations_)) {
|
||||
mbar_phase = explicit_phase.value();
|
||||
}
|
||||
// NOTE(wt): Decide GemmInst and compute warp partition on Python side
|
||||
auto prim_func = Downcast<PrimFunc>(
|
||||
(*f)(tvm::ffi::GetRef<GemmPy>(this), T.layout_map, T.target,
|
||||
T.thread_bounds, T.thread_var, mbar_phase));
|
||||
ICHECK(prim_func->attrs.defined());
|
||||
auto global_symbol =
|
||||
prim_func->attrs.GetAttr<tvm::ffi::String>("global_symbol");
|
||||
ICHECK(global_symbol.has_value());
|
||||
if (prim_func->body.as<BlockRealizeNode>()) {
|
||||
BlockRealize block_realize = Downcast<BlockRealize>(prim_func->body);
|
||||
auto block = block_realize->block;
|
||||
{
|
||||
BlockNode *n = block.CopyOnWrite();
|
||||
n->name_hint = global_symbol.value();
|
||||
n->annotations.Set(tl::attr::kLexicalAllocScope,
|
||||
IntImm(DataType::Int(32), 1));
|
||||
}
|
||||
return BlockRealize(block_realize->iter_values, block_realize->predicate,
|
||||
block);
|
||||
}
|
||||
// warp with block realize node
|
||||
Map<String, ObjectRef> block_annotations;
|
||||
block_annotations.Set(tl::attr::kLexicalAllocScope,
|
||||
IntImm(DataType::Int(32), 1));
|
||||
return BlockRealize(
|
||||
/*iter_values=*/Array<PrimExpr>(),
|
||||
/*predicate=*/const_true(),
|
||||
/*block=*/
|
||||
Block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{},
|
||||
/*name_hint=*/global_symbol.value(), prim_func->body,
|
||||
/*init=*/Optional<Stmt>(), /*alloc_buffers=*/{},
|
||||
/*match_buffers=*/{}, /*annotations=*/block_annotations));
|
||||
} else {
|
||||
LOG(FATAL) << "No lower function found for gemm_py";
|
||||
return Stmt(); // This line will never be reached due to LOG(FATAL), but
|
||||
// satisfies compiler
|
||||
}
|
||||
}
|
||||
|
||||
LayoutMap GemmPyNode::InferLayout(const LayoutInferArgs &T,
|
||||
InferLevel level) const {
|
||||
if (completed_)
|
||||
return {};
|
||||
LayoutMap results;
|
||||
|
||||
if (const auto f = ffi::Function::GetGlobal("tl.gemm_py.infer_layout")) {
|
||||
results = Downcast<LayoutMap>(
|
||||
(*f)(tvm::ffi::GetRef<GemmPy>(this), T.target, T.thread_bounds));
|
||||
// Bind all fragment layouts with the provided thread range
|
||||
for (auto kv : results) {
|
||||
const Buffer &buf = kv.first;
|
||||
const Layout &layout = kv.second;
|
||||
if (auto frag = layout.as<Fragment>()) {
|
||||
results.Set(buf, frag.value()->BindThreadRange(T.thread_bounds));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
LOG(FATAL) << "No infer layout function found for gemm_py";
|
||||
}
|
||||
|
||||
completed_ = true;
|
||||
return results;
|
||||
}
|
||||
|
||||
TIR_REGISTER_TL_TILE_OP(GemmPy, gemm_py)
|
||||
.set_num_inputs(5)
|
||||
.set_attr<TCallEffectKind>("TCallEffectKind",
|
||||
Integer(CallEffectKind::kOpaque));
|
||||
|
||||
TVM_REGISTER_OP("tl.tileop.wgmma_gemm_py")
|
||||
.set_attr<TScriptPrinterName>("TScriptPrinterName", "wgmma_gemm_py")
|
||||
.set_attr<OpBuilderFunc>("TLOpBuilder",
|
||||
[](Array<PrimExpr> args,
|
||||
Map<String, ObjectRef> annotations) {
|
||||
Map<String, ObjectRef> ann = annotations;
|
||||
ann.Set("is_wgmma",
|
||||
IntImm(DataType::Int(32), 1));
|
||||
return GemmPy(args, ann);
|
||||
})
|
||||
.set_num_inputs(5)
|
||||
.set_attr<TCallEffectKind>("TCallEffectKind",
|
||||
Integer(CallEffectKind::kOpaque));
|
||||
|
||||
TVM_REGISTER_OP("tl.tileop.tcgen05_gemm_py")
|
||||
.set_attr<TScriptPrinterName>("TScriptPrinterName", "tcgen05_gemm_py")
|
||||
.set_attr<OpBuilderFunc>("TLOpBuilder",
|
||||
[](Array<PrimExpr> args,
|
||||
Map<String, ObjectRef> annotations) {
|
||||
Map<String, ObjectRef> ann = annotations;
|
||||
ann.Set("is_tcgen05",
|
||||
IntImm(DataType::Int(32), 1));
|
||||
return GemmPy(args, ann);
|
||||
})
|
||||
.set_num_inputs(5)
|
||||
.set_attr<TCallEffectKind>("TCallEffectKind",
|
||||
Integer(CallEffectKind::kOpaque));
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { GemmPyNode::RegisterReflection(); }
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("tl.GemmPyGemmInst",
|
||||
[](GemmPy gemm_py, int block_size, Target target) {
|
||||
return gemm_py->getGemmInst(block_size, target);
|
||||
});
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def(
|
||||
"tl.get_tcgen5_mma_meta", [](int M, int N, int K, DataType ab_dtype,
|
||||
DataType c_dtype, bool disable_2cta) {
|
||||
auto [success, meta] =
|
||||
GetTCGEN5MMAMeta(M, N, K, ab_dtype, c_dtype, disable_2cta);
|
||||
Array<Integer> result;
|
||||
if (success) {
|
||||
result.push_back(Integer(meta.atom_m));
|
||||
result.push_back(Integer(meta.atom_n));
|
||||
result.push_back(Integer(meta.atom_k));
|
||||
result.push_back(Integer(meta.enable_ws));
|
||||
result.push_back(Integer(meta.enable_2cta));
|
||||
}
|
||||
return result;
|
||||
});
|
||||
refl::GlobalDef().def(
|
||||
"tl.get_tcgen5_instr_desc",
|
||||
[](int atom_m, int atom_n, int atom_k, DataType ab_dtype,
|
||||
DataType c_dtype, bool a_is_k_major, bool b_is_k_major, int scale_in_a,
|
||||
int scale_in_b) {
|
||||
uint32_t desc = GetTCGEN5InstrDesc(atom_m, atom_n, atom_k, ab_dtype,
|
||||
c_dtype, a_is_k_major, b_is_k_major,
|
||||
scale_in_a, scale_in_b);
|
||||
return Integer(static_cast<int64_t>(desc));
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace tl
|
||||
} // namespace tvm
|
||||
@@ -1,99 +0,0 @@
|
||||
/*!
|
||||
* \file tl/op/gemm_py.h
|
||||
* \brief Define gemm operator.
|
||||
*
|
||||
*/
|
||||
|
||||
#ifndef TVM_TL_OP_GEMM_PY_H_
|
||||
#define TVM_TL_OP_GEMM_PY_H_
|
||||
|
||||
#include "gemm.h"
|
||||
#include "operator.h"
|
||||
|
||||
namespace tvm {
|
||||
|
||||
namespace tl {
|
||||
|
||||
using namespace tir;
|
||||
|
||||
class GemmPyNode : public TileOperatorNode {
|
||||
public:
|
||||
bool checkWgmma() const;
|
||||
bool allowTcgen5Mma(Target target) const;
|
||||
bool allowWgmma(int block_size, Target target) const;
|
||||
tir::Buffer a_, b_, c_;
|
||||
// BufferRegion for A, B and C
|
||||
BufferRegion aRegion_, bRegion_, cRegion_;
|
||||
bool transA_, transB_;
|
||||
int m_, n_, k_;
|
||||
int strideA_, strideB_;
|
||||
int offsetA_, offsetB_;
|
||||
PrimExpr clearAccum_ = const_false();
|
||||
tir::BufferLoad mbar_; // mbar is optional, only used for TCGEN5MMA
|
||||
Array<PrimExpr> cCoords_;
|
||||
// k_pack please ref to bitblas/tl/mfma_macro_generator.py::k_pack
|
||||
// only will be enabled under cdna mfma instructions
|
||||
int kPack_ = 1;
|
||||
int wgWait_ = 0;
|
||||
bool isWgmma_ = false;
|
||||
bool isTcgen05_ = false;
|
||||
mutable GemmWarpPolicy policy_;
|
||||
Map<String, ObjectRef> annotations_;
|
||||
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tl.GemmPy", GemmPyNode, TileOperatorNode);
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<GemmPyNode>()
|
||||
.def_ro("a", &GemmPyNode::a_)
|
||||
.def_ro("b", &GemmPyNode::b_)
|
||||
.def_ro("c", &GemmPyNode::c_)
|
||||
.def_ro("aRegion", &GemmPyNode::aRegion_)
|
||||
.def_ro("bRegion", &GemmPyNode::bRegion_)
|
||||
.def_ro("cRegion", &GemmPyNode::cRegion_)
|
||||
.def_ro("transA", &GemmPyNode::transA_)
|
||||
.def_ro("transB", &GemmPyNode::transB_)
|
||||
.def_ro("m", &GemmPyNode::m_)
|
||||
.def_ro("n", &GemmPyNode::n_)
|
||||
.def_ro("k", &GemmPyNode::k_)
|
||||
.def_ro("strideA", &GemmPyNode::strideA_)
|
||||
.def_ro("strideB", &GemmPyNode::strideB_)
|
||||
.def_ro("offsetA", &GemmPyNode::offsetA_)
|
||||
.def_ro("offsetB", &GemmPyNode::offsetB_)
|
||||
.def_ro("clearAccum", &GemmPyNode::clearAccum_)
|
||||
.def_ro("mbar", &GemmPyNode::mbar_)
|
||||
.def_ro("cCoords", &GemmPyNode::cCoords_)
|
||||
.def_ro("kPack", &GemmPyNode::kPack_)
|
||||
.def_ro("wgWait", &GemmPyNode::wgWait_)
|
||||
.def_ro("isWgmma", &GemmPyNode::isWgmma_)
|
||||
.def_ro("isTcgen05", &GemmPyNode::isTcgen05_)
|
||||
.def_ro("policy", &GemmPyNode::policy_)
|
||||
.def_ro("annotations", &GemmPyNode::annotations_);
|
||||
}
|
||||
|
||||
Stmt Lower(const LowerArgs &T, arith::Analyzer *analyzer) const override;
|
||||
LayoutMap InferLayout(const LayoutInferArgs &T,
|
||||
InferLevel level) const override;
|
||||
AccessRegions GetAccessRegions() const override;
|
||||
|
||||
TileOperator Clone() const;
|
||||
|
||||
// Target GEMM instruction
|
||||
GemmInst getGemmInst(int block_size, Target target) const;
|
||||
|
||||
private:
|
||||
mutable bool completed_ = false;
|
||||
};
|
||||
|
||||
class GemmPy : public TileOperator {
|
||||
public:
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GemmPy, TileOperator, GemmPyNode);
|
||||
TVM_DLL GemmPy(Array<PrimExpr> args,
|
||||
Map<String, ObjectRef> annotations = Map<String, ObjectRef>());
|
||||
static const Op &Get();
|
||||
};
|
||||
|
||||
} // namespace tl
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_TL_OP_GEMM_PY_H_
|
||||
+1
-1
@@ -4,7 +4,7 @@
|
||||
*
|
||||
*/
|
||||
|
||||
// TODO: @botbw: remove redundant code with gemm_py.h
|
||||
// TODO: @botbw: remove redundant code with gemm.h
|
||||
|
||||
#ifndef TVM_TL_OP_GEMM_SP_PY_H_
|
||||
#define TVM_TL_OP_GEMM_SP_PY_H_
|
||||
|
||||
@@ -19,7 +19,6 @@
|
||||
#include "../op/builtin.h"
|
||||
#include "../op/copy.h"
|
||||
#include "../op/gemm.h"
|
||||
#include "../op/gemm_py.h"
|
||||
#include "../op/operator.h"
|
||||
#include "../op/region.h"
|
||||
#include "../op/utils.h"
|
||||
@@ -404,8 +403,7 @@ private:
|
||||
auto tile_op = ParseOperator(call);
|
||||
return tile_op.defined() && (tile_op.as<CopyNode>() != nullptr ||
|
||||
tile_op.as<Conv2DIm2ColOpNode>() != nullptr ||
|
||||
tile_op.as<GemmNode>() != nullptr ||
|
||||
tile_op.as<GemmPyNode>() != nullptr);
|
||||
tile_op.as<GemmNode>() != nullptr);
|
||||
}
|
||||
|
||||
PrimExpr phase_expr_;
|
||||
|
||||
@@ -3,11 +3,8 @@
|
||||
* \brief Lower 2SM TCGEN5MMA and related on Blackwell target
|
||||
*
|
||||
* This pass runs before LowerTileOp. At that point the IR still has T.gemm
|
||||
* (tl_gemm / tl.tileop.gemm_py Call), not the lowered tl::tcgen5mma_gemm_ss/ts.
|
||||
* We detect Gemm ops that will be lowered to TCGEN5MMA with use_2cta and set
|
||||
* block attr.
|
||||
*
|
||||
* Tilelang gemm defaults to v2 (GemmPyNode); we only support v2, not v1 (Gemm).
|
||||
* (tl.tileop.gemm Call), not the lowered tl::tcgen5mma_gemm_ss/ts. We detect
|
||||
* Gemm ops that will be lowered to TCGEN5MMA with use_2cta and set block attr.
|
||||
*/
|
||||
|
||||
// todo: consider mixture of 1cta/2cta tcgen5mma in the same kernel
|
||||
@@ -20,7 +17,7 @@
|
||||
#include <tvm/tir/stmt_functor.h>
|
||||
#include <tvm/tir/transform.h>
|
||||
|
||||
#include "../op/gemm_py.h"
|
||||
#include "../op/gemm.h"
|
||||
#include "../op/operator.h"
|
||||
#include "../target/utils.h"
|
||||
|
||||
@@ -62,9 +59,8 @@ static bool HasValidClusterDimsFor2Cta(const Stmt &body) {
|
||||
|
||||
/**
|
||||
* \brief Detect 2SM TCGEN5MMA in the kernel (before LowerTileOp).
|
||||
* Looks for T.gemm (tl_gemm() Call); if it will be lowered to TCGEN5MMA with
|
||||
* use_2cta, sets the flag for the mutator to add block attr.
|
||||
* Only supports v2 (GemmPy); v1 (Gemm) is ignored.
|
||||
* Looks for T.gemm (tl.tileop.gemm Call); if it will be lowered to TCGEN5MMA
|
||||
* with use_2cta, sets the flag for the mutator to add block attr.
|
||||
*/
|
||||
class Tcgen5_2SmLower : public StmtExprMutator {
|
||||
public:
|
||||
@@ -76,7 +72,7 @@ private:
|
||||
Stmt VisitStmt_(const EvaluateNode *op) final {
|
||||
if (const CallNode *call = op->value.as<CallNode>()) {
|
||||
TileOperator tile_op = ParseOperator(ffi::GetRef<Stmt>(op));
|
||||
if (tile_op.defined() && tile_op.as<GemmPy>()) {
|
||||
if (tile_op.defined() && tile_op.as<Gemm>()) {
|
||||
// Check if the user explicitly requested 2CTA via the use_2cta
|
||||
// annotation on the Call node (set by T.tcgen05_gemm(use_2cta=True)).
|
||||
if (call->annotations.count(attr::kUse2Cta)) {
|
||||
|
||||
@@ -106,7 +106,7 @@ private:
|
||||
}
|
||||
// Step 5. Materialize a lexical scope boundary only for blocks that were
|
||||
// explicitly marked by an earlier semantic lowering pass (for example
|
||||
// gemm_py/gemm_sp_py). We intentionally avoid re-inferring this from the
|
||||
// gemm/gemm_sp_py). We intentionally avoid re-inferring this from the
|
||||
// lowered alloc_buffers here because provenance has already been blurred by
|
||||
// earlier allocation planning/hoisting passes.
|
||||
if (HasLexicalAllocScopeAnnotation(new_block->annotations)) {
|
||||
|
||||
@@ -31,7 +31,6 @@
|
||||
#include "../op/copy.h"
|
||||
#include "../op/fill.h"
|
||||
#include "../op/gemm.h"
|
||||
#include "../op/gemm_py.h"
|
||||
#include "../op/operator.h"
|
||||
#include "../op/region.h"
|
||||
#include "../op/utils.h"
|
||||
@@ -839,8 +838,7 @@ private:
|
||||
auto tile_op = ParseOperator(call);
|
||||
return tile_op.defined() && (tile_op.as<CopyNode>() != nullptr ||
|
||||
tile_op.as<Conv2DIm2ColOpNode>() != nullptr ||
|
||||
tile_op.as<GemmNode>() != nullptr ||
|
||||
tile_op.as<GemmPyNode>() != nullptr);
|
||||
tile_op.as<GemmNode>() != nullptr);
|
||||
}
|
||||
|
||||
PrimExpr phase_expr_;
|
||||
|
||||
@@ -45,6 +45,7 @@ def run_tilelang_if_range(M=128, N=128, block_M=32, block_N=32, dtype=T.float16)
|
||||
torch.testing.assert_close(b, ref_b, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
@tilelang.testing.requires_cuda
|
||||
def test_tilelang_if_range():
|
||||
run_tilelang_if_range(M=128, N=128, block_M=32, block_N=32)
|
||||
|
||||
|
||||
@@ -21,6 +21,7 @@ def _kernel_min_blocks_per_sm():
|
||||
return main
|
||||
|
||||
|
||||
@tilelang.testing.requires_cuda
|
||||
def test_annotate_min_blocks_per_sm_launch_bounds():
|
||||
"""Codegen should emit the second __launch_bounds__ argument from the annotation."""
|
||||
src = _kernel_min_blocks_per_sm.get_kernel_source()
|
||||
|
||||
@@ -94,12 +94,14 @@ def run_tilelang_transpose_square(M=256, block_M=128, dtype=T.float16):
|
||||
print(f"PASS: square transpose M={M}, block_M={block_M}")
|
||||
|
||||
|
||||
@tilelang.testing.requires_cuda
|
||||
def test_tilelang_transpose():
|
||||
run_tilelang_transpose(M=128, N=128, block_M=128, block_N=128)
|
||||
run_tilelang_transpose(M=256, N=256, block_M=128, block_N=128)
|
||||
run_tilelang_transpose(M=128, N=256, block_M=128, block_N=256)
|
||||
|
||||
|
||||
@tilelang.testing.requires_cuda
|
||||
def test_tilelang_transpose_square():
|
||||
run_tilelang_transpose_square(M=128, block_M=128)
|
||||
run_tilelang_transpose_square(M=256, block_M=128)
|
||||
@@ -107,5 +109,4 @@ def test_tilelang_transpose_square():
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_tilelang_transpose()
|
||||
test_tilelang_transpose_square()
|
||||
tilelang.testing.main()
|
||||
|
||||
@@ -1,624 +0,0 @@
|
||||
import tilelang.language as T
|
||||
from tilelang import tvm as tvm
|
||||
import tilelang.testing
|
||||
from tilelang.utils import determine_fp8_type
|
||||
import pytest
|
||||
|
||||
|
||||
def matmul(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype)
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
if trans_A:
|
||||
T.copy(A[k * block_K, by * block_M], A_shared)
|
||||
else:
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
if trans_B:
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
else:
|
||||
T.copy(B[k * block_K, bx * block_N], B_shared)
|
||||
T.gemm_v2(A_shared, B_shared, C_local, trans_A, trans_B)
|
||||
# T.gemm(A_shared, B_shared, C_local, trans_A, trans_B)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def run_gemm_ss(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=3,
|
||||
num_threads=128,
|
||||
):
|
||||
program = matmul(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
|
||||
kernel = tilelang.compile(
|
||||
program,
|
||||
out_idx=[2],
|
||||
pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
|
||||
)
|
||||
|
||||
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
||||
|
||||
def ref_program(A, B):
|
||||
import torch
|
||||
|
||||
if trans_A:
|
||||
A = A.T
|
||||
if trans_B:
|
||||
B = B.T
|
||||
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
||||
C = C.to(torch.__getattribute__(out_dtype))
|
||||
return C
|
||||
|
||||
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Temporarily disabling until GEMM SS issues are resolved")
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
(512, 1024, 768, False, True, T.float16, T.float16, T.float32, 128, 128, 32, 2, 128),
|
||||
(512, 1024, 768, False, False, T.float16, T.float16, T.float32, 128, 128, 32, 2, 128),
|
||||
(512, 1024, 768, True, False, T.float16, T.float16, T.float32, 128, 128, 32, 2, 128),
|
||||
(512, 1024, 768, True, True, T.float16, T.float16, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 16, 32, False, True, T.float16, T.float16, T.float32, 128, 16, 32, 0, 128),
|
||||
(128, 128, 128, False, True, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, False, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, False, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, False, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, True, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, False, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_ss(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_ss(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Temporarily disabling until GEMM SS issues are resolved")
|
||||
@tilelang.testing.requires_cuda
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
(128, 128, 128, True, True, T.float8_e5m2, T.float8_e5m2, T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_ss_fp8_cuda(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_ss(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Temporarily disabling until GEMM SS issues are resolved")
|
||||
@tilelang.testing.requires_rocm
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
(128, 128, 128, True, True, determine_fp8_type("e5m2"), determine_fp8_type("e5m2"), T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, determine_fp8_type(), determine_fp8_type(), T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_ss_fp8_rocm(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_ss(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
def matmul_rs(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
A_frag_shape = A_shared_shape
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype)
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
||||
A_frag = T.alloc_fragment(A_frag_shape, in_dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
T.annotate_layout(
|
||||
{
|
||||
A_shared: tilelang.layout.make_swizzled_layout(A_shared),
|
||||
}
|
||||
)
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
if trans_A:
|
||||
T.copy(A[k * block_K, by * block_M], A_shared)
|
||||
else:
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
if trans_B:
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
else:
|
||||
T.copy(B[k * block_K, bx * block_N], B_shared)
|
||||
T.copy(A_shared, A_frag)
|
||||
T.gemm_v2(A_frag, B_shared, C_local, trans_A, trans_B)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def run_gemm_rs(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=3,
|
||||
num_threads=128,
|
||||
):
|
||||
program = matmul_rs(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
|
||||
kernel = tilelang.compile(
|
||||
program,
|
||||
out_idx=[2],
|
||||
pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
|
||||
)
|
||||
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
||||
|
||||
def ref_program(A, B):
|
||||
import torch
|
||||
|
||||
if trans_A:
|
||||
A = A.T
|
||||
if trans_B:
|
||||
B = B.T
|
||||
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
||||
C = C.to(torch.__getattribute__(out_dtype))
|
||||
return C
|
||||
|
||||
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Temporarily disabling until GEMM RS issues are resolved")
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
(512, 1024, 768, False, False, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(512, 1024, 768, False, True, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(512, 1024, 768, True, False, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(512, 1024, 768, True, True, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(128, 16, 32, False, True, T.float16, T.float16, T.float32, 128, 16, 32, 0, 128),
|
||||
(128, 128, 128, False, True, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, False, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, False, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, False, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, True, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, False, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_rs(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_rs(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Temporarily disabling until GEMM RS issues are resolved")
|
||||
@tilelang.testing.requires_cuda
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
(128, 128, 128, True, True, T.float8_e5m2, T.float8_e5m2, T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_rs_fp8_cuda(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_rs(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Temporarily disabling until GEMM RS issues are resolved")
|
||||
@tilelang.testing.requires_rocm
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
(128, 128, 128, True, True, T.float8_e5m2fnuz, T.float8_e5m2fnuz, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, T.float8_e4m3fnuz, T.float8_e4m3fnuz, T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_rs_fp8_rocm(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_rs(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
def matmul_sr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
B_frag_shape = B_shared_shape
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype)
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
||||
B_frag = T.alloc_fragment(B_frag_shape, in_dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
T.annotate_layout(
|
||||
{
|
||||
B_shared: tilelang.layout.make_swizzled_layout(B_shared),
|
||||
}
|
||||
)
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
if trans_A:
|
||||
T.copy(A[k * block_K, by * block_M], A_shared)
|
||||
else:
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
if trans_B:
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
else:
|
||||
T.copy(B[k * block_K, bx * block_N], B_shared)
|
||||
T.copy(B_shared, B_frag)
|
||||
T.gemm_v2(A_shared, B_frag, C_local, trans_A, trans_B)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def run_gemm_sr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=3,
|
||||
num_threads=128,
|
||||
):
|
||||
program = matmul_sr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
|
||||
kernel = tilelang.compile(
|
||||
program,
|
||||
out_idx=[2],
|
||||
pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
|
||||
)
|
||||
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
||||
|
||||
def ref_program(A, B):
|
||||
import torch
|
||||
|
||||
if trans_A:
|
||||
A = A.T
|
||||
if trans_B:
|
||||
B = B.T
|
||||
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
||||
C = C.to(torch.__getattribute__(out_dtype))
|
||||
return C
|
||||
|
||||
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
(512, 1024, 768, False, False, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(512, 1024, 768, False, True, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(512, 1024, 768, True, False, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(512, 1024, 768, True, True, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(128, 16, 32, False, True, T.float16, T.float16, T.float32, 128, 16, 32, 0, 128),
|
||||
(128, 128, 32, False, True, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 32, False, False, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 32, True, False, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 32, True, True, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, False, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, True, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, False, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_sr(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_sr(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
@tilelang.testing.requires_cuda
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
(128, 128, 128, True, True, T.float8_e5m2, T.float8_e5m2, T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_sr_fp8_cuda(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_sr(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
@tilelang.testing.requires_rocm
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
# TODO: There is precision problem needs to repair
|
||||
# (128, 128, 128, True, True, determine_fp8_type("e5m2"), determine_fp8_type("e5m2"), T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, determine_fp8_type(), determine_fp8_type(), T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_sr_fp8_rocm(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_sr(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
def matmul_rr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
accum_dtype,
|
||||
num_stages,
|
||||
threads,
|
||||
):
|
||||
A_shape = (K, M) if trans_A else (M, K)
|
||||
B_shape = (N, K) if trans_B else (K, N)
|
||||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
||||
A_frag_shape = A_shared_shape
|
||||
B_frag_shape = B_shared_shape
|
||||
|
||||
import tilelang.language as T
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Tensor(A_shape, in_dtype),
|
||||
B: T.Tensor(B_shape, in_dtype),
|
||||
C: T.Tensor((M, N), out_dtype),
|
||||
):
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
||||
A_shared = T.alloc_shared(A_shared_shape, in_dtype)
|
||||
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
||||
A_frag = T.alloc_fragment(A_frag_shape, in_dtype)
|
||||
B_frag = T.alloc_fragment(B_frag_shape, in_dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
T.annotate_layout(
|
||||
{
|
||||
A_shared: tilelang.layout.make_swizzled_layout(A_shared),
|
||||
B_shared: tilelang.layout.make_swizzled_layout(B_shared),
|
||||
}
|
||||
)
|
||||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||||
if trans_A:
|
||||
T.copy(A[k * block_K, by * block_M], A_shared)
|
||||
else:
|
||||
T.copy(A[by * block_M, k * block_K], A_shared)
|
||||
if trans_B:
|
||||
T.copy(B[bx * block_N, k * block_K], B_shared)
|
||||
else:
|
||||
T.copy(B[k * block_K, bx * block_N], B_shared)
|
||||
T.copy(A_shared, A_frag)
|
||||
T.copy(B_shared, B_frag)
|
||||
T.gemm_v2(A_frag, B_frag, C_local, trans_A, trans_B)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def run_gemm_rr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
num_stages=3,
|
||||
num_threads=128,
|
||||
):
|
||||
program = matmul_rr(
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_M,
|
||||
block_N,
|
||||
block_K,
|
||||
trans_A,
|
||||
trans_B,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
dtypeAccum,
|
||||
num_stages,
|
||||
num_threads,
|
||||
)
|
||||
|
||||
kernel = tilelang.compile(
|
||||
program,
|
||||
out_idx=[2],
|
||||
pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
|
||||
)
|
||||
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
||||
|
||||
def ref_program(A, B):
|
||||
import torch
|
||||
|
||||
if trans_A:
|
||||
A = A.T
|
||||
if trans_B:
|
||||
B = B.T
|
||||
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
||||
C = C.to(torch.__getattribute__(out_dtype))
|
||||
return C
|
||||
|
||||
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
(512, 1024, 768, False, False, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(512, 1024, 768, False, True, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(512, 1024, 768, True, False, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(512, 1024, 768, True, True, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128),
|
||||
(512, 1024, 768, False, True, T.bfloat16, T.bfloat16, T.float, 128, 256, 32, 2, 128),
|
||||
# TODO: There is precision problem when num_stages=2 on ROCm
|
||||
# (128, 16, 128, False, True, T.float16, T.float16, T.float32, 128, 16, 32, 2, 128)
|
||||
# (128, 16, 128, False, True, T.int8, T.int8, T.int32, 128, 16, 32, 2, 128),
|
||||
(128, 128, 128, False, True, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, False, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, False, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, T.int8, T.int8, T.int32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, False, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, False, True, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, False, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, T.float, T.float, T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_rr(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_rr(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
@tilelang.testing.requires_cuda
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
(128, 128, 128, True, True, T.float8_e5m2, T.float8_e5m2, T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_rr_fp8_cuda(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_rr(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
@tilelang.testing.requires_rocm
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads",
|
||||
[
|
||||
# TODO: There is precision problem needs to repair
|
||||
# (128, 128, 128, True, True, T.float8_e5m2fnuz, T.float8_e5m2fnuz, T.float32, 128, 128, 32, 2, 128),
|
||||
(128, 128, 128, True, True, determine_fp8_type(), determine_fp8_type(), T.float32, 128, 128, 32, 2, 128),
|
||||
],
|
||||
)
|
||||
def test_gemm_rr_fp8_rocm(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads):
|
||||
run_gemm_rr(M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tilelang.testing.main()
|
||||
@@ -272,10 +272,6 @@ class Environment:
|
||||
"TILELANG_CLEANUP_TEMP_FILES", "0"
|
||||
) # cleanup temporary compiler files/dirs after compilation (default: keep for debugging)
|
||||
|
||||
# Kernel selection options
|
||||
# Default to GEMM v2; set to "1"/"true"/"yes"/"on" to force v1
|
||||
TILELANG_USE_GEMM_V1 = EnvVar("TILELANG_USE_GEMM_V1", "0")
|
||||
|
||||
# Auto-tuning settings
|
||||
TILELANG_AUTO_TUNING_DISABLE_CACHE = EnvVar("TILELANG_AUTO_TUNING_DISABLE_CACHE", "0")
|
||||
TILELANG_AUTO_TUNING_CPU_UTILITIES = EnvVar("TILELANG_AUTO_TUNING_CPU_UTILITIES", "0.9") # percent of CPUs used
|
||||
@@ -327,14 +323,6 @@ class Environment:
|
||||
def should_cleanup_temp_files(self) -> bool:
|
||||
return str(self.TILELANG_CLEANUP_TEMP_FILES).lower() in ("1", "true", "yes", "on")
|
||||
|
||||
def use_gemm_v1(self) -> bool:
|
||||
"""Return True if GEMM v1 should be used based on env.
|
||||
|
||||
Controlled by `TILELANG_USE_GEMM_V1`. Truthy values are one of
|
||||
{"1", "true", "yes", "on"} (case-insensitive).
|
||||
"""
|
||||
return str(self.TILELANG_USE_GEMM_V1).lower() in ("1", "true", "yes", "on")
|
||||
|
||||
def get_default_target(self) -> str:
|
||||
"""Get default compilation target from environment."""
|
||||
return self.TILELANG_DEFAULT_TARGET
|
||||
|
||||
@@ -45,10 +45,6 @@ class GemmSPWarpPolicy(Node, Scriptable):
|
||||
return self.m_warp, self.n_warp
|
||||
|
||||
|
||||
@tvm_ffi.register_object("tl.Gemm")
|
||||
class Gemm(Node, Scriptable): ...
|
||||
|
||||
|
||||
@tvm_ffi.register_object("tl.GemmSP")
|
||||
class GemmSP(Node, Scriptable): ...
|
||||
|
||||
|
||||
@@ -80,7 +80,7 @@ def resolve_execution_backend(requested: str | None, target: Target) -> str:
|
||||
if is_cutedsl_target(target):
|
||||
return "cutedsl"
|
||||
kind = _target_kind(target)
|
||||
if kind == "cuda" or kind == "metal":
|
||||
if kind == "cuda" or kind == "metal" or kind == "hip":
|
||||
choice = "tvm_ffi"
|
||||
else:
|
||||
choice = "cython"
|
||||
|
||||
@@ -56,7 +56,7 @@ from .allocate import (
|
||||
from tvm.script.parser.tir import allocate as allocate # noqa: F401
|
||||
from .copy_op import copy, async_copy, tma_copy, transpose, c2d_im2col # noqa: F401
|
||||
from tilelang.tileop.base import GemmWarpPolicy # noqa: F401
|
||||
from .gemm_op import gemm, gemm_v1, gemm_v2, wgmma_gemm, tcgen05_gemm # noqa: F401
|
||||
from .gemm_op import gemm, wgmma_gemm, tcgen05_gemm # noqa: F401
|
||||
from .experimental.gemm_sp import gemm_sp, gemm_sp_v2 # noqa: F401
|
||||
from .fill_op import fill, clear # noqa: F401
|
||||
from .reduce_op import (
|
||||
|
||||
@@ -16,7 +16,6 @@ from tilelang.utils.language import (
|
||||
from tilelang.language.utils import (
|
||||
buffer_region_to_tile_region,
|
||||
)
|
||||
from tilelang.env import env as _env
|
||||
|
||||
|
||||
def _gemm_impl(
|
||||
@@ -146,62 +145,6 @@ def _gemm_impl(
|
||||
)
|
||||
|
||||
|
||||
# Public wrappers
|
||||
def gemm_v1(
|
||||
A: BufferLikeType,
|
||||
B: BufferLikeType,
|
||||
C: BufferLikeType,
|
||||
transpose_A: bool = False,
|
||||
transpose_B: bool = False,
|
||||
policy: GemmWarpPolicy = GemmWarpPolicy.Square,
|
||||
clear_accum: bool = False,
|
||||
k_pack: int = 1,
|
||||
mbar: BarrierType | None = None,
|
||||
) -> tir.PrimExpr:
|
||||
"""Synchronous GEMM v1: use op tl.gemm."""
|
||||
return _gemm_impl(
|
||||
"tl.tileop.gemm",
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
transpose_A,
|
||||
transpose_B,
|
||||
policy,
|
||||
clear_accum,
|
||||
k_pack,
|
||||
0,
|
||||
mbar,
|
||||
)
|
||||
|
||||
|
||||
# experimental currently, for fast compilation
|
||||
def gemm_v2(
|
||||
A: BufferLikeType,
|
||||
B: BufferLikeType,
|
||||
C: BufferLikeType,
|
||||
transpose_A: bool = False,
|
||||
transpose_B: bool = False,
|
||||
policy: GemmWarpPolicy = GemmWarpPolicy.Square,
|
||||
clear_accum: bool = False,
|
||||
k_pack: int = 1,
|
||||
mbar: BarrierType | None = None,
|
||||
) -> tir.PrimExpr:
|
||||
"""Synchronous GEMM v2: use op tl.gemm_py."""
|
||||
return _gemm_impl(
|
||||
"tl.tileop.gemm_py",
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
transpose_A,
|
||||
transpose_B,
|
||||
policy,
|
||||
clear_accum,
|
||||
k_pack,
|
||||
0,
|
||||
mbar,
|
||||
)
|
||||
|
||||
|
||||
def gemm(
|
||||
A: BufferLikeType,
|
||||
B: BufferLikeType,
|
||||
@@ -239,8 +182,19 @@ def gemm(
|
||||
Returns:
|
||||
tir.Call: A handle to the GEMM operation.
|
||||
"""
|
||||
impl = gemm_v1 if _env.use_gemm_v1() else gemm_v2
|
||||
return impl(A, B, C, transpose_A, transpose_B, policy, clear_accum, k_pack, mbar)
|
||||
return _gemm_impl(
|
||||
"tl.tileop.gemm",
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
transpose_A,
|
||||
transpose_B,
|
||||
policy,
|
||||
clear_accum,
|
||||
k_pack,
|
||||
0,
|
||||
mbar,
|
||||
)
|
||||
|
||||
|
||||
def wgmma_gemm(
|
||||
@@ -264,7 +218,7 @@ def wgmma_gemm(
|
||||
"""
|
||||
|
||||
return _gemm_impl(
|
||||
"tl.tileop.wgmma_gemm_py",
|
||||
"tl.tileop.wgmma_gemm",
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
@@ -306,7 +260,7 @@ def tcgen05_gemm(
|
||||
|
||||
ann = {"use_2cta": int(use_2cta)} if use_2cta else None
|
||||
return _gemm_impl(
|
||||
"tl.tileop.tcgen05_gemm_py",
|
||||
"tl.tileop.tcgen05_gemm",
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
from .base import GemmWarpPolicy # noqa: F401
|
||||
from .gemm import GemmPy # noqa: F401
|
||||
from .gemm import Gemm # noqa: F401
|
||||
from .gemm_sp import GemmSPPy # noqa: F401
|
||||
|
||||
@@ -17,15 +17,15 @@ from tilelang import _ffi_api
|
||||
from tilelang.utils.target import target_is_volta
|
||||
|
||||
|
||||
@tvm_ffi.register_global_func("tl.gemm_py.infer_layout")
|
||||
def gemm_py_infer_layout(gemm_py: GemmMMA, target: Target, thread_bounds: Range):
|
||||
@tvm_ffi.register_global_func("tl.gemm.infer_layout")
|
||||
def gemm_infer_layout(gemm: GemmMMA, target: Target, thread_bounds: Range):
|
||||
thread_nums = thread_bounds.extent
|
||||
return gemm_py.infer_layout(target, thread_nums)
|
||||
return gemm.infer_layout(target, thread_nums)
|
||||
|
||||
|
||||
@tvm_ffi.register_global_func("tl.gemm_py.lower")
|
||||
def gemm_py_lower(
|
||||
gemm_py: GemmMMA,
|
||||
@tvm_ffi.register_global_func("tl.gemm.lower")
|
||||
def gemm_lower(
|
||||
gemm: GemmMMA,
|
||||
layout_map,
|
||||
target: Target,
|
||||
thread_bounds: Range,
|
||||
@@ -33,12 +33,12 @@ def gemm_py_lower(
|
||||
mbar_phase_expr: tir.PrimExpr,
|
||||
):
|
||||
# We pass thread_bounds rather than thread_extents because tcgen5mma need to check this
|
||||
stmt = gemm_py.lower(layout_map, target, thread_bounds, thread_var, mbar_phase_expr)
|
||||
stmt = gemm.lower(layout_map, target, thread_bounds, thread_var, mbar_phase_expr)
|
||||
return stmt
|
||||
|
||||
|
||||
@tvm_ffi.register_object("tl.GemmPy")
|
||||
class GemmPy(Node, Scriptable):
|
||||
@tvm_ffi.register_object("tl.Gemm")
|
||||
class Gemm(Node, Scriptable):
|
||||
# FFI fields (LLVM/MLIR-style lowerCamel via reflection):
|
||||
# a, b, c, aPtr, bPtr, cPtr, m, n, k, transA, transB,
|
||||
# strideA, strideB, offsetA, offsetB, clearAccum, kPack, wgWait, policy
|
||||
@@ -159,7 +159,7 @@ class GemmPy(Node, Scriptable):
|
||||
Returns:
|
||||
GemmInst: The selected GEMM instruction type
|
||||
"""
|
||||
return GemmInst(_ffi_api.GemmPyGemmInst(self, int(thread_nums), target))
|
||||
return GemmInst(_ffi_api.GemmGetGemmInst(self, int(thread_nums), target))
|
||||
|
||||
def _get_implementation_class(self, gemm_inst: GemmInst, target: Target):
|
||||
"""Get the appropriate implementation class for the given GEMM instruction.
|
||||
|
||||
Reference in New Issue
Block a user