You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by cs...@apache.org on 2024/06/20 05:45:53 UTC
(tvm) branch main updated: [Dlight] Use 16x32 spatial x reduction thread extents in GEMV scheduling (#17082)
This is an automated email from the ASF dual-hosted git repository.
csullivan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 269a4f7e62 [Dlight] Use 16x32 spatial x reduction thread extents in GEMV scheduling (#17082)
269a4f7e62 is described below
commit 269a4f7e621cf6cf359fe223c3bbdc17986321bc
Author: Chris Sullivan <cs...@octo.ai>
AuthorDate: Wed Jun 19 22:45:47 2024 -0700
[Dlight] Use 16x32 spatial x reduction thread extents in GEMV scheduling (#17082)
Change to use 16x32 spatial x reduction thread extents regardless of workload size. This works around a lowering bug which I haven't tracked down yet.
Currently when the spatial dimension is larger than the reduction dimension, it uses a 4x64 thread layout. This implies two warps in the reduction dimension corresponding to blockDim.x=64. An illegal cuda instruction is encountered in the second warp during the __shfl_down_sync for the remainder portion of the computation (suspect interaction with rfactor). It appears the mask calculation used for this remainder shfl is incorrect and is causing the error. Specifically it occurs on the [...]
Changing the thread extents to 16x32 (one warp along the reduction dimension) works around the issue. It also improves performance for the tested shapes by ~10%.
---
python/tvm/dlight/gpu/gemv.py | 5 +-
.../python/codegen/test_target_codegen_cuda_fp8.py | 121 ++++++++++++++++++++
tests/python/dlight/test_gpu_gemv.py | 126 ++++++++++-----------
3 files changed, 185 insertions(+), 67 deletions(-)
diff --git a/python/tvm/dlight/gpu/gemv.py b/python/tvm/dlight/gpu/gemv.py
index ce1c5986e1..2bcb8563a2 100644
--- a/python/tvm/dlight/gpu/gemv.py
+++ b/python/tvm/dlight/gpu/gemv.py
@@ -445,10 +445,7 @@ class GEMV(GPUScheduleRule):
UNROLL = 256
SUPPORT_WARP_SHUFFLE = True
if isinstance(len_S, int):
- if len_S > len_R:
- TS, TR = 4, 64
- else:
- TS, TR = 16, 32
+ TS, TR = 16, 32
else:
TS, TR = 1, 64
elif target.kind.name == "metal":
diff --git a/tests/python/codegen/test_target_codegen_cuda_fp8.py b/tests/python/codegen/test_target_codegen_cuda_fp8.py
index adcb05839b..c22f3f01a8 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fp8.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fp8.py
@@ -33,6 +33,13 @@ from tvm.runtime import NDArray
from tvm.target import Target
from tvm.topi.utils import get_const_tuple
+from tvm.script import ir as I, relax as R, tir as T
+
+try:
+ import ml_dtypes
+except ImportError:
+ ml_dtypes = None
+
@tvm.testing.requires_cuda_compute_version(9)
def test_e4m3_conversions():
@@ -814,5 +821,119 @@ def test_const(dtype):
tvm.build(mod, target="cuda")
+num_experts = 8
+reduce_size = 1792
+spatial_size = 4096
+
+
+@tvm.testing.requires_cuda_compute_version(9)
+@pytest.mark.skipif(ml_dtypes is None, reason="Requires ml_dtypes to be installed")
+def test_moe_gemv_shfl_down_illegal_instr():
+ global num_experts
+ global reduce_size
+ global spatial_size
+
+ @I.ir_module
+ class SingleBatchMoE_float8_e4m3:
+ @T.prim_func(private=True)
+ def moe_dequantize_gemv(
+ x_handle: T.handle,
+ w: T.Buffer((num_experts, spatial_size, reduce_size), "e4m3_float8"),
+ scale: T.Buffer((1,), "float16"),
+ indptr: T.Buffer((1, 2), "int32"),
+ o: T.Buffer((2, spatial_size), "float16"),
+ ):
+ T.func_attr({"op_pattern": 4, "tir.noalias": T.bool(True)})
+ num_seq = T.int64()
+ x = T.match_buffer(x_handle, (num_seq, reduce_size), "float16")
+ for expert_id in T.thread_binding(2, thread="blockIdx.y"):
+ with T.block("gemv_o"):
+ e = T.axis.spatial(2, expert_id)
+ T.reads(
+ w[indptr[0, e], 0:spatial_size, 0:reduce_size],
+ indptr[0, e],
+ scale[0],
+ x[e, 0:reduce_size],
+ )
+ T.writes(o[e, 0:spatial_size])
+ y = T.alloc_buffer((spatial_size, reduce_size), "float16")
+ for i1, i2 in T.grid(spatial_size, reduce_size):
+ with T.block("dequantize"):
+ i, j = T.axis.remap("SS", [i1, i2])
+ T.reads(w[indptr[0, e], i, j], indptr[0, e], scale[0])
+ T.writes(y[i, j])
+ y[i, j] = T.Cast("float16", w[indptr[0, e], i, j]) * scale[0]
+ for i1, i2 in T.grid(spatial_size, reduce_size):
+ with T.block("gemv"):
+ i, j = T.axis.remap("SR", [i1, i2])
+ T.reads(x[e, j], y[i, j])
+ T.writes(o[e, i])
+ with T.init():
+ o[e, i] = T.float16(0)
+ o[e, i] = o[e, i] + x[e, j] * y[i, j]
+
+ @R.function
+ def main(
+ x: R.Tensor(("num_seq", reduce_size), dtype="float16"),
+ indptr: R.Tensor((1, 2), dtype="int32"),
+ weight: R.Tensor((num_experts, spatial_size, reduce_size), dtype="e4m3_float8"),
+ scale: R.Tensor((1,), dtype="float32"),
+ ) -> R.Tensor((2, spatial_size), dtype="float16"):
+ num_seq = T.int64()
+ R.func_attr({"num_input": 2})
+ cls = SingleBatchMoE_float8_e4m3
+ with R.dataflow():
+ astype: R.Tensor((1,), dtype="float16") = R.astype(scale, dtype="float16")
+ lv = R.call_tir(
+ cls.moe_dequantize_gemv,
+ (x, weight, astype, indptr),
+ out_sinfo=R.Tensor((2, spatial_size), dtype="float16"),
+ )
+ gv: R.Tensor((2, spatial_size), dtype="float16") = lv
+ R.output(gv)
+ return gv
+
+ def _pipeline(mod: tvm.ir.IRModule) -> tvm.ir.IRModule:
+ seq = tvm.transform.Sequential(
+ [
+ tvm.relax.transform.LegalizeOps(),
+ tvm.dlight.ApplyDefaultSchedule(
+ tvm.dlight.gpu.Matmul(),
+ tvm.dlight.gpu.GEMV(),
+ tvm.dlight.gpu.Reduction(),
+ tvm.dlight.gpu.GeneralReduction(),
+ tvm.dlight.gpu.Fallback(),
+ ),
+ ]
+ )
+ mod = seq(mod)
+ return mod
+
+ mod = SingleBatchMoE_float8_e4m3
+
+ target = tvm.target.Target("cuda")
+ with tvm.transform.PassContext(config={"relax.backend.use_cuda_graph": False}) and target:
+ mod = _pipeline(mod)
+ rt_mod = tvm.relax.build(mod, target=target)
+ dev = tvm.cuda(0)
+
+ x_data = np.zeros((1, reduce_size), dtype=np.float16)
+ x = tvm.nd.array(x_data, device=dev)
+
+ indptr_data = np.zeros((1, 2), dtype=np.int32)
+ indptr = tvm.nd.array(indptr_data, device=dev)
+
+ weight_data = np.zeros((num_experts, spatial_size, reduce_size), dtype="float8_e4m3fn")
+ weight = tvm.nd.array(weight_data, device=dev)
+
+ scale_data = np.zeros((1,), dtype=np.float32)
+ scale = tvm.nd.array(scale_data, device=dev)
+
+ vm = relax.VirtualMachine(rt_mod, dev)
+ # Ensure this runs without failure. Utilizing dlight thread extents TS, TR = 4, 64
+ # in GEMV scheduling will yield: CUDA: an illegal instruction was encountered.
+ vm["main"](x, indptr, weight, scale)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/dlight/test_gpu_gemv.py b/tests/python/dlight/test_gpu_gemv.py
index 20cb703f7f..79c5ab3a12 100644
--- a/tests/python/dlight/test_gpu_gemv.py
+++ b/tests/python/dlight/test_gpu_gemv.py
@@ -305,73 +305,73 @@ def test_decode_gemv1():
def expected(lv571: T.Buffer((22016, 512), "uint32"), lv572: T.Buffer((22016, 128), "float16"), lv1654: T.Buffer((1, 1, 4096), "float16"), var_NT_matmul_intermediate: T.Buffer((1, 1, 22016), "float16")):
T.func_attr({"tir.is_scheduled": 1, "tir.noalias": T.bool(True)})
# with T.block("root"):
- var_NT_matmul_intermediate_rf_local = T.alloc_buffer((256, 1, 1, 22016), "float16", scope="local")
- var_NT_matmul_intermediate_rf_local_1 = T.alloc_buffer((64, 1, 1, 22016), "float16", scope="local")
+ var_NT_matmul_intermediate_rf_local = T.alloc_buffer((128, 1, 1, 22016), "float16", scope="local")
+ var_NT_matmul_intermediate_rf_local_1 = T.alloc_buffer((32, 1, 1, 22016), "float16", scope="local")
lv571_local = T.alloc_buffer((22016, 512), "uint32", scope="local")
lv1654_shared = T.alloc_buffer((1, 1, 4096), "float16", scope="shared")
- for u_fused_ax0_fused_fused_0 in T.thread_binding(5504, thread="blockIdx.x"):
- for u_fused_ax0_fused_fused_1 in T.thread_binding(4, thread="threadIdx.y"):
- for ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 in T.thread_binding(64, thread="threadIdx.x"):
+ for u_fused_ax0_fused_fused_0 in T.thread_binding(1376, thread="blockIdx.x"):
+ for u_fused_ax0_fused_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
+ for ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 in T.thread_binding(32, thread="threadIdx.x"):
for ax0, ax1 in T.grid(1, 1):
- for ax2_0 in T.serial(2, annotations={"pragma_unroll_explicit": 256, "pragma_vectorize": 1}):
- for ax2_1 in T.thread_binding(4, thread="threadIdx.y"):
- for ax2_2 in T.thread_binding(64, thread="threadIdx.x"):
+ for ax2_0 in T.serial(1, annotations={"pragma_unroll_explicit": 256, "pragma_vectorize": 1}):
+ for ax2_1 in T.thread_binding(16, thread="threadIdx.y"):
+ for ax2_2 in T.thread_binding(32, thread="threadIdx.x"):
for ax2_3 in T.vectorized(8):
with T.block("lv1654_shared"):
v0, v1 = T.axis.remap("SS", [ax0, ax1])
- v2 = T.axis.spatial(4096, ax2_0 * 2048 + ax2_1 * 512 + ax2_2 * 8 + ax2_3)
+ v2 = T.axis.spatial(4096, ax2_0 * 4096 + ax2_1 * 256 + ax2_2 * 8 + ax2_3)
T.reads(lv1654[v0, v1, v2])
T.writes(lv1654_shared[v0, v1, v2])
lv1654_shared[v0, v1, v2] = lv1654[v0, v1, v2]
for u_fused_ax0_fused_fused_2_init in range(1):
for ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1_init in T.vectorized(4):
with T.block("NT_matmul_rf_init"):
- vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused = T.axis.spatial(256, ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1_init)
- v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 4 + u_fused_ax0_fused_fused_1 + u_fused_ax0_fused_fused_2_init)
+ vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused = T.axis.spatial(128, ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1_init)
+ v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 16 + u_fused_ax0_fused_fused_1 + u_fused_ax0_fused_fused_2_init)
T.reads()
T.writes(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0])
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0] = T.float16(0)
- for ax1_0_fused_ax1_1_fused_0 in T.serial(8, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0_fused_ax1_1_fused_0 in T.serial(16, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
for ax0_ax1_fused_0 in range(1):
for ax0_ax1_fused_1 in T.vectorized(1):
with T.block("lv571_local"):
- v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 4 + u_fused_ax0_fused_fused_1)
- v1 = T.axis.spatial(512, ax1_0_fused_ax1_1_fused_0 * 64 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0)
+ v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 16 + u_fused_ax0_fused_fused_1)
+ v1 = T.axis.spatial(512, ax1_0_fused_ax1_1_fused_0 * 32 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0)
T.reads(lv571[v0, v1])
T.writes(lv571_local[v0, v1])
lv571_local[v0, v1] = lv571[v0, v1]
for u_fused_ax0_fused_fused_2, ax1_0_fused_ax1_1_fused_2 in T.grid(1, 2):
for ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1 in T.vectorized(4):
with T.block("NT_matmul_rf_update"):
- vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused = T.axis.spatial(256, ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1)
- v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 4 + u_fused_ax0_fused_fused_1 + u_fused_ax0_fused_fused_2)
+ vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused = T.axis.spatial(128, ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1)
+ v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 16 + u_fused_ax0_fused_fused_1 + u_fused_ax0_fused_fused_2)
vax1_0_fused_ax1_1_fused_0, vax1_0_fused_ax1_1_fused_2 = T.axis.remap("RR", [ax1_0_fused_ax1_1_fused_0, ax1_0_fused_ax1_1_fused_2])
- T.reads(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0], lv1654_shared[0, 0, vax1_0_fused_ax1_1_fused_0 * 512 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 * 8 + vax1_0_fused_ax1_1_fused_2 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused % 4], lv571_local[v0, vax1_0_fused_ax1_1_fused_0 * 64 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 + va [...]
+ T.reads(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0], lv1654_shared[0, 0, vax1_0_fused_ax1_1_fused_0 * 256 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 * 8 + vax1_0_fused_ax1_1_fused_2 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused % 4], lv571_local[v0, vax1_0_fused_ax1_1_fused_0 * 32 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 + va [...]
T.writes(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0])
- var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0] = var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0] + lv1654_shared[0, 0, vax1_0_fused_ax1_1_fused_0 * 512 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 * 8 + vax1_0_fused_ax1_1_fused_2 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused % 4] * ((T.Cast("float1 [...]
- for ax2_fused_0 in T.thread_binding(4, thread="threadIdx.y"):
- for ax0 in T.thread_binding(64, thread="threadIdx.x"):
- for ax2_fused_1_0 in T.serial(1, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
- for ax2_fused_1_1 in T.vectorized(1):
+ var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0] = var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0] + lv1654_shared[0, 0, vax1_0_fused_ax1_1_fused_0 * 256 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 * 8 + vax1_0_fused_ax1_1_fused_2 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused % 4] * ((T.Cast("float1 [...]
+ for ax2_fused_0_ax2_fused_1_fused in T.thread_binding(16, thread="threadIdx.y"):
+ for ax0 in T.thread_binding(32, thread="threadIdx.x"):
+ for ax2_fused_2_0 in T.serial(1, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_2_1 in T.vectorized(1):
with T.block("NT_matmul_rf_init"):
- vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 = T.axis.spatial(64, ax0)
- v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 4 + ax2_fused_0 + ax2_fused_1_0 + ax2_fused_1_1)
+ vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 = T.axis.spatial(32, ax0)
+ v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 16 + ax2_fused_0_ax2_fused_1_fused + ax2_fused_2_0 + ax2_fused_2_1)
T.reads()
T.writes(var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0])
var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0] = T.float16(0)
for ax1 in range(4):
with T.block("NT_matmul_rf_update"):
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1 = T.axis.remap("SR", [ax0, ax1])
- v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 4 + ax2_fused_0 + ax2_fused_1_0 + ax2_fused_1_1)
+ v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 16 + ax2_fused_0_ax2_fused_1_fused + ax2_fused_2_0 + ax2_fused_2_1)
T.reads(var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0], var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1, 0, 0, v0])
T.writes(var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0])
var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0] = var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0] + var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1, 0, 0, v0]
- for ax1_fused_1 in range(1):
- for ax1_fused_0 in T.thread_binding(4, thread="threadIdx.y"):
- for ax0 in T.thread_binding(64, thread="threadIdx.x"):
+ for ax1_fused_2 in range(1):
+ for ax1_fused_0_ax1_fused_1_fused in T.thread_binding(16, thread="threadIdx.y"):
+ for ax0 in T.thread_binding(32, thread="threadIdx.x"):
with T.block("NT_matmul"):
- vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 = T.axis.reduce(64, ax0)
- v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 4 + ax1_fused_0 + ax1_fused_1)
+ vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 = T.axis.reduce(32, ax0)
+ v0 = T.axis.spatial(22016, u_fused_ax0_fused_fused_0 * 16 + ax1_fused_0_ax1_fused_1_fused + ax1_fused_2)
T.reads(var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0])
T.writes(var_NT_matmul_intermediate[0, 0, v0])
with T.init():
@@ -421,82 +421,82 @@ def test_decode_gemv2():
T.func_attr({"tir.is_scheduled": 1, "tir.noalias": T.bool(True)})
# with T.block("root"):
var_NT_matmul_intermediate_local = T.alloc_buffer((1, 1, 32000), "float16", scope="local")
- var_NT_matmul_intermediate_rf_local = T.alloc_buffer((256, 1, 1, 32000), "float16", scope="local")
- var_NT_matmul_intermediate_rf_local_1 = T.alloc_buffer((64, 1, 1, 32000), "float16", scope="local")
+ var_NT_matmul_intermediate_rf_local = T.alloc_buffer((128, 1, 1, 32000), "float16", scope="local")
+ var_NT_matmul_intermediate_rf_local_1 = T.alloc_buffer((32, 1, 1, 32000), "float16", scope="local")
lv771_local = T.alloc_buffer((32000, 512), "uint32", scope="local")
lv3216_shared = T.alloc_buffer((1, 1, 4096), "float16", scope="shared")
- for u_fused_ax0_fused_fused_0 in T.thread_binding(8000, thread="blockIdx.x"):
- for u_fused_ax0_fused_fused_1 in T.thread_binding(4, thread="threadIdx.y"):
- for ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 in T.thread_binding(64, thread="threadIdx.x"):
+ for u_fused_ax0_fused_fused_0 in T.thread_binding(2000, thread="blockIdx.x"):
+ for u_fused_ax0_fused_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
+ for ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 in T.thread_binding(32, thread="threadIdx.x"):
for ax0, ax1 in T.grid(1, 1):
- for ax2_0 in T.serial(2, annotations={"pragma_unroll_explicit": 256, "pragma_vectorize": 1}):
- for ax2_1 in T.thread_binding(4, thread="threadIdx.y"):
- for ax2_2 in T.thread_binding(64, thread="threadIdx.x"):
+ for ax2_0 in T.serial(1, annotations={"pragma_unroll_explicit": 256, "pragma_vectorize": 1}):
+ for ax2_1 in T.thread_binding(16, thread="threadIdx.y"):
+ for ax2_2 in T.thread_binding(32, thread="threadIdx.x"):
for ax2_3 in T.vectorized(8):
with T.block("lv3216_shared"):
v0, v1 = T.axis.remap("SS", [ax0, ax1])
- v2 = T.axis.spatial(4096, ax2_0 * 2048 + ax2_1 * 512 + ax2_2 * 8 + ax2_3)
+ v2 = T.axis.spatial(4096, ax2_0 * 4096 + ax2_1 * 256 + ax2_2 * 8 + ax2_3)
T.reads(lv3216[v0, v1, v2])
T.writes(lv3216_shared[v0, v1, v2])
lv3216_shared[v0, v1, v2] = lv3216[v0, v1, v2]
for u_fused_ax0_fused_fused_2_init in range(1):
for ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1_init in T.vectorized(4):
with T.block("NT_matmul_rf_init"):
- vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused = T.axis.spatial(256, ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1_init)
- v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 4 + u_fused_ax0_fused_fused_1 + u_fused_ax0_fused_fused_2_init)
+ vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused = T.axis.spatial(128, ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1_init)
+ v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 16 + u_fused_ax0_fused_fused_1 + u_fused_ax0_fused_fused_2_init)
T.reads()
T.writes(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0])
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0] = T.float16(0)
- for ax1_0_fused_ax1_1_fused_0 in T.serial(8, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0_fused_ax1_1_fused_0 in T.serial(16, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
for ax0_ax1_fused_0 in range(1):
for ax0_ax1_fused_1 in T.vectorized(1):
with T.block("lv771_local"):
- v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 4 + u_fused_ax0_fused_fused_1)
- v1 = T.axis.spatial(512, ax1_0_fused_ax1_1_fused_0 * 64 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0)
+ v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 16 + u_fused_ax0_fused_fused_1)
+ v1 = T.axis.spatial(512, ax1_0_fused_ax1_1_fused_0 * 32 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0)
T.reads(lv771[v0, v1])
T.writes(lv771_local[v0, v1])
lv771_local[v0, v1] = lv771[v0, v1]
for u_fused_ax0_fused_fused_2, ax1_0_fused_ax1_1_fused_2 in T.grid(1, 2):
for ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1 in T.vectorized(4):
with T.block("NT_matmul_rf_update"):
- vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused = T.axis.spatial(256, ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1)
- v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 4 + u_fused_ax0_fused_fused_1 + u_fused_ax0_fused_fused_2)
+ vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused = T.axis.spatial(128, ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1)
+ v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 16 + u_fused_ax0_fused_fused_1 + u_fused_ax0_fused_fused_2)
vax1_0_fused_ax1_1_fused_0, vax1_0_fused_ax1_1_fused_2 = T.axis.remap("RR", [ax1_0_fused_ax1_1_fused_0, ax1_0_fused_ax1_1_fused_2])
- T.reads(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0], lv3216_shared[0, 0, vax1_0_fused_ax1_1_fused_0 * 512 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 * 8 + vax1_0_fused_ax1_1_fused_2 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused % 4], lv771_local[v0, vax1_0_fused_ax1_1_fused_0 * 64 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 + va [...]
+ T.reads(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0], lv3216_shared[0, 0, vax1_0_fused_ax1_1_fused_0 * 256 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 * 8 + vax1_0_fused_ax1_1_fused_2 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused % 4], lv771_local[v0, vax1_0_fused_ax1_1_fused_0 * 32 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 + va [...]
T.writes(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0])
- var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0] = var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0] + lv3216_shared[0, 0, vax1_0_fused_ax1_1_fused_0 * 512 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 * 8 + vax1_0_fused_ax1_1_fused_2 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused % 4] * ((T.Cast("float1 [...]
- for ax2_fused_0 in T.thread_binding(4, thread="threadIdx.y"):
- for ax0 in T.thread_binding(64, thread="threadIdx.x"):
- for ax2_fused_1_0 in T.serial(1, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
- for ax2_fused_1_1 in T.vectorized(1):
+ var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0] = var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused, 0, 0, v0] + lv3216_shared[0, 0, vax1_0_fused_ax1_1_fused_0 * 256 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 * 8 + vax1_0_fused_ax1_1_fused_2 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused % 4] * ((T.Cast("float1 [...]
+ for ax2_fused_0_ax2_fused_1_fused in T.thread_binding(16, thread="threadIdx.y"):
+ for ax0 in T.thread_binding(32, thread="threadIdx.x"):
+ for ax2_fused_2_0 in T.serial(1, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_2_1 in T.vectorized(1):
with T.block("NT_matmul_rf_init"):
- vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 = T.axis.spatial(64, ax0)
- v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 4 + ax2_fused_0 + ax2_fused_1_0 + ax2_fused_1_1)
+ vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 = T.axis.spatial(32, ax0)
+ v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 16 + ax2_fused_0_ax2_fused_1_fused + ax2_fused_2_0 + ax2_fused_2_1)
T.reads()
T.writes(var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0])
var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0] = T.float16(0)
for ax1 in range(4):
with T.block("NT_matmul_rf_update"):
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1 = T.axis.remap("SR", [ax0, ax1])
- v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 4 + ax2_fused_0 + ax2_fused_1_0 + ax2_fused_1_1)
+ v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 16 + ax2_fused_0_ax2_fused_1_fused + ax2_fused_2_0 + ax2_fused_2_1)
T.reads(var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0], var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1, 0, 0, v0])
T.writes(var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0])
var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0] = var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0] + var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 * 4 + vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_1, 0, 0, v0]
- for ax1_fused_1 in range(1):
- for ax1_fused_0 in T.thread_binding(4, thread="threadIdx.y"):
- for ax0 in T.thread_binding(64, thread="threadIdx.x"):
+ for ax1_fused_2 in range(1):
+ for ax1_fused_0_ax1_fused_1_fused in T.thread_binding(16, thread="threadIdx.y"):
+ for ax0 in T.thread_binding(32, thread="threadIdx.x"):
with T.block("NT_matmul"):
- vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 = T.axis.reduce(64, ax0)
- v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 4 + ax1_fused_0 + ax1_fused_1)
+ vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 = T.axis.reduce(32, ax0)
+ v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 16 + ax1_fused_0_ax1_fused_1_fused + ax1_fused_2)
T.reads(var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0])
T.writes(var_NT_matmul_intermediate_local[0, 0, v0])
with T.init():
var_NT_matmul_intermediate_local[0, 0, v0] = T.float16(0)
var_NT_matmul_intermediate_local[0, 0, v0] = var_NT_matmul_intermediate_local[0, 0, v0] + var_NT_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0, 0, 0, v0]
- for ax0_fused_0 in T.thread_binding(4, thread="threadIdx.y"):
- for ax0_fused_1 in range(1):
+ for ax0_fused_0_ax0_fused_1_fused in T.thread_binding(16, thread="threadIdx.y"):
+ for ax0_fused_2 in range(1):
with T.block("compute"):
- v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 4 + ax0_fused_0 + ax0_fused_1)
+ v0 = T.axis.spatial(32000, u_fused_ax0_fused_fused_0 * 16 + ax0_fused_0_ax0_fused_1_fused + ax0_fused_2)
T.reads(var_NT_matmul_intermediate_local[0, 0, v0])
T.writes(p_output0_intermediate[0, 0, v0])
p_output0_intermediate[0, 0, v0] = T.Cast("float32", var_NT_matmul_intermediate_local[0, 0, v0])