chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,249 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=missing-function-docstring
|
||||
"""Tests for the priority=0 ``copy/fallback`` dispatch — scalar single-thread
|
||||
emit picked when every higher-priority variant rejects.
|
||||
|
||||
The cases here are *intentionally* shaped so ``gmem_smem`` rejects (region
|
||||
element count doesn't divide ``thread_cnt``) and ``reg`` / ``ld_stmatrix`` /
|
||||
... don't apply (scope pair mismatch). The dispatcher should land on
|
||||
fallback, the emit should pick one active thread, and the round-trip
|
||||
``A_gmem → A_smem → B_gmem`` should match.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
|
||||
# Force the fallback dispatch to register before any test compiles a kernel.
|
||||
# Without this import, in fresh pytest workers the `copy/fallback` variant
|
||||
# isn't yet registered when the dispatcher snapshots its registry.
|
||||
from tvm.tirx.cuda.operator.tile_primitive.copy import fallback as _fallback_module # noqa: F401
|
||||
from tvm.tirx.layout import S, TileLayout
|
||||
|
||||
|
||||
def _round_trip_shapes_and_threads():
|
||||
"""Cases where ``gmem_smem`` rejects on ``n_elements % thread_cnt``.
|
||||
|
||||
Per task: ``(scope, n_threads, shape, why_fallback)``. ``shape`` is small
|
||||
enough that scalar emit is fine, and chosen so the higher-priority
|
||||
variants can't accept it (size doesn't divide thread_cnt).
|
||||
"""
|
||||
return [
|
||||
# warp scope, 32 threads, 24 elements (4x6) → 24 % 32 != 0.
|
||||
("warp", 32, (4, 6), "4*6=24 ∤ 32"),
|
||||
# warp scope, 32 threads, 8 elements (1x8) → 8 % 32 != 0.
|
||||
("warp", 32, (1, 8), "1*8=8 ∤ 32"),
|
||||
# warpgroup scope, 128 threads, 24 elements (4x6) → 24 % 128 != 0.
|
||||
("warpgroup", 128, (4, 6), "4*6=24 ∤ 128"),
|
||||
# warpgroup scope, 128 threads, 32 elements (4x8) → 32 % 128 != 0.
|
||||
("warpgroup", 128, (4, 8), "4*8=32 ∤ 128"),
|
||||
# cta scope, 256 threads, 32 elements (4x8) → 32 % 256 != 0.
|
||||
("cta", 256, (4, 8), "4*8=32 ∤ 256"),
|
||||
# cta scope, 1024 threads, 64 elements (8x8) → 64 % 1024 != 0.
|
||||
# Mimics the test_partial_reduction sparse-write-back pattern.
|
||||
("cta", 1024, (8, 8), "8*8=64 ∤ 1024"),
|
||||
]
|
||||
|
||||
|
||||
def _build_round_trip_kernel(scope, n_threads, shape, dtype):
|
||||
"""``Tx.copy(A_smem, A); Tx.copy(B, A_smem)`` at the given scope. Both
|
||||
copies hit the same predicates; both should fall to fallback."""
|
||||
s_layout = TileLayout(S[shape])
|
||||
full = tuple(slice(0, d) for d in shape)
|
||||
|
||||
# Each scope variant inserts an explicit ``cta_sync`` between the two
|
||||
# copies — fallback's emit no longer sneaks one in, so the writer/reader
|
||||
# pair on ``A_smem`` would otherwise race.
|
||||
if scope == "warp":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.lane_id([32])
|
||||
T.thread_id([n_threads])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
Tx.warp.copy(A_smem[full], A[full])
|
||||
T.cuda.cta_sync()
|
||||
Tx.warp.copy(B[full], A_smem[full])
|
||||
|
||||
elif scope == "warpgroup":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([n_threads // 128])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
T.thread_id_in_wg([128])
|
||||
T.thread_id([n_threads])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
Tx.wg.copy(A_smem[full], A[full])
|
||||
T.cuda.cta_sync()
|
||||
Tx.wg.copy(B[full], A_smem[full])
|
||||
|
||||
elif scope == "cta":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warp_id([n_threads // 32])
|
||||
T.lane_id([32])
|
||||
T.thread_id([n_threads])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
Tx.cta.copy(A_smem[full], A[full])
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(B[full], A_smem[full])
|
||||
|
||||
else:
|
||||
raise ValueError(f"unsupported scope {scope!r}")
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
@pytest.mark.parametrize(
|
||||
"scope,n_threads,shape,why",
|
||||
[
|
||||
pytest.param(s, n, sh, w, id=f"{s}-{n}-{'x'.join(map(str, sh))}")
|
||||
for s, n, sh, w in _round_trip_shapes_and_threads()
|
||||
],
|
||||
)
|
||||
def test_fallback_round_trip(scope, n_threads, shape, why):
|
||||
"""End-to-end: compile + run + compare. Failure means either the
|
||||
dispatcher didn't pick fallback (silent crash earlier) or fallback's
|
||||
emit is wrong (mismatch on B vs A)."""
|
||||
del why
|
||||
dtype = "float32"
|
||||
kernel = _build_round_trip_kernel(scope, n_threads, shape, dtype)
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target, pytest.warns(UserWarning, match="copy/fallback"):
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
np_dtype = tvm.testing.np_dtype_from_str(dtype)
|
||||
A_np = tvm.testing.generate_random_array(dtype, shape)
|
||||
B_np = np.zeros(shape, dtype=np_dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
compiled(A, B)
|
||||
np.testing.assert_array_equal(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
def test_fallback_thread_scope():
|
||||
"""``T.thread()`` — single thread, no gate. Either ``gmem_smem`` picks
|
||||
it up (n_elements % 1 == 0) or ``fallback`` does — both end up emitting
|
||||
a sensible single-thread copy. We only check the round trip is correct,
|
||||
not which variant fired."""
|
||||
shape = (4, 6)
|
||||
dtype = "float32"
|
||||
s_layout = TileLayout(S[shape])
|
||||
full = tuple(slice(0, d) for d in shape)
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.thread_id([1])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
Tx.copy(A_smem[full], A[full])
|
||||
T.cuda.cta_sync()
|
||||
Tx.copy(B[full], A_smem[full])
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
np_dtype = tvm.testing.np_dtype_from_str(dtype)
|
||||
A_np = tvm.testing.generate_random_array(dtype, shape)
|
||||
B_np = np.zeros(shape, dtype=np_dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
compiled(A, B)
|
||||
np.testing.assert_array_equal(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
def test_fallback_emits_gate():
|
||||
"""Compiled CUDA source must contain a single-thread gate so only one
|
||||
active thread executes the scalar copy (not all of them, which would
|
||||
work but be racy + wasteful and indicate gate elision)."""
|
||||
shape = (4, 6)
|
||||
dtype = "float32"
|
||||
s_layout = TileLayout(S[shape])
|
||||
full = tuple(slice(0, d) for d in shape)
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warp_id([8]) # 256 threads => 8 warps
|
||||
T.lane_id([32])
|
||||
T.thread_id([256])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
Tx.cta.copy(A_smem[full], A[full])
|
||||
Tx.cta.copy(B[full], A_smem[full])
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target, pytest.warns(UserWarning, match="copy/fallback"):
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
src = "".join(im.inspect_source() for im in compiled.mod.imports)
|
||||
# The gate compiles to something like ``if (((int)threadIdx.x) == 0)``.
|
||||
# We don't pin the exact spelling; just require an equality predicate
|
||||
# against threadIdx.x somewhere in the source.
|
||||
assert "threadIdx.x" in src
|
||||
# At least one ``== 0`` (or ``== <literal>``) comparison must exist for
|
||||
# the single-thread gate.
|
||||
assert "== 0" in src, "fallback emit didn't produce a tid==0 gate; src:\n" + src[:2000]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,578 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=missing-function-docstring
|
||||
"""Round-trip tests for the ``gmem_smem`` copy dispatch (synthesized partition).
|
||||
|
||||
Pipeline: A_gmem --G2S--> A_smem --S2G--> B_gmem. If either direction is
|
||||
wrong the round trip leaves B mismatched against A.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout
|
||||
|
||||
|
||||
def _build_kernel(scope, n_threads, shape, dtype):
|
||||
s_layout = TileLayout(S[shape])
|
||||
full_slices = tuple(slice(0, d) for d in shape)
|
||||
|
||||
if scope == "warpgroup":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([n_threads // 128])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
T.thread_id_in_wg([128])
|
||||
T.thread_id([n_threads])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
Tx.wg.copy(A_smem[full_slices], A[full_slices])
|
||||
T.cuda.cta_sync()
|
||||
Tx.wg.copy(B[full_slices], A_smem[full_slices])
|
||||
|
||||
elif scope == "warp":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.lane_id([32])
|
||||
T.thread_id([n_threads])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
Tx.warp.copy(A_smem[full_slices], A[full_slices])
|
||||
T.cuda.cta_sync()
|
||||
Tx.warp.copy(B[full_slices], A_smem[full_slices])
|
||||
|
||||
elif scope == "cta":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warp_id([n_threads // 32])
|
||||
T.lane_id([32])
|
||||
T.thread_id([n_threads])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
Tx.cta.copy(A_smem[full_slices], A[full_slices])
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(B[full_slices], A_smem[full_slices])
|
||||
else:
|
||||
raise ValueError(f"unsupported scope {scope!r}")
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
# (scope, n_threads, shape) — shape chosen so total / T / vec_len > 1 with at
|
||||
# least one outer round, and total is divisible by T*vec_len.
|
||||
TASKS = [
|
||||
("warp", 32, (32, 32)), # 1024 total, T=32, vec 8 → outer 4
|
||||
("warp", 32, (32, 64)), # 2048 total, outer 8
|
||||
("warpgroup", 128, (128, 32)), # 4096, outer 4
|
||||
("warpgroup", 128, (128, 64)), # 8192, outer 8
|
||||
("warpgroup", 128, (256, 16)), # 4096, outer 4
|
||||
("cta", 256, (256, 32)), # 8192, T=256, vec 8 → outer 4
|
||||
("cta", 256, (512, 16)), # 8192, outer 4
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
@pytest.mark.parametrize(
|
||||
"scope,n_threads,shape",
|
||||
[pytest.param(*t, id=f"{t[0]}-{t[1]}-{'x'.join(map(str, t[2]))}") for t in TASKS],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["float16", "float32", "uint8"])
|
||||
def test_gmem_smem_roundtrip(scope, n_threads, shape, dtype):
|
||||
kernel = _build_kernel(scope, n_threads, shape, dtype)
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
np_dtype = tvm.testing.np_dtype_from_str(dtype)
|
||||
A_np = tvm.testing.generate_random_array(dtype, shape)
|
||||
B_np = np.zeros(shape, dtype=np_dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
compiled(A, B)
|
||||
np.testing.assert_array_equal(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Migrated from test_copy_sync.py: sync G↔S copy via the user-facing
|
||||
# Tx.copy() (which dispatches to gmem_smem).
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"task",
|
||||
[
|
||||
# A[0:128, 0:32] -> A_smem[0:128, 0:32] -> B[0:128, 0:32]
|
||||
(
|
||||
(128, 32),
|
||||
(128, 32),
|
||||
((0, 128), (0, 32)),
|
||||
32,
|
||||
TileLayout(S[128, 32]),
|
||||
TileLayout(S[128, 32]),
|
||||
TileLayout(S[128, 32]),
|
||||
),
|
||||
# A[32:64, 32:64] -> A_smem[0:32, 0:32] -> B[32:64, 32:64]
|
||||
(
|
||||
(64, 64),
|
||||
(32, 32),
|
||||
((32, 64), (32, 64)),
|
||||
32,
|
||||
TileLayout(S[64, 64]),
|
||||
TileLayout(S[64, 64]),
|
||||
TileLayout(S[32, 32]),
|
||||
),
|
||||
# A[0:1, 0:32, 0:32] -> A_smem[0:32, 0:32] -> B[0:1, 0:32, 0:32]
|
||||
(
|
||||
(4, 32, 32),
|
||||
(32, 32),
|
||||
((0, 1), (0, 32), (0, 32)),
|
||||
32,
|
||||
TileLayout(S[4, 32, 32]),
|
||||
TileLayout(S[4, 32, 32]),
|
||||
TileLayout(S[32, 32]),
|
||||
),
|
||||
# A[0:8, 0:8] -> A_smem[0:8, 0:8] -> B[0:8, 0:8]
|
||||
(
|
||||
(16, 16),
|
||||
(8, 8),
|
||||
((0, 8), (0, 8)),
|
||||
32,
|
||||
TileLayout(S[16, 16]),
|
||||
TileLayout(S[16, 16]),
|
||||
TileLayout(S[8, 8]),
|
||||
),
|
||||
# A[32:96, 256:512] -> A_smem[0:32, 0:256] -> B[32:96, 256:512] (swizzled)
|
||||
(
|
||||
(96, 512),
|
||||
(32, 256),
|
||||
((16, 48), (256, 512)),
|
||||
32,
|
||||
TileLayout(S[96, 512]),
|
||||
TileLayout(S[96, 512]),
|
||||
ComposeLayout(SwizzleLayout(3, 3, 3), TileLayout(S[8, 64]))
|
||||
.tile_to((16, 128), (8, 64))
|
||||
.tile_to((32, 256), (16, 128)),
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
@pytest.mark.parametrize(
|
||||
"dtype", ["int8", "float8_e4m3fn", "float8_e5m2", "float16", "bfloat16", "float32"]
|
||||
)
|
||||
@pytest.mark.parametrize("scope", ["cta", "thread"])
|
||||
def test_copy_g2s_s2g(task, dtype, scope):
|
||||
g_shape, s_shape, g_region, thread_cnt, layoutA, layoutB, layoutS = task
|
||||
|
||||
r_smem = tuple(slice(None) for _ in range(len(s_shape)))
|
||||
r_gmem = tuple(slice(g_region[i][0], g_region[i][1]) for i in range(len(g_shape)))
|
||||
|
||||
if scope == "thread":
|
||||
thread_cnt = 1
|
||||
|
||||
@T.prim_func
|
||||
def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, g_shape, dtype, layout=layoutA)
|
||||
B = T.match_buffer(B_ptr, g_shape, dtype, layout=layoutB)
|
||||
|
||||
T.device_entry()
|
||||
T.cta_id([2])
|
||||
T.thread_id([thread_cnt])
|
||||
|
||||
A_smem = T.alloc_buffer(s_shape, dtype, scope="shared", layout=layoutS)
|
||||
# `scope` is parametrized at runtime; select the scope namespace
|
||||
# dynamically (T.cta / T.thread) instead of a literal prefix.
|
||||
getattr(Tx, scope).copy(A_smem[r_smem], A[r_gmem])
|
||||
T.cuda.cta_sync()
|
||||
getattr(Tx, scope).copy(B[r_gmem], A_smem[r_smem])
|
||||
|
||||
np_dtype = tvm.testing.np_dtype_from_str(dtype)
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy_sync})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
np.random.seed(0)
|
||||
A_np = tvm.testing.generate_random_array(dtype, g_shape)
|
||||
B_np = np.zeros(g_shape, dtype=np_dtype)
|
||||
|
||||
B_ref = B_np.copy()
|
||||
B_ref[r_gmem] = A_np[r_gmem]
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
np.testing.assert_allclose(B_ref, B.numpy())
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Regression tests for known correctness gaps in ``align_layouts_gs``.
|
||||
#
|
||||
# These are intentionally algorithm-level (no GPU runtime) and currently
|
||||
# XFAIL; flipping them to passing is the contract for the upcoming swizzle /
|
||||
# alignment fix.
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _align(
|
||||
g_layout, g_shape, s_layout, s_shape, elem_bits, thread_cnt, g_region=None, s_region=None
|
||||
):
|
||||
from tvm.tirx.cuda.operator.tile_primitive.copy._common import align_layouts_gs
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
if g_region is None:
|
||||
g_region = [(0, d) for d in g_shape]
|
||||
if s_region is None:
|
||||
s_region = [(0, d) for d in s_shape]
|
||||
with target:
|
||||
return align_layouts_gs(
|
||||
g_layout,
|
||||
g_shape,
|
||||
g_region,
|
||||
s_layout,
|
||||
s_shape,
|
||||
s_region,
|
||||
elem_bits,
|
||||
thread_cnt,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.xfail(
|
||||
reason="align_layouts_gs ignores swizzle chunk size; "
|
||||
"_extract_tile strips the swizzle wrap before vec_len pick."
|
||||
)
|
||||
@pytest.mark.parametrize("per_element,expected_max_vec", [(2, 4), (1, 2), (0, 1)])
|
||||
def test_swizzled_smem_vec_len_must_fit_chunk(per_element, expected_max_vec):
|
||||
"""``SwizzleLayout(per_element, ...)`` keeps the bottom ``per_element``
|
||||
bits unswizzled. vec must stay within that chunk or it crosses an XOR
|
||||
boundary and reads/writes the wrong physical bytes."""
|
||||
shape = (32, 32) # 1024 fp16 elements total
|
||||
g_layout = TileLayout(S[shape])
|
||||
s_layout = ComposeLayout(SwizzleLayout(per_element, 3, 3), TileLayout(S[shape]))
|
||||
_g, _s, vec_len = _align(g_layout, shape, s_layout, shape, elem_bits=16, thread_cnt=32)
|
||||
chunk_elems = 1 << per_element
|
||||
assert vec_len <= chunk_elems, (
|
||||
f"vec_len={vec_len} crosses swizzle chunk size={chunk_elems} "
|
||||
f"(SwizzleLayout(per_element={per_element}, ...))"
|
||||
)
|
||||
|
||||
|
||||
def test_unaligned_strides_must_clamp_vec_len():
|
||||
"""G layout with row stride 20 (non-multiple of vec_len=8) → tid=2's
|
||||
base offset = 20 elements * 2 bytes = 40 bytes, which is not 16-byte
|
||||
aligned for a 128-bit vec ld/st (uint4 reinterpret crashes)."""
|
||||
shape = (2, 16)
|
||||
# row stride 20 (instead of 16) — leaves 4-elem gap between rows.
|
||||
g_layout = TileLayout(S[(2, 16) : (20, 1)])
|
||||
s_layout = TileLayout(S[(2, 16) : (20, 1)])
|
||||
_g, s_p, vec_len = _align(g_layout, shape, s_layout, shape, elem_bits=16, thread_cnt=4)
|
||||
# All non-vec strides must be multiples of vec_len so per-thread / per-round
|
||||
# starting offset stays vec-aligned. vec iter is s_p.shard[-1] (always
|
||||
# stride=1 by construction).
|
||||
for it in s_p.shard[:-1]:
|
||||
stride = int(it.stride)
|
||||
assert stride % vec_len == 0, (
|
||||
f"stride={stride} not a multiple of vec_len={vec_len}; "
|
||||
f"per-thread / per-round offset will be misaligned for the vec ld/st"
|
||||
)
|
||||
|
||||
|
||||
def test_unaligned_region_offset_must_clamp_vec_len():
|
||||
"""Slicing the gmem region at a non-vec-aligned column (e.g. col 3 in
|
||||
fp16) means the per-thread base offset starts at 3 elements = 6 bytes,
|
||||
which is not 16/8/4-byte aligned — vec_len must drop to 1."""
|
||||
shape = (4, 16)
|
||||
g_layout = TileLayout(S[(4, 32)]) # full buffer is 4x32 fp16
|
||||
s_layout = TileLayout(S[(4, 16)])
|
||||
# Take cols [3, 19) — start offset 3 (odd for any vec_len > 1 in fp16).
|
||||
g_region = [(0, 4), (3, 19)]
|
||||
s_region = [(0, 4), (0, 16)]
|
||||
_g, _s, vec_len = _align(
|
||||
g_layout,
|
||||
(4, 32),
|
||||
s_layout,
|
||||
(4, 16),
|
||||
elem_bits=16,
|
||||
thread_cnt=4,
|
||||
g_region=g_region,
|
||||
s_region=s_region,
|
||||
)
|
||||
assert 3 % vec_len == 0, (
|
||||
f"vec_len={vec_len} doesn't divide the region's starting column 3; "
|
||||
f"per-thread base offset will be misaligned for the vec ld/st"
|
||||
)
|
||||
|
||||
|
||||
def test_swizzled_smem_emit_must_be_swizzle_aware():
|
||||
"""Codegen-level: emitted S address should go through the SwizzleLayout's
|
||||
Apply so the XOR scrambling is honored. Currently emit uses
|
||||
``s_buf.ptr_to([0,..,0]) + linear_offset`` which only matches a
|
||||
non-swizzled storage layout."""
|
||||
import tvm
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout
|
||||
|
||||
shape = (128, 32)
|
||||
s_layout = ComposeLayout(SwizzleLayout(3, 3, 3), TileLayout(S[shape]))
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, "float16")
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([1])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
T.thread_id_in_wg([128])
|
||||
T.thread_id([128])
|
||||
A_smem = T.alloc_buffer(shape, "float16", scope="shared", layout=s_layout)
|
||||
Tx.wg.copy(A_smem[0:128, 0:32], A[0:128, 0:32])
|
||||
|
||||
# NB: pin sm_90 explicitly — the default cuda target falls back to sm_50
|
||||
# when no GPU is detected, which nvcc 13+ rejects. Codegen happens before
|
||||
# nvcc; if the whole tvm.compile pipeline fails, we never see the source.
|
||||
target = tvm.target.Target({"kind": "cuda", "arch": "sm_90"})
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
src = "".join(im.inspect_source() for im in compiled.mod.imports)
|
||||
|
||||
# If emit is swizzle-aware, two ways it shows up in the generated
|
||||
# CUDA:
|
||||
# 1. fallback path emits ``swizzle.apply(linear)``, which lowers
|
||||
# to a ``^`` (XOR) somewhere in the S-offset computation
|
||||
# (typically on a separate ``s_off_ptr[0] = ...`` line, not on
|
||||
# the ``tvm_builtin_pointer_offset`` line itself).
|
||||
# 2. fast path precomputes a ``signed_strides[N]`` register array
|
||||
# (one per binary outer iter), so each per-iter offset is a
|
||||
# sum of those strides — fingerprintable by the ``1 - 2 *``
|
||||
# sign-computation idiom emit_init writes.
|
||||
# XOR-less code paired with no signed_strides init means swizzle
|
||||
# was silently dropped.
|
||||
has_xor = "^" in src
|
||||
has_signed_strides_init = "1 - 2 *" in src or "(1 - 2 *" in src
|
||||
assert has_xor or has_signed_strides_init, (
|
||||
"emitted s_ptr address shows no swizzle handling — no XOR (fallback "
|
||||
"path) and no signed_strides init (fast path)"
|
||||
)
|
||||
|
||||
|
||||
def test_layout_permute_copy_preserves_smem_strides():
|
||||
"""Regression for the MMA-style K-tiled SMEM layout (``tcgen05_mma_ss_no_tma``):
|
||||
|
||||
``Tx.copy(A_smem, A)`` where A is plain row-major and A_smem uses the
|
||||
K-tiled MMA layout. The two layouts cover the same byte range but map
|
||||
``(i, j)`` to *different* physical offsets:
|
||||
|
||||
A : ``A[i, j] → i*K + j`` (row-major)
|
||||
A_smem : ``A_smem[i, j] → i*8 + (j//8)*1024 + (j%8)`` (K-tiled)
|
||||
|
||||
Earlier ``align_layouts_gs`` sorted+canonicalized BOTH sides
|
||||
independently. A_smem's three iters all chain by stride
|
||||
(8*128 == 1024, 1*8 == 8), so ``FuseContiguousShardIters`` collapsed
|
||||
A_smem to ``[(8192, 1)]`` — same as A's canonical form. The partition
|
||||
then synthesized identical ``s_p`` and ``g_p`` strides, ``apply``
|
||||
emitted ``tid*8`` on *both* sides, and the copy treated A_smem as
|
||||
row-major. MMA descriptors that re-read A_smem with the K-tiled
|
||||
formula then saw 99% wrong elements.
|
||||
|
||||
Fix: ``align_layouts_gs`` only sorts+canonicalizes G. S is grouped by
|
||||
G's pre-sort iter extents and permuted by G's stride-desc permutation;
|
||||
S's per-iter strides are preserved end-to-end.
|
||||
|
||||
This test asserts the structural property: ``s_p.apply`` for any
|
||||
``tid > 0`` must produce an offset that differs from G's row-major
|
||||
``tid * vec_len`` — proving S kept its K-tiled stride 1024.
|
||||
"""
|
||||
from tvm.tirx import Var as _TirVar
|
||||
from tvm.tirx.expr import IntImm as _IntImm
|
||||
|
||||
M, K = 128, 64
|
||||
# Plain row-major GMEM.
|
||||
g_layout = TileLayout(S[M, K])
|
||||
# K-tiled SMEM: 3D shape (M, K//8, 8) with strides (8, M*8, 1) — the
|
||||
# MMA descriptor's canonical SWIZZLE=0 layout from
|
||||
# tests/.../codegen/test_codegen_blackwell.py::test_tcgen05_mma_ss_no_tma.
|
||||
s_layout = TileLayout(S[(M, K // 8, 8) : (8, M * 8, 1)])
|
||||
|
||||
g_p, s_p, vec_len = _align(
|
||||
g_layout,
|
||||
(M, K),
|
||||
s_layout,
|
||||
(M, K),
|
||||
elem_bits=16,
|
||||
thread_cnt=128,
|
||||
)
|
||||
|
||||
# vec_len must reach 8 (fp16 → 128-bit vec ld/st).
|
||||
assert vec_len == 8, f"expected vec_len=8 for K-tiled fp16 MMA layout, got {vec_len}"
|
||||
|
||||
# S must keep at least one iter with stride 1024 (the K-tile jump
|
||||
# between 8-elem columns). After the fix, s_p.shard has 4 iters with
|
||||
# strides [128, 8, 1024, 1]; the old (broken) code collapsed to 3
|
||||
# iters all matching g_p's row-major strides [1024, 8, 1].
|
||||
s_strides = [int(it.stride) for it in s_p.shard]
|
||||
assert 1024 in s_strides, (
|
||||
f"s_p strides {s_strides} lost the K-tile stride 1024 — "
|
||||
f"align_layouts_gs collapsed A_smem to row-major and the copy "
|
||||
f"will write A_smem in the wrong layout"
|
||||
)
|
||||
|
||||
# Codegen-level check: s_p.apply on (f=0, tid, v=0) must depend on
|
||||
# ``tid % 8`` (the K-tile jump), not just ``tid * 8`` (row-major).
|
||||
# We pin this by evaluating apply for a couple of concrete tids.
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
apply_shape = [_IntImm("int32", 8), _IntImm("int32", 128), _IntImm("int32", 8)]
|
||||
tid_var = _TirVar("tid", "int32")
|
||||
s_off_expr = s_p.apply(
|
||||
_IntImm("int32", 0),
|
||||
tid_var,
|
||||
_IntImm("int32", 0),
|
||||
shape=apply_shape,
|
||||
)["m"]
|
||||
g_off_expr = g_p.apply(
|
||||
_IntImm("int32", 0),
|
||||
tid_var,
|
||||
_IntImm("int32", 0),
|
||||
shape=apply_shape,
|
||||
)["m"]
|
||||
|
||||
# G is row-major: g_off(tid) = tid * 8.
|
||||
# S is K-tiled : s_off(tid) = (tid // 8) * 8 + (tid % 8) * 1024.
|
||||
# For tid=1 the two MUST differ — they're identical iff S was
|
||||
# collapsed to row-major (the regression).
|
||||
from tvm.tirx import stmt_functor
|
||||
|
||||
analyzer = tvm.arith.Analyzer()
|
||||
s_off_at_1 = analyzer.simplify(
|
||||
stmt_functor.substitute(s_off_expr, {tid_var: _IntImm("int32", 1)})
|
||||
)
|
||||
g_off_at_1 = analyzer.simplify(
|
||||
stmt_functor.substitute(g_off_expr, {tid_var: _IntImm("int32", 1)})
|
||||
)
|
||||
assert int(s_off_at_1) == 1024, (
|
||||
f"s_p.apply at tid=1 produced offset {s_off_at_1}, expected 1024 "
|
||||
f"(K-tile jump). S was collapsed to row-major somewhere."
|
||||
)
|
||||
assert int(g_off_at_1) == 8, (
|
||||
f"g_p.apply at tid=1 produced offset {g_off_at_1}, expected 8 (row-major)"
|
||||
)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Fast-path firing test (positive). Pairs with the var_bounds wiring inside
|
||||
# ``gmem_smem._emit_gmem_smem``.
|
||||
#
|
||||
# Setup: warp-scope 32x64 fp16 G2S/S2G with 128b swizzled SMEM. The outer
|
||||
# iter stride is ``thread_cnt * vec_len = 32 * 8 = 256``, which puts the
|
||||
# binary-split bj's at {5, 6, 7} — well above the swizzle XOR region (so
|
||||
# Case 1.D, signed_stride = +T). The (C1) analyzer check
|
||||
# ``bit_bj(s_off // C) == 0`` needs the placeholder var bounded to
|
||||
# laneid ∈ [0, 32); the dispatch passes ``var_bounds`` so it can discharge,
|
||||
# recognizer accepts, and emit lowers to the
|
||||
# ``base_off + sum_j bit_j(f) · signed_strides[j]`` precomputed form.
|
||||
# ----------------------------------------------------------------------------
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
def test_gmem_smem_swizzle_fast_path_fires_with_var_bounds():
|
||||
"""Warp-scope 32x64 fp16 G2S/S2G with 128b swizzled SMEM. Fast path
|
||||
must fire: a 3-slot ``v_<n>[]`` signed_strides buffer + bit-select adds
|
||||
per outer iter, no per-iter ``swizzle.apply`` XOR splice in the hot path."""
|
||||
import re
|
||||
|
||||
swizzle = SwizzleLayout(3, 3, 3)
|
||||
shape = (32, 64)
|
||||
g_layout = TileLayout(S[shape])
|
||||
s_layout = ComposeLayout(swizzle, TileLayout(S[shape]))
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, "float16", layout=g_layout)
|
||||
B = T.match_buffer(B_ptr, shape, "float16", layout=g_layout)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.lane_id([32])
|
||||
T.thread_id([32])
|
||||
smem = T.alloc_buffer(shape, "float16", scope="shared", layout=s_layout)
|
||||
Tx.warp.copy(smem, A[:, :])
|
||||
T.cuda.cta_sync()
|
||||
Tx.warp.copy(B[:, :], smem)
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
ex = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
src = ex.mod.imports[0].inspect_source()
|
||||
|
||||
bitsel = re.findall(r"& 1\) \* v_\d+\[", src)
|
||||
v_decls = re.findall(r"alignas\(\d+\) int v_\d+\[(\d+)\]", src)
|
||||
assert bitsel, (
|
||||
"expected fast-path ``(bit & 1) * v_<n>[i]`` adds; if missing, "
|
||||
"var_bounds wiring may have regressed"
|
||||
)
|
||||
assert "3" in v_decls, (
|
||||
f"expected at least one 3-slot signed_strides buffer for bjs "
|
||||
f"[7, 6, 5]; got decl sizes {v_decls}"
|
||||
)
|
||||
|
||||
# Round-trip correctness.
|
||||
A_np = np.arange(32 * 64, dtype="float16").reshape(shape)
|
||||
B_np = np.zeros(shape, dtype="float16")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, device=dev)
|
||||
B = tvm.runtime.tensor(B_np, device=dev)
|
||||
ex(A, B)
|
||||
np.testing.assert_allclose(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,544 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=missing-function-docstring
|
||||
"""Round-trip tests for the ``ldstmatrix`` copy dispatch.
|
||||
|
||||
Pipeline:
|
||||
ld direction: A_gmem → A_smem (per-thread init) → R_local (T.copy dispatch
|
||||
under test) → B_gmem (per-thread write).
|
||||
st direction: A_gmem → R_local (per-thread init) → A_smem (T.copy dispatch
|
||||
under test) → B_gmem (per-thread write).
|
||||
|
||||
Both directions must round-trip ``A == B``. Layout strides are constructed
|
||||
so that:
|
||||
- trans=False S layout matches step-9's row-major spec (8→p, 4→2, num→q, 2→1).
|
||||
- trans=True S layout matches step-9's col-major spec (8→1, 4→2p, num→q, 2→p).
|
||||
|
||||
Uniform shape ``(*scope_outer, 8, 4, num, 2)`` is used for every num (including
|
||||
num=1, which gets an extent-1 placeholder for the num atom).
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout, laneid, tid_in_wg, tx
|
||||
|
||||
|
||||
def _compile_src(kernel):
|
||||
target = tvm.target.Target("cuda")
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
with target:
|
||||
compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
return compiled, compiled.mod.imports[0].inspect_source()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Layout builders.
|
||||
# ---------------------------------------------------------------------------
|
||||
def _r_layout_warp(num):
|
||||
return TileLayout(S[(8, 4, num, 2) : (4 @ laneid, 1 @ laneid, 2, 1)])
|
||||
|
||||
|
||||
def _r_layout_warpgroup(num):
|
||||
return TileLayout(S[(4, 8, 4, num, 2) : (32 @ tid_in_wg, 4 @ tid_in_wg, 1 @ tid_in_wg, 2, 1)])
|
||||
|
||||
|
||||
def _r_layout_cta(num):
|
||||
return TileLayout(S[(4, 8, 4, num, 2) : (32 @ tx, 4 @ tx, 1 @ tx, 2, 1)])
|
||||
|
||||
|
||||
def _s_layout_warp(num, trans):
|
||||
if not trans:
|
||||
return TileLayout(S[(8, 4, num, 2) : (num * 8, 2, 8, 1)])
|
||||
return TileLayout(S[(8, 4, num, 2) : (1, 2 * num * 8, 8, num * 8)])
|
||||
|
||||
|
||||
def _s_layout_warpgroup_or_cta(num, trans):
|
||||
if not trans:
|
||||
return TileLayout(S[(4, 8, 4, num, 2) : (64 * num, num * 8, 2, 8, 1)])
|
||||
return TileLayout(S[(4, 8, 4, num, 2) : (64 * num, 1, 16 * num, 8, 8 * num)])
|
||||
|
||||
|
||||
# 128b swizzle for fp16 (p=3 ⇒ 8 fp16 chunk; sw=at=3 ⇒ 8-row swizzle period).
|
||||
_SWIZZLE_128B = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3)
|
||||
|
||||
|
||||
def _maybe_wrap_swizzle(tile_layout, enable: bool):
|
||||
if not enable:
|
||||
return tile_layout
|
||||
return ComposeLayout(_SWIZZLE_128B, tile_layout)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Warp scope kernel builder.
|
||||
# ---------------------------------------------------------------------------
|
||||
def _build_warp_kernel(num, direction, trans, swizzle=False):
|
||||
r_layout = _r_layout_warp(num)
|
||||
s_layout = _maybe_wrap_swizzle(_s_layout_warp(num, trans), swizzle)
|
||||
s_shape = (8, 4, num, 2)
|
||||
full = (slice(0, 8), slice(0, 4), slice(0, num), slice(0, 2))
|
||||
M, N = 8, num * 8
|
||||
|
||||
def _coord(row, cp, t, w):
|
||||
# Map per-thread layout coord (row, cp, t, w) to gmem (row, col).
|
||||
if not trans:
|
||||
return row, t * 8 + cp * 2 + w
|
||||
return cp * 2 + w, row + t * 8
|
||||
|
||||
# fmt: off
|
||||
if direction == "ld":
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (M, N), "float16")
|
||||
B = T.match_buffer(B_ptr, (M, N), "float16")
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.lane_id([32])
|
||||
tid = T.thread_id([32])
|
||||
A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
|
||||
row = tid // 4
|
||||
cp = tid % 4
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(row, cp, t, w)
|
||||
A_smem[row, cp, t, w] = A[gr, gc]
|
||||
T.cuda.cta_sync()
|
||||
R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
|
||||
Tx.warp.copy(R_local[full], A_smem[full])
|
||||
r_view = R_local.local()
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(row, cp, t, w)
|
||||
B[gr, gc] = r_view[t * 2 + w]
|
||||
else: # direction == "st"
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (M, N), "float16")
|
||||
B = T.match_buffer(B_ptr, (M, N), "float16")
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.lane_id([32])
|
||||
tid = T.thread_id([32])
|
||||
A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
|
||||
row = tid // 4
|
||||
cp = tid % 4
|
||||
R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
|
||||
r_view = R_local.local()
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(row, cp, t, w)
|
||||
r_view[t * 2 + w] = A[gr, gc]
|
||||
Tx.warp.copy(A_smem[full], R_local[full])
|
||||
T.cuda.cta_sync()
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(row, cp, t, w)
|
||||
B[gr, gc] = A_smem[row, cp, t, w]
|
||||
# fmt: on
|
||||
return kernel, (M, N)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Warpgroup scope kernel builder. 4 warps stacked vertically.
|
||||
# ---------------------------------------------------------------------------
|
||||
def _build_warpgroup_kernel(num, direction, trans, swizzle=False):
|
||||
r_layout = _r_layout_warpgroup(num)
|
||||
s_layout = _maybe_wrap_swizzle(_s_layout_warpgroup_or_cta(num, trans), swizzle)
|
||||
s_shape = (4, 8, 4, num, 2)
|
||||
full = (slice(0, 4), slice(0, 8), slice(0, 4), slice(0, num), slice(0, 2))
|
||||
M, N = 32, num * 8
|
||||
|
||||
def _coord(wid, row, cp, t, w):
|
||||
if not trans:
|
||||
return wid * 8 + row, t * 8 + cp * 2 + w
|
||||
return wid * 8 + cp * 2 + w, row + t * 8
|
||||
|
||||
# fmt: off
|
||||
if direction == "ld":
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (M, N), "float16")
|
||||
B = T.match_buffer(B_ptr, (M, N), "float16")
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([1])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
T.thread_id_in_wg([128])
|
||||
tid = T.thread_id([128])
|
||||
A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
|
||||
wid = tid // 32
|
||||
lid = tid % 32
|
||||
row = lid // 4
|
||||
cp = lid % 4
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(wid, row, cp, t, w)
|
||||
A_smem[wid, row, cp, t, w] = A[gr, gc]
|
||||
T.cuda.cta_sync()
|
||||
R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
|
||||
Tx.wg.copy(R_local[full], A_smem[full])
|
||||
r_view = R_local.local()
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(wid, row, cp, t, w)
|
||||
B[gr, gc] = r_view[t * 2 + w]
|
||||
else:
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (M, N), "float16")
|
||||
B = T.match_buffer(B_ptr, (M, N), "float16")
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([1])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
T.thread_id_in_wg([128])
|
||||
tid = T.thread_id([128])
|
||||
A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
|
||||
wid = tid // 32
|
||||
lid = tid % 32
|
||||
row = lid // 4
|
||||
cp = lid % 4
|
||||
R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
|
||||
r_view = R_local.local()
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(wid, row, cp, t, w)
|
||||
r_view[t * 2 + w] = A[gr, gc]
|
||||
Tx.wg.copy(A_smem[full], R_local[full])
|
||||
T.cuda.cta_sync()
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(wid, row, cp, t, w)
|
||||
B[gr, gc] = A_smem[wid, row, cp, t, w]
|
||||
# fmt: on
|
||||
return kernel, (M, N)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CTA scope kernel builder. Same geometry as warpgroup, but R uses ``tx``.
|
||||
# ---------------------------------------------------------------------------
|
||||
def _build_cta_kernel(num, direction, trans, swizzle=False):
|
||||
r_layout = _r_layout_cta(num)
|
||||
s_layout = _maybe_wrap_swizzle(_s_layout_warpgroup_or_cta(num, trans), swizzle)
|
||||
s_shape = (4, 8, 4, num, 2)
|
||||
full = (slice(0, 4), slice(0, 8), slice(0, 4), slice(0, num), slice(0, 2))
|
||||
M, N = 32, num * 8
|
||||
|
||||
def _coord(wid, row, cp, t, w):
|
||||
if not trans:
|
||||
return wid * 8 + row, t * 8 + cp * 2 + w
|
||||
return wid * 8 + cp * 2 + w, row + t * 8
|
||||
|
||||
# fmt: off
|
||||
if direction == "ld":
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (M, N), "float16")
|
||||
B = T.match_buffer(B_ptr, (M, N), "float16")
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warp_id([4])
|
||||
T.lane_id([32])
|
||||
tid = T.thread_id([128])
|
||||
A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
|
||||
wid = tid // 32
|
||||
lid = tid % 32
|
||||
row = lid // 4
|
||||
cp = lid % 4
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(wid, row, cp, t, w)
|
||||
A_smem[wid, row, cp, t, w] = A[gr, gc]
|
||||
T.cuda.cta_sync()
|
||||
R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
|
||||
Tx.cta.copy(R_local[full], A_smem[full])
|
||||
r_view = R_local.local()
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(wid, row, cp, t, w)
|
||||
B[gr, gc] = r_view[t * 2 + w]
|
||||
else:
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (M, N), "float16")
|
||||
B = T.match_buffer(B_ptr, (M, N), "float16")
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warp_id([4])
|
||||
T.lane_id([32])
|
||||
tid = T.thread_id([128])
|
||||
A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
|
||||
wid = tid // 32
|
||||
lid = tid % 32
|
||||
row = lid // 4
|
||||
cp = lid % 4
|
||||
R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
|
||||
r_view = R_local.local()
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(wid, row, cp, t, w)
|
||||
r_view[t * 2 + w] = A[gr, gc]
|
||||
Tx.cta.copy(A_smem[full], R_local[full])
|
||||
T.cuda.cta_sync()
|
||||
for t in range(num):
|
||||
for w in range(2):
|
||||
gr, gc = _coord(wid, row, cp, t, w)
|
||||
B[gr, gc] = A_smem[wid, row, cp, t, w]
|
||||
# fmt: on
|
||||
return kernel, (M, N)
|
||||
|
||||
|
||||
_BUILDERS = {
|
||||
"warp": _build_warp_kernel,
|
||||
"warpgroup": _build_warpgroup_kernel,
|
||||
"cta": _build_cta_kernel,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scope", ["warp", "warpgroup", "cta"])
|
||||
@pytest.mark.parametrize("trans", [False, True])
|
||||
@pytest.mark.parametrize("direction", ["ld", "st"])
|
||||
@pytest.mark.parametrize("num", [1, 2, 4])
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
def test_ldstmatrix(scope, trans, direction, num):
|
||||
kernel, (M, N) = _BUILDERS[scope](num, direction, trans)
|
||||
compiled, src = _compile_src(kernel)
|
||||
|
||||
inst = "ldmatrix" if direction == "ld" else "stmatrix"
|
||||
trans_inst = ".trans" if trans else ""
|
||||
expected = f"{inst}.sync.aligned.m8n8.x{num}{trans_inst}.shared.b16"
|
||||
assert expected in src, f"{expected} not emitted; src=\n{src}"
|
||||
|
||||
A_np = np.arange(M * N, dtype="float16").reshape(M, N)
|
||||
B_np = np.zeros((M, N), dtype="float16")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, device=dev)
|
||||
B = tvm.runtime.tensor(B_np, device=dev)
|
||||
compiled(A, B)
|
||||
np.testing.assert_allclose(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Swizzled-S round-trip. Verifies the ldstmatrix dispatch's swizzle fast
|
||||
# path (when recognized) and slow path (fallback) both produce correct
|
||||
# A == B. The 128b swizzle (p=sw=at=3) is the most common fp16 SMEM
|
||||
# swizzle; with it the dispatch's per-tile S offset goes through
|
||||
# ``swizzle.apply`` (or its precomputed signed-stride lowering) per mm.
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.parametrize("scope", ["warp", "warpgroup", "cta"])
|
||||
@pytest.mark.parametrize("trans", [False, True])
|
||||
@pytest.mark.parametrize("direction", ["ld", "st"])
|
||||
@pytest.mark.parametrize("num", [1, 2, 4])
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
def test_ldstmatrix_swizzle(scope, trans, direction, num):
|
||||
kernel, (M, N) = _BUILDERS[scope](num, direction, trans, swizzle=True)
|
||||
compiled, src = _compile_src(kernel)
|
||||
|
||||
inst = "ldmatrix" if direction == "ld" else "stmatrix"
|
||||
trans_inst = ".trans" if trans else ""
|
||||
expected = f"{inst}.sync.aligned.m8n8.x{num}{trans_inst}.shared.b16"
|
||||
assert expected in src, f"{expected} not emitted; src=\n{src}"
|
||||
|
||||
A_np = np.arange(M * N, dtype="float16").reshape(M, N)
|
||||
B_np = np.zeros((M, N), dtype="float16")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, device=dev)
|
||||
B = tvm.runtime.tensor(B_np, device=dev)
|
||||
compiled(A, B)
|
||||
np.testing.assert_allclose(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multi-iter outer (m_outer > 1) fast-path round-trip. The existing 128b
|
||||
# swizzle tests above all have m_outer = 1 (only base_off, no signed_strides
|
||||
# bits). These two cover the non-trivial outer iter case:
|
||||
#
|
||||
# pow2 case (32x64): R outermost mem ext = 4 (pow2).
|
||||
# m_outer iters on S 6D seg 3 = [(4, 512), (2, 32)].
|
||||
# Binary split → bjs [7, 6, 2]. All BitIter.
|
||||
# signed_strides buffer has 3 slots.
|
||||
#
|
||||
# linear case (40x64): R outermost mem ext = 5 (non-pow2). Stride lands
|
||||
# the outermost S 6D seg 3 iter at (5, 512). 512 is
|
||||
# exactly 2^(p+at+sw) = swizzle period → Case 1.D pure,
|
||||
# so the LinearIter relaxation accepts it.
|
||||
# outer_iters = [LinearIter(5, 512), BitIter(2, ...)].
|
||||
# Inner BitIter (bj=2 Case 1.A) is the only slot in
|
||||
# signed_strides; the outer LinearIter contributes
|
||||
# ``c * 512`` per mm as a compile-time constant.
|
||||
# ---------------------------------------------------------------------------
|
||||
def _build_multi_iter_kernel(outer_ext: int):
|
||||
"""Warp + R=(outer_ext, 8, 2, 4, 4, 2):(16, 4@laneid, 8, 2, 1@laneid, 1)
|
||||
+ 333 swizzle on S. Mem strides 16/8/2/1 on extents outer_ext/2/4/2 are
|
||||
bijective for outer_ext ∈ {4, 5}: max = (outer_ext-1)*16 + 8 + 6 + 1
|
||||
= outer_ext*16 - 1, matching extent product = 16*outer_ext."""
|
||||
shape = (outer_ext, 8, 2, 4, 4, 2)
|
||||
r_layout = TileLayout(S[shape : (16, 4 @ laneid, 8, 2, 1 @ laneid, 1)])
|
||||
s_layout = SwizzleLayout(3, 3, 3)
|
||||
full = tuple(slice(0, e) for e in shape)
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, "float16")
|
||||
B = T.match_buffer(B_ptr, shape, "float16")
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.lane_id([32])
|
||||
tid = T.thread_id([32])
|
||||
A_smem = T.alloc_buffer(shape, "float16", scope="shared", layout=s_layout)
|
||||
for a in range(outer_ext):
|
||||
for c in range(2):
|
||||
for d in range(4):
|
||||
for e in range(2):
|
||||
A_smem[a, tid // 4, c, d, tid % 4, e] = A[a, tid // 4, c, d, tid % 4, e]
|
||||
T.cuda.cta_sync()
|
||||
R_local = T.alloc_buffer(shape, "float16", scope="local", layout=r_layout)
|
||||
Tx.warp.copy(R_local[full], A_smem[full])
|
||||
r_view = R_local.local()
|
||||
for a in range(outer_ext):
|
||||
for c in range(2):
|
||||
for d in range(4):
|
||||
for e in range(2):
|
||||
B[a, tid // 4, c, d, tid % 4, e] = r_view[a * 16 + c * 8 + d * 2 + e]
|
||||
|
||||
return kernel, shape
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
def test_ldstmatrix_swizzle_multi_iter_pow2():
|
||||
"""32x64 fp16 warp; outer m_outer split into multiple BitIters (no
|
||||
LinearIter). Fast path must fire with a 3-slot signed_strides buffer."""
|
||||
import re
|
||||
|
||||
kernel, shape = _build_multi_iter_kernel(outer_ext=4)
|
||||
compiled, src = _compile_src(kernel)
|
||||
assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src
|
||||
|
||||
# Fast-path fingerprint: 3-slot signed_strides + bit-select uses.
|
||||
assert re.search(r"alignas\(\d+\) int v_\d+\[3\]", src), (
|
||||
"expected 3-slot signed_strides buffer for bjs [7, 6, 2]"
|
||||
)
|
||||
bitsel = re.findall(r"& 1\) \* v_\d+\[", src)
|
||||
assert bitsel, "fast-path bit-select pattern '& 1) * v_<n>[' missing"
|
||||
|
||||
n_elem = 1
|
||||
for e in shape:
|
||||
n_elem *= e
|
||||
A_np = np.arange(n_elem, dtype="float16").reshape(shape)
|
||||
B_np = np.zeros(shape, dtype="float16")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, device=dev)
|
||||
B = tvm.runtime.tensor(B_np, device=dev)
|
||||
compiled(A, B)
|
||||
np.testing.assert_allclose(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
def test_ldstmatrix_tcgen05_warpgroup_atom_emits_ldmatrix():
|
||||
"""Regression: a warpgroup ``.16x256b`` tcgen05 register atom loaded from a
|
||||
128B-swizzled SMEM tile must dispatch to ``ldmatrix.x4``.
|
||||
|
||||
The atom's laneid is split across two shard iters (stride 4 and stride 1)
|
||||
plus ``wid_in_wg`` — it only fuses to a single ``tid_in_wg`` thread axis
|
||||
after the slice. With the redundant pre-slice ``canonicalize()`` removed,
|
||||
Step 2 relies on ``TileLayout.Slice``'s built-in global pre-canonicalize, so
|
||||
the slice no longer leaves an ill-formed split-laneid sub-layout that
|
||||
``GetScope`` would reject (which silently fell back to a scalar reg path).
|
||||
Compile-only (no GPU): asserts the instruction appears in generated source.
|
||||
"""
|
||||
from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout
|
||||
from tvm.tirx.layout import tcgen05_atom_layout
|
||||
|
||||
m, k = 64, 64
|
||||
# bf16 payload carrying the *fp32* atom layout (mirrors the tf32-prenorm
|
||||
# cast warp's ``a_bf16``: one element per 32-bit slot).
|
||||
reg_layout = tcgen05_atom_layout("16x256b", (m, k), "float32")
|
||||
smem_layout = mma_shared_layout("bfloat16", 3, (m, k))
|
||||
|
||||
@T.prim_func
|
||||
def kernel(s_ptr: T.handle) -> None:
|
||||
smem = T.match_buffer(s_ptr, (m, k), "bfloat16", scope="shared", layout=smem_layout)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([1])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
T.thread_id_in_wg([128])
|
||||
a_reg = T.alloc_buffer((m, k), "bfloat16", scope="local", layout=reg_layout)
|
||||
Tx.wg.copy(a_reg, smem)
|
||||
|
||||
_, src = _compile_src(kernel)
|
||||
assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src, (
|
||||
f"warpgroup tcgen05 atom did not emit ldmatrix.x4; src=\n{src}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
def test_ldstmatrix_swizzle_multi_iter_linear():
|
||||
"""40x64 fp16 warp; outer ext=5 is non-pow2 but stride lands on swizzle
|
||||
period (Case 1.D pure) so the LinearIter relaxation fires. Pattern has
|
||||
a 1-slot signed_strides (inner BitIter bj=2 Case 1.A); outer iter
|
||||
contributes ``c * 512`` per mm as a compile-time constant."""
|
||||
import re
|
||||
|
||||
kernel, shape = _build_multi_iter_kernel(outer_ext=5)
|
||||
compiled, src = _compile_src(kernel)
|
||||
assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src
|
||||
|
||||
# Fast-path fingerprint: 1-slot signed_strides (just the inner BitIter).
|
||||
assert re.search(r"alignas\(\d+\) int v_\d+\[1\]", src), (
|
||||
"expected 1-slot signed_strides buffer (only the inner Case-1.A bj=2)"
|
||||
)
|
||||
bitsel = re.findall(r"& 1\) \* v_\d+\[", src)
|
||||
assert bitsel, "fast-path bit-select pattern missing"
|
||||
|
||||
n_elem = 1
|
||||
for e in shape:
|
||||
n_elem *= e
|
||||
A_np = np.arange(n_elem, dtype="float16").reshape(shape)
|
||||
B_np = np.zeros(shape, dtype="float16")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, device=dev)
|
||||
B = tvm.runtime.tensor(B_np, device=dev)
|
||||
compiled(A, B)
|
||||
np.testing.assert_allclose(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
@@ -0,0 +1,629 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=missing-function-docstring
|
||||
"""Round-trip tests for the ``reg`` copy dispatch.
|
||||
|
||||
R = per-thread local (register). The dispatch handles round-trips between R
|
||||
and any non-R buffer (``shared*`` or ``global``); ``non_r_scope`` parametrize
|
||||
toggles which side is exercised.
|
||||
|
||||
Self-contained: each thread direct-stores its row into the non-R buffer (no
|
||||
G2S / G2L dispatch needed because each thread writes its own address), the
|
||||
dispatch does the inbound copy into R and the outbound copy back, then each
|
||||
thread reads its row into ``B``. Round-trip mismatch ⇒ at least one direction
|
||||
is wrong.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx.layout import S, TileLayout, laneid, tid_in_wg, tx
|
||||
|
||||
|
||||
def _r_layout(scope, shape):
|
||||
if scope == "warpgroup":
|
||||
return TileLayout(S[shape : (1 @ tid_in_wg, 1)])
|
||||
if scope == "warp":
|
||||
return TileLayout(S[shape : (1 @ laneid, 1)])
|
||||
if scope == "cta":
|
||||
return TileLayout(S[shape : (1 @ tx, 1)])
|
||||
raise ValueError(f"unsupported scope {scope!r}")
|
||||
|
||||
|
||||
def _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope):
|
||||
"""Build a kernel that round-trips data through R via ``non_r_scope``.
|
||||
|
||||
``non_r_scope == "shared"``: ``A_smem`` is allocated inside the kernel.
|
||||
Kernel signature: ``kernel(B_ptr)``.
|
||||
|
||||
``non_r_scope == "global"``: a separate gmem ``A`` is the staging area.
|
||||
Kernel signature: ``kernel(A_ptr, B_ptr)``.
|
||||
"""
|
||||
shape = (n_threads, k)
|
||||
full_slices = (slice(0, n_threads), slice(0, k))
|
||||
r_layout = _r_layout(scope, shape)
|
||||
|
||||
if non_r_scope == "shared":
|
||||
s_layout = TileLayout(S[shape])
|
||||
|
||||
if scope == "warpgroup":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(B_ptr: T.handle) -> None:
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([n_threads // 128])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
T.thread_id_in_wg([128])
|
||||
tid = T.thread_id([n_threads])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
for kk in range(k):
|
||||
A_smem[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
|
||||
T.cuda.cta_sync()
|
||||
R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
|
||||
Tx.wg.copy(R_local[full_slices], A_smem[full_slices])
|
||||
for kk in range(k):
|
||||
A_smem[tid, kk] = T.cast(0, dtype)
|
||||
T.cuda.cta_sync()
|
||||
Tx.wg.copy(A_smem[full_slices], R_local[full_slices])
|
||||
T.cuda.cta_sync()
|
||||
for kk in range(k):
|
||||
B[tid, kk] = A_smem[tid, kk]
|
||||
|
||||
elif scope == "warp":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(B_ptr: T.handle) -> None:
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.lane_id([32])
|
||||
tid = T.thread_id([n_threads])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
for kk in range(k):
|
||||
A_smem[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
|
||||
T.cuda.cta_sync()
|
||||
R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
|
||||
Tx.warp.copy(R_local[full_slices], A_smem[full_slices])
|
||||
for kk in range(k):
|
||||
A_smem[tid, kk] = T.cast(0, dtype)
|
||||
T.cuda.cta_sync()
|
||||
Tx.warp.copy(A_smem[full_slices], R_local[full_slices])
|
||||
T.cuda.cta_sync()
|
||||
for kk in range(k):
|
||||
B[tid, kk] = A_smem[tid, kk]
|
||||
|
||||
elif scope == "cta":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(B_ptr: T.handle) -> None:
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warp_id([n_threads // 32])
|
||||
T.lane_id([32])
|
||||
tid = T.thread_id([n_threads])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
|
||||
for kk in range(k):
|
||||
A_smem[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
|
||||
T.cuda.cta_sync()
|
||||
R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
|
||||
Tx.cta.copy(R_local[full_slices], A_smem[full_slices])
|
||||
for kk in range(k):
|
||||
A_smem[tid, kk] = T.cast(0, dtype)
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(A_smem[full_slices], R_local[full_slices])
|
||||
T.cuda.cta_sync()
|
||||
for kk in range(k):
|
||||
B[tid, kk] = A_smem[tid, kk]
|
||||
|
||||
return kernel
|
||||
|
||||
if non_r_scope == "global":
|
||||
if scope == "warpgroup":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([n_threads // 128])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
T.thread_id_in_wg([128])
|
||||
tid = T.thread_id([n_threads])
|
||||
for kk in range(k):
|
||||
A[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
|
||||
T.cuda.cta_sync()
|
||||
R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
|
||||
Tx.wg.copy(R_local[full_slices], A[full_slices])
|
||||
for kk in range(k):
|
||||
A[tid, kk] = T.cast(0, dtype)
|
||||
T.cuda.cta_sync()
|
||||
Tx.wg.copy(A[full_slices], R_local[full_slices])
|
||||
T.cuda.cta_sync()
|
||||
for kk in range(k):
|
||||
B[tid, kk] = A[tid, kk]
|
||||
|
||||
elif scope == "warp":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.lane_id([32])
|
||||
tid = T.thread_id([n_threads])
|
||||
for kk in range(k):
|
||||
A[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
|
||||
T.cuda.cta_sync()
|
||||
R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
|
||||
Tx.warp.copy(R_local[full_slices], A[full_slices])
|
||||
for kk in range(k):
|
||||
A[tid, kk] = T.cast(0, dtype)
|
||||
T.cuda.cta_sync()
|
||||
Tx.warp.copy(A[full_slices], R_local[full_slices])
|
||||
T.cuda.cta_sync()
|
||||
for kk in range(k):
|
||||
B[tid, kk] = A[tid, kk]
|
||||
|
||||
elif scope == "cta":
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warp_id([n_threads // 32])
|
||||
T.lane_id([32])
|
||||
tid = T.thread_id([n_threads])
|
||||
for kk in range(k):
|
||||
A[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
|
||||
T.cuda.cta_sync()
|
||||
R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
|
||||
Tx.cta.copy(R_local[full_slices], A[full_slices])
|
||||
for kk in range(k):
|
||||
A[tid, kk] = T.cast(0, dtype)
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(A[full_slices], R_local[full_slices])
|
||||
T.cuda.cta_sync()
|
||||
for kk in range(k):
|
||||
B[tid, kk] = A[tid, kk]
|
||||
|
||||
return kernel
|
||||
|
||||
raise ValueError(f"unsupported non_r_scope {non_r_scope!r}")
|
||||
|
||||
|
||||
def _expected(shape, dtype):
|
||||
n, k = shape
|
||||
np_dtype = tvm.testing.np_dtype_from_str(dtype)
|
||||
out = np.empty(shape, dtype=np_dtype)
|
||||
for t in range(n):
|
||||
for kk in range(k):
|
||||
out[t, kk] = (t * 100 + kk + 1) % 256 if dtype == "uint8" else t * 100 + kk + 1
|
||||
return out
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
@pytest.mark.parametrize("non_r_scope", ["shared", "global"])
|
||||
@pytest.mark.parametrize(
|
||||
"scope,n_threads,k",
|
||||
[
|
||||
("warpgroup", 128, 16),
|
||||
("warpgroup", 128, 32),
|
||||
("warpgroup", 128, 8),
|
||||
("warp", 32, 8),
|
||||
("warp", 32, 16),
|
||||
("cta", 256, 8),
|
||||
("cta", 256, 16),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["float16", "float32", "uint8"])
|
||||
def test_reg_roundtrip(scope, n_threads, k, dtype, non_r_scope):
|
||||
shape = (n_threads, k)
|
||||
kernel = _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope)
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
np_dtype = tvm.testing.np_dtype_from_str(dtype)
|
||||
B_np = np.zeros(shape, dtype=np_dtype)
|
||||
expected = _expected(shape, dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
if non_r_scope == "shared":
|
||||
compiled(B)
|
||||
else:
|
||||
A_np = np.zeros(shape, dtype=np_dtype)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
compiled(A, B)
|
||||
np.testing.assert_array_equal(B.numpy(), expected)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Migrated from test_copy_sync.py: sync G↔L copy via Tx.copy() (L = local =
|
||||
# per-thread register, so it dispatches to the reg variant).
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"task",
|
||||
[
|
||||
# A[3:4, 8:16, 8:16] -> A_local[0:8, 0:8] -> B[3:4, 8:16, 8:16]
|
||||
(
|
||||
(4, 16, 16), # g_shape
|
||||
(8, 8), # l_shape
|
||||
((3, 4), (8, 16), (8, 16)), # g_region
|
||||
1, # thread_cnt
|
||||
TileLayout(S[4, 16, 16]), # layoutA
|
||||
TileLayout(S[4, 16, 16]), # layoutB
|
||||
TileLayout(S[8, 8]), # layoutLocal
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
@pytest.mark.parametrize(
|
||||
"dtype", ["int8", "float8_e4m3fn", "float8_e5m2", "float16", "bfloat16", "float32"]
|
||||
)
|
||||
def test_copy_g2l_l2g_vec_load(task, dtype):
|
||||
g_shape, l_shape, g_region, thread_cnt, layoutA, layoutB, layoutLocal = task
|
||||
|
||||
r_lmem = tuple(slice(None) for _ in range(len(l_shape)))
|
||||
r_gmem = tuple(slice(g_region[i][0], g_region[i][1]) for i in range(len(g_shape)))
|
||||
|
||||
@T.prim_func
|
||||
def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, g_shape, dtype, layout=layoutA)
|
||||
B = T.match_buffer(B_ptr, g_shape, dtype, layout=layoutB)
|
||||
|
||||
T.device_entry()
|
||||
T.cta_id([2])
|
||||
T.thread_id([thread_cnt])
|
||||
A_local = T.alloc_buffer(l_shape, dtype, scope="local", layout=layoutLocal)
|
||||
Tx.copy(A_local[r_lmem], A[r_gmem])
|
||||
Tx.copy(B[r_gmem], A_local[r_lmem])
|
||||
|
||||
np_dtype = tvm.testing.np_dtype_from_str(dtype)
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy_sync})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
np.random.seed(0)
|
||||
A_np = tvm.testing.generate_random_array(dtype, g_shape)
|
||||
B_np = np.zeros(g_shape, dtype=np_dtype)
|
||||
|
||||
B_ref = B_np.copy()
|
||||
B_ref[r_gmem] = A_np[r_gmem]
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
np.testing.assert_allclose(B_ref, B.numpy())
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
def test_reg_copy_wg_local_to_swizzled_shared_uses_swizzle_fastpath():
|
||||
"""Regression: R→S copy where R has a ``wg_local_layout`` (thread iter
|
||||
``1 @ tid_in_wg``) must pick the widest vec PTX ``st.shared.v4`` AND use the
|
||||
swizzle fast path (precomputed ``signed_strides`` + per-iter
|
||||
bit-select), not the per-iter ``swizzle.apply()`` fallback.
|
||||
|
||||
Two distinct bugs this test guards against:
|
||||
|
||||
(1) ``_choose_vec_len`` used to include R-side thread-iter strides in
|
||||
its alignment check. ``wg_local_layout``'s thread iter has stride 1;
|
||||
a vec=8 (16-byte) alignment check on ``1 % 8 != 0`` would reject
|
||||
every wider variant and fall to scalar ``copy_16b``. Thread-axis
|
||||
strides are partition-coord (virtual), not storage-physical, so they
|
||||
must be excluded.
|
||||
|
||||
(2) Even at the widest vec, if the outer loop is a runtime serial
|
||||
(Python ``range`` doesn't actually unroll in TVMScript) the swizzle
|
||||
fast path's per-iter constant-fold can't kick in and the
|
||||
``tvm_builtin_pointer_offset`` swizzle XOR ends up recomputed every
|
||||
iteration. Loop must be ``T.unroll``.
|
||||
"""
|
||||
from tvm.tirx.layout import SwizzleLayout, wg_local_layout
|
||||
|
||||
N_THREADS, EPI_N = 128, 64
|
||||
g_shape = (N_THREADS, EPI_N)
|
||||
g_layout = TileLayout(S[g_shape])
|
||||
# 128b swizzle on the SMEM side (per_element=3 ⇒ 8 fp16 atom width).
|
||||
smem_layout = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3)
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, g_shape, "float16", layout=g_layout)
|
||||
B = T.match_buffer(B_ptr, g_shape, "float16", layout=g_layout)
|
||||
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.thread_id([N_THREADS])
|
||||
tid = T.thread_id_in_wg([N_THREADS])
|
||||
reg = T.alloc_buffer(g_shape, "float16", scope="local", layout=wg_local_layout(EPI_N))
|
||||
smem = T.alloc_buffer(g_shape, "float16", scope="shared", layout=smem_layout)
|
||||
|
||||
# Populate the per-thread slice via .local() (decomposes the wg
|
||||
# thread-axis layout into a per-thread 1D view).
|
||||
reg_local = reg.local(EPI_N)
|
||||
for i in T.serial(EPI_N):
|
||||
reg_local[i] = A[tid, i]
|
||||
Tx.wg.copy(smem, reg)
|
||||
T.cuda.cta_sync()
|
||||
for i in T.serial(EPI_N):
|
||||
B[tid, i] = smem[tid, i]
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
ex = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
src = ex.mod.imports[0].inspect_source()
|
||||
|
||||
# (1) Widest variant: 8 fp16 elements per call (16 bytes → v4.u32 st).
|
||||
assert "tvm_builtin_ptx_st" in src, (
|
||||
"expected PTX st in generated CUDA, alignment check fell back to a narrower variant"
|
||||
)
|
||||
assert "st.shared.v4" in src, "expected 128b vector store (st.shared.v4.u32)"
|
||||
assert "tvm_builtin_copy_" not in src, (
|
||||
"copy_xxb helpers appeared — reg dispatch should use PTX ld/st only"
|
||||
)
|
||||
# (2) Swizzle fast path fingerprint:
|
||||
# * emit_init allocates a size-N int buffer of "signed strides".
|
||||
# * emit_iter_offset uses bit-select * signed-stride: ``(bit) * v[i]``
|
||||
# where ``bit = (f >> M) & 1``.
|
||||
# The fallback (per-iter ``swizzle.apply(s_off + ds_per_iter)``) has no
|
||||
# such bit-select * signed-stride pattern.
|
||||
import re
|
||||
|
||||
bitsel_pattern = re.findall(r"& 1\) \* v_\d+\[", src)
|
||||
assert bitsel_pattern, (
|
||||
"fast-path bit-select pattern '& 1) * v_<n>[' not found; "
|
||||
"looks like emit_iter_offset's fast path didn't fire."
|
||||
)
|
||||
|
||||
|
||||
# --- tcgen05 D epilogue deposit (tf32_hc_prenorm_gemm) -----------------------
|
||||
# Production op: ``Tx.warpgroup.copy(smem_cd_mma, d_reg)`` after ``tcgen05.ld``
|
||||
# pulls the M=64 accumulator fragment from TMEM into ``d_reg``, then deposits it
|
||||
# into 128B-swizzled MMA SMEM for the subsequent TMA store to gmem D.
|
||||
_TCGEN05_D_ATOM = "16x256b"
|
||||
_TCGEN05_D_SHAPE = (64, 64)
|
||||
_TCGEN05_D_DTYPE = "float32"
|
||||
_TCGEN05_D_SWIZZLE = 3 # SwizzleMode 128B → mma_shared_layout(..., 3, shape)
|
||||
_TCGEN05_D_SLICE = (slice(0, 64), slice(0, 64))
|
||||
|
||||
|
||||
def _tcgen05_d_epilogue_layouts():
|
||||
from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout
|
||||
from tvm.tirx.layout import tcgen05_atom_layout
|
||||
|
||||
m, n = _TCGEN05_D_SHAPE
|
||||
reg_layout = tcgen05_atom_layout(_TCGEN05_D_ATOM, (m, n), _TCGEN05_D_DTYPE)
|
||||
smem_layout = mma_shared_layout(_TCGEN05_D_DTYPE, _TCGEN05_D_SWIZZLE, (m, n))
|
||||
return m, n, reg_layout, smem_layout
|
||||
|
||||
|
||||
def _build_tcgen05_d_epilogue_deposit():
|
||||
"""``Tx.wg.copy(smem[slice], d_reg[slice])``: R (tcgen05 atom) → S (128B swizzle)."""
|
||||
from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout
|
||||
from tvm.tirx.layout import tcgen05_atom_layout
|
||||
|
||||
m, n = _TCGEN05_D_SHAPE
|
||||
smem_layout = mma_shared_layout(_TCGEN05_D_DTYPE, _TCGEN05_D_SWIZZLE, (m, n))
|
||||
reg_layout = tcgen05_atom_layout(_TCGEN05_D_ATOM, (m, n), _TCGEN05_D_DTYPE)
|
||||
sl_m, sl_n = _TCGEN05_D_SLICE
|
||||
|
||||
@T.prim_func
|
||||
def deposit(
|
||||
d_reg: T.Buffer((m, n), _TCGEN05_D_DTYPE, scope="local", layout=reg_layout),
|
||||
) -> None:
|
||||
smem_cd_mma = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="shared", layout=smem_layout)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([1])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
T.thread_id_in_wg([128])
|
||||
Tx.wg.copy(smem_cd_mma[sl_m, sl_n], d_reg[sl_m, sl_n])
|
||||
|
||||
return deposit
|
||||
|
||||
|
||||
def _compile_tcgen05_d_epilogue_deposit():
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": _build_tcgen05_d_epilogue_deposit()})
|
||||
return tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
|
||||
def test_reg_copy_tcgen05_d_epilogue_deposit_layout_pairing():
|
||||
"""Pre-fix bug: canonical ``r_p`` collapses atom ``m`` groups and drops S pairings.
|
||||
|
||||
Copy: ``Tx.wg.copy(smem_cd_mma[0:64,0:64], d_reg[0:64,0:64])`` (R→S).
|
||||
``d_reg``: ``(64,64)`` fp32 ``tcgen05_atom_layout("16x256b", ...)``.
|
||||
``smem_cd_mma``: ``(64,64)`` fp32 ``mma_shared_layout(..., swizzle=128B)``.
|
||||
"""
|
||||
from tvm.backend.cuda.operator.tile_primitive.copy.reg import (
|
||||
_split_thread_loop,
|
||||
align_layouts_raw,
|
||||
)
|
||||
from tvm.tirx.exec_scope import ExecScope
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext
|
||||
|
||||
m, n, reg_layout, smem_layout = _tcgen05_d_epilogue_layouts()
|
||||
region = [(0, m), (0, n)]
|
||||
sctx = DispatchContext(
|
||||
tvm.target.Target("cuda"), ExecScope("warpgroup"), {}, {}, scope_kind="warpgroup"
|
||||
)
|
||||
with sctx.target:
|
||||
r_sliced = reg_layout.slice([m, n], region)
|
||||
s_sliced = smem_layout.slice([m, n], region)
|
||||
r_p, s_p, s_seps, r_perm = align_layouts_raw(r_sliced, s_sliced, region)
|
||||
|
||||
r_iters, s_groups = _split_thread_loop(r_perm, s_p, s_seps)
|
||||
r_iters_bug, s_groups_bug = _split_thread_loop(r_p, s_p, s_seps)
|
||||
mem_extents = [int(it.extent) for it in r_iters]
|
||||
bug_extents = [int(it.extent) for it in r_iters_bug]
|
||||
|
||||
# Fixed path: 3 register m-groups stay 1:1 with 3 S-side groups.
|
||||
assert mem_extents == [8, 2, 2]
|
||||
assert len(r_iters) == len(s_groups) == 3
|
||||
# Pre-fix path (what _split_thread_loop used to take): 1 fused m-iter, only
|
||||
# the first S group is paired — the other two are silently dropped.
|
||||
assert bug_extents == [32]
|
||||
assert len(r_iters_bug) == len(s_groups_bug) == 1
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
def test_reg_copy_tcgen05_d_epilogue_deposit_codegen():
|
||||
"""``reg`` dispatch lowers the D epilogue deposit; pre-fix loop only covered half.
|
||||
|
||||
Without the ``r_perm`` fix the emitted outer loop ran ``f < 8`` (one fused
|
||||
register group) instead of ``f < 16`` (three atom m-groups x vec tail), and
|
||||
the swizzled SMEM stores landed in the wrong (row, col) slots.
|
||||
"""
|
||||
import re
|
||||
|
||||
ex = _compile_tcgen05_d_epilogue_deposit()
|
||||
src = ex.mod.imports[0].inspect_source()
|
||||
|
||||
assert "copy/fallback" not in src, "reg dispatch must not fall back to scalar copy"
|
||||
assert "tvm_builtin_copy_" not in src, "reg dispatch should emit PTX ld/st only"
|
||||
assert "tvm_builtin_ptx_st" in src
|
||||
assert "st.shared.v2.u32" in src, "fp32 vec=2 → 8B shared store per outer iter"
|
||||
|
||||
loop = re.search(r"for \(int f = 0; f < (\d+)", src)
|
||||
assert loop is not None, "expected reg copy outer loop in generated CUDA"
|
||||
assert loop.group(1) == "16", (
|
||||
f"fixed pairing emits 16 outer stores (pre-fix bug collapsed to 8); got f < {loop.group(1)}"
|
||||
)
|
||||
|
||||
|
||||
def _tcgen05_16x256b_row_col(tid_wg: T.int32, lane: T.int32, reg_idx: T.int32):
|
||||
"""Map ``(tid_in_wg, reg)`` → logical ``(row, col)`` for ``.16x256b`` fp32 atom."""
|
||||
t0 = lane & T.int32(3)
|
||||
t1 = lane >> 2
|
||||
v0p = reg_idx & T.int32(1)
|
||||
va = (reg_idx >> 1) & T.int32(1)
|
||||
vb = reg_idx >> 2
|
||||
wid = tid_wg >> 5
|
||||
row = t1 + T.int32(8) * va + T.int32(16) * wid
|
||||
col = v0p + T.int32(2) * t0 + T.int32(8) * vb
|
||||
return row, col
|
||||
|
||||
|
||||
def _build_tcgen05_d_epilogue_deposit_roundtrip():
|
||||
"""Fill ``d_reg``, R→S deposit, S→R reload, dump via ``.local()`` to gmem."""
|
||||
from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout
|
||||
from tvm.tirx.layout import tcgen05_atom_layout
|
||||
|
||||
m, n = _TCGEN05_D_SHAPE
|
||||
regs_per_thread = 32 # ``.16x256b.x8`` fp32: 4 regs/slot x rep=8
|
||||
smem_layout = mma_shared_layout(_TCGEN05_D_DTYPE, _TCGEN05_D_SWIZZLE, (m, n))
|
||||
reg_layout = tcgen05_atom_layout(_TCGEN05_D_ATOM, (m, n), _TCGEN05_D_DTYPE)
|
||||
sl_m, sl_n = _TCGEN05_D_SLICE
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (m, n), _TCGEN05_D_DTYPE)
|
||||
B = T.match_buffer(B_ptr, (m, n), _TCGEN05_D_DTYPE)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([1])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
tid_wg = T.thread_id_in_wg([128])
|
||||
lane = T.lane_id([32])
|
||||
d_reg = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="local", layout=reg_layout)
|
||||
d_reg_out = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="local", layout=reg_layout)
|
||||
smem_cd_mma = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="shared", layout=smem_layout)
|
||||
reg_in = d_reg.local(regs_per_thread)
|
||||
reg_out = d_reg_out.local(regs_per_thread)
|
||||
for r in T.serial(regs_per_thread):
|
||||
row, col = _tcgen05_16x256b_row_col(tid_wg, lane, T.cast(r, "int32"))
|
||||
reg_in[r] = A[row, col]
|
||||
Tx.wg.copy(smem_cd_mma[sl_m, sl_n], d_reg[sl_m, sl_n])
|
||||
T.cuda.cta_sync()
|
||||
Tx.wg.copy(d_reg_out[sl_m, sl_n], smem_cd_mma[sl_m, sl_n])
|
||||
for r in T.serial(regs_per_thread):
|
||||
row, col = _tcgen05_16x256b_row_col(tid_wg, lane, T.cast(r, "int32"))
|
||||
B[row, col] = reg_out[r]
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
def test_reg_copy_tcgen05_d_epilogue_deposit_gpu():
|
||||
"""GPU: R→S deposit + S→R reload must preserve the Layout-F register tile.
|
||||
|
||||
Host fills gmem ``A[row,col]=row*100+col``; each thread scatters into ``d_reg``
|
||||
via the ``.16x256b`` (tid,reg)→(row,col) map, runs production
|
||||
``Tx.wg.copy(smem_cd_mma, d_reg)`` then the inverse
|
||||
``Tx.wg.copy(d_reg_out, smem_cd_mma)``, and dumps ``d_reg_out`` back to gmem
|
||||
``B`` through ``.local()``. Pre-fix pairing dropped 2/3 of the S groups —
|
||||
``max|B-A|`` was hundreds, not 0.
|
||||
"""
|
||||
m, n = _TCGEN05_D_SHAPE
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.compile(
|
||||
tvm.IRModule({"main": _build_tcgen05_d_epilogue_deposit_roundtrip()}),
|
||||
target=target,
|
||||
tir_pipeline="tirx",
|
||||
)
|
||||
|
||||
rows = np.arange(m, dtype=np.int32)[:, None]
|
||||
cols = np.arange(n, dtype=np.int32)[None, :]
|
||||
a_np = (rows * 100 + cols).astype(np.float32)
|
||||
b_np = np.zeros((m, n), dtype=np.float32)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
a = tvm.runtime.tensor(a_np, dev)
|
||||
b = tvm.runtime.tensor(b_np, dev)
|
||||
mod(a, b)
|
||||
np.testing.assert_allclose(b.numpy(), a_np, rtol=0, atol=0)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,443 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Tests for the generic swizzle-aware iter pattern in
|
||||
``cuda/copy/_swizzle_iter.py``.
|
||||
|
||||
Two layers:
|
||||
|
||||
* **Recognizer tests** check that ``try_recognize`` returns the expected
|
||||
``SwizzlePattern`` (or rejects) for each of conditions (a)+(b)+(c).
|
||||
* **Numeric correctness tests** verify the proof empirically: for many
|
||||
``(M0, k)`` samples, the formula
|
||||
``apply(M0) + sum_{j : bit_j(k)=1} signed_strides[j]``
|
||||
equals ``apply(M0 + ds_k)`` computed by the layout's own Apply formula.
|
||||
Plus a per-thread-sign-matters test that would fail for a constant-sign
|
||||
implementation, ensuring the test isn't trivially satisfied.
|
||||
|
||||
All algorithm-level (no GPU needed). End-to-end emit is tested in
|
||||
``test_gmem_smem.py::test_swizzled_smem_emit_must_be_swizzle_aware``.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
from tvm.tirx import Var as _TirVar
|
||||
from tvm.tirx.cuda.operator.tile_primitive.copy._swizzle_iter import (
|
||||
get_swizzle,
|
||||
try_recognize,
|
||||
)
|
||||
from tvm.tirx.expr import IntImm as _IntImm
|
||||
from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Pure-Python reference: SwizzleLayout's Apply, plus the proof's formula.
|
||||
# Used as ground truth — both must agree for the proof to hold.
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def py_swizzle_apply(M: int, p: int, sw: int, at: int) -> int:
|
||||
"""Pure-Python reimplementation of SwizzleLayoutNode::Apply (swizzle_inner=True):
|
||||
phys = swz_q * C + (M mod C)
|
||||
q = M / C; swz_q = q XOR ((q & outer_mask) >> at)
|
||||
"""
|
||||
C = 1 << p
|
||||
q = M // C
|
||||
outer_mask = ((1 << sw) - 1) << at
|
||||
swz_q = q ^ ((q & outer_mask) >> at)
|
||||
return swz_q * C + (M % C)
|
||||
|
||||
|
||||
def py_signed_strides(
|
||||
M0: int, p: int, sw: int, at: int, bit_positions: list[int], iter_strides_elems: list[int]
|
||||
) -> list[int]:
|
||||
"""Pure-Python reimplementation of emit_init's formula. Mirrors:
|
||||
if bj >= sw: sigma_bj = +1 (mid_bits)
|
||||
else : sigma_bj = 1 - 2 * bit_(at+bj)(M0/C) (chunk_bits)
|
||||
signed_strides[j] = sigma_bj * iter_strides_elems[j]
|
||||
"""
|
||||
C = 1 << p
|
||||
q = M0 // C
|
||||
out: list[int] = []
|
||||
for bj, stride in zip(bit_positions, iter_strides_elems):
|
||||
if bj >= sw:
|
||||
out.append(stride)
|
||||
else:
|
||||
row_bit = (q >> (at + bj)) & 1
|
||||
sigma = 1 - 2 * row_bit
|
||||
out.append(sigma * stride)
|
||||
return out
|
||||
|
||||
|
||||
def py_iter_offset(base_off: int, k: int, signed_strides: list[int]) -> int:
|
||||
"""Formula sum: base_off + sum_{j : bit_(n-1-j)(k)=1} signed_strides[j]."""
|
||||
n = len(signed_strides)
|
||||
off = base_off
|
||||
for j in range(n):
|
||||
if (k >> (n - 1 - j)) & 1:
|
||||
off += signed_strides[j]
|
||||
return off
|
||||
|
||||
|
||||
def py_outer_ds(k: int, iter_extents: list[int], iter_strides: list[int]) -> int:
|
||||
"""Decode a flat outer index k into per-iter coords (matching
|
||||
_flat_outer_coords) and sum coord_i * stride_i for the corresponding ds."""
|
||||
coords: list[int] = []
|
||||
rem = k
|
||||
for ext in reversed(iter_extents):
|
||||
coords.append(rem % ext)
|
||||
rem //= ext
|
||||
coords.reverse()
|
||||
return sum(c * s for c, s in zip(coords, iter_strides))
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Recognizer tests — verify try_recognize accepts / rejects under (a)+(b)+(c).
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_swizzle_extracts_from_compose():
|
||||
sw = SwizzleLayout(3, 3, 3)
|
||||
assert get_swizzle(sw) is not None
|
||||
assert get_swizzle(ComposeLayout(sw, TileLayout(S[(64, 64)]))) is not None
|
||||
assert get_swizzle(TileLayout(S[(64, 64)])) is None
|
||||
|
||||
|
||||
def test_recognize_nvfp4_case():
|
||||
"""nvfp4's epilogue: SwizzleLayout(3,3,3), iter extents [2,2,2] strides
|
||||
[8,16,32], M0 = tid * 64 (each thread starts at col 0 of one row;
|
||||
row_stride 64 = 8 chunks, ensures chunk bits of M0/C are zero for all
|
||||
iter bit positions)."""
|
||||
sw = SwizzleLayout(3, 3, 3)
|
||||
tid = _TirVar("tid", "int32")
|
||||
# M0 = tid * 64 → M0/C = tid * 8 → bits 0,1,2 are 0 (since multiplied by 8).
|
||||
M0 = tid * _IntImm("int32", 64)
|
||||
pat = try_recognize(sw, [2, 2, 2], [8, 16, 32], M0)
|
||||
assert pat is not None
|
||||
assert pat.bit_positions == [0, 1, 2]
|
||||
assert pat.iter_strides_elems == [8, 16, 32]
|
||||
assert pat.n_binary_iters == 3
|
||||
|
||||
|
||||
def test_recognize_binary_split():
|
||||
"""A single outer iter with extent=4 stride=8 splits into two binary
|
||||
iters with strides 16 and 8 (outermost first, matching _flat_outer_coords)."""
|
||||
sw = SwizzleLayout(3, 3, 3)
|
||||
tid = _TirVar("tid", "int32")
|
||||
M0 = tid * _IntImm("int32", 64)
|
||||
pat = try_recognize(sw, [4], [8], M0)
|
||||
assert pat is not None
|
||||
# Split: stride 8*2 = 16 (outermost), stride 8 (innermost) → bits [1, 0]
|
||||
assert pat.bit_positions == [1, 0]
|
||||
assert pat.iter_strides_elems == [16, 8]
|
||||
|
||||
|
||||
def test_recognize_mid_bits():
|
||||
"""SwizzleLayout(p=4, sw=2, at=4): chunk bits [0,2), mid bits [2,4),
|
||||
row bits [4,6). An iter at bj=2 lives in mid_bits → sigma is always +1
|
||||
(i.e., the recognizer accepts and the sign formula won't read row bits)."""
|
||||
sw = SwizzleLayout(4, 2, 4) # C=16, mid_bits cover bits 2..3
|
||||
tid = _TirVar("tid", "int32")
|
||||
# M0/C must have bit 2 == 0. Pick row_stride = 64 (= 4*C) so M0/C = tid*4
|
||||
# which has zeros at bit 0,1, and bit 2 is bit 0 of tid... hmm that varies.
|
||||
# Use row_stride = 128 (= 8*C, contributes 4 to M0/C per tid → bit 2 of M0/C
|
||||
# depends on whether tid is even/odd — not zero. Instead use row_stride such
|
||||
# that M0/C is provably 0 mod 8 = 0 at bits 0..2: row_stride = 256 (= 16*C)
|
||||
# → M0/C = tid*16 → bits 0..3 all 0. iter_mask = bit 2, divisor = C*4 = 64.
|
||||
M0 = tid * _IntImm("int32", 256)
|
||||
pat = try_recognize(sw, [2], [64], M0) # stride 64 = C * 2^2 → bj=2 (mid)
|
||||
assert pat is not None
|
||||
assert pat.bit_positions == [2]
|
||||
assert pat.iter_strides_elems == [64]
|
||||
|
||||
|
||||
def test_reject_not_chunk_aligned():
|
||||
"""Condition (a): stride must be a multiple of C."""
|
||||
sw = SwizzleLayout(3, 3, 3) # C=8
|
||||
tid = _TirVar("tid", "int32")
|
||||
M0 = tid * _IntImm("int32", 64)
|
||||
# stride 4 is not a multiple of C=8 → reject.
|
||||
assert try_recognize(sw, [2], [4], M0) is None
|
||||
|
||||
|
||||
def test_reject_carries_into_row_bits():
|
||||
"""Condition (b): bj < at. A binary iter with stride C * 2^at lands at
|
||||
bj=at, which would change the row bits → reject."""
|
||||
sw = SwizzleLayout(3, 3, 3) # at=3, so max bj = 2
|
||||
tid = _TirVar("tid", "int32")
|
||||
M0 = tid * _IntImm("int32", 64)
|
||||
# Strides 8,16,32 OK (bj=0,1,2); 64 → bj=3 → reject.
|
||||
assert try_recognize(sw, [2, 2, 2, 2], [8, 16, 32, 64], M0) is None
|
||||
|
||||
|
||||
def test_reject_chunk_overlap():
|
||||
"""Condition (c): (M0/C) must have 0 bits at all iter-bit positions per
|
||||
thread. If M0 = tid * 8 (so M0/C = tid), then bit 0 of M0/C is bit 0 of
|
||||
tid — analyzer can't prove this is 0 across all threads, so reject."""
|
||||
sw = SwizzleLayout(3, 3, 3) # C=8
|
||||
tid = _TirVar("tid", "int32")
|
||||
# M0 = tid * C = tid * 8 → M0/C = tid → bit 0 NOT provably zero.
|
||||
M0 = tid * _IntImm("int32", 8)
|
||||
assert try_recognize(sw, [2], [8], M0) is None
|
||||
|
||||
|
||||
def test_recognize_no_outer_iters():
|
||||
"""Degenerate case: no outer iter at all. Recognizer returns a trivial
|
||||
pattern (empty bit_positions). Emit will use base_off alone."""
|
||||
sw = SwizzleLayout(3, 3, 3)
|
||||
tid = _TirVar("tid", "int32")
|
||||
M0 = tid * _IntImm("int32", 64)
|
||||
pat = try_recognize(sw, [], [], M0)
|
||||
assert pat is not None
|
||||
assert pat.n_binary_iters == 0
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Numeric correctness — the PROOF. The formula must equal apply(M0 + ds_k)
|
||||
# for all sampled (M0, k) and for non-trivial M0 values per thread.
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"p,sw,at,iter_extents,iter_strides,row_stride",
|
||||
[
|
||||
# nvfp4-like (p=sw=at=3, 3 binary iters covering one swizzle row)
|
||||
(3, 3, 3, [2, 2, 2], [8, 16, 32], 64),
|
||||
# single binary iter at chunk_bit position
|
||||
(3, 3, 3, [2], [8], 64),
|
||||
# split-from-extent-4 (one outer becomes two binary)
|
||||
(3, 3, 3, [4], [8], 64),
|
||||
# mid_bits region
|
||||
(4, 2, 4, [2], [64], 256),
|
||||
# mix: one chunk_bit + one mid_bit
|
||||
(
|
||||
3,
|
||||
2,
|
||||
4,
|
||||
[2, 2],
|
||||
[8, 32],
|
||||
256,
|
||||
), # C=8, sw=2, at=4 → bj_max=3 for stride 64 → use 32 (bj=2 in mid)
|
||||
],
|
||||
)
|
||||
def test_formula_matches_apply_under_conditions(
|
||||
p,
|
||||
sw,
|
||||
at,
|
||||
iter_extents,
|
||||
iter_strides,
|
||||
row_stride,
|
||||
):
|
||||
"""For every (M0, k) sample, the signed-strides formula must equal
|
||||
py_swizzle_apply(M0 + ds_k). Sweeps multiple per-thread M0 values to
|
||||
catch any per-thread-sign bug (a constant-sign impl would fail here)."""
|
||||
swizzle = SwizzleLayout(p, sw, at)
|
||||
tid = _TirVar("tid", "int32")
|
||||
M0_template = tid * _IntImm("int32", row_stride)
|
||||
pat = try_recognize(swizzle, iter_extents, iter_strides, M0_template)
|
||||
assert pat is not None, (
|
||||
f"recognizer rejected supposedly-valid case p={p},sw={sw},at={at} "
|
||||
f"iter_extents={iter_extents} iter_strides={iter_strides} row_stride={row_stride}"
|
||||
)
|
||||
|
||||
total_iters = 1
|
||||
for ext in iter_extents:
|
||||
total_iters *= ext
|
||||
|
||||
# Per-thread sweep: pick concrete tid values that span a few rows of the
|
||||
# swizzle atom. Tid = 0 alone would hide the sign issue (M0/C row bits all
|
||||
# 0 → all signs +1); larger tids exercise the sign-flip branches.
|
||||
for tid_val in [0, 1, 3, 5, 7, 13, 21]:
|
||||
M0 = tid_val * row_stride
|
||||
base_off = py_swizzle_apply(M0, p, sw, at)
|
||||
ss = py_signed_strides(
|
||||
M0,
|
||||
p,
|
||||
sw,
|
||||
at,
|
||||
pat.bit_positions,
|
||||
pat.iter_strides_elems,
|
||||
)
|
||||
for k in range(total_iters):
|
||||
ds_k = py_outer_ds(k, iter_extents, iter_strides)
|
||||
ground_truth = py_swizzle_apply(M0 + ds_k, p, sw, at)
|
||||
formula = py_iter_offset(base_off, k, ss)
|
||||
assert formula == ground_truth, (
|
||||
f"formula mismatch: p={p},sw={sw},at={at} "
|
||||
f"iter_extents={iter_extents} iter_strides={iter_strides} "
|
||||
f"tid={tid_val} M0={M0} k={k} ds_k={ds_k} "
|
||||
f"apply(M0+ds_k)={ground_truth} formula={formula} "
|
||||
f"signed_strides={ss}"
|
||||
)
|
||||
|
||||
|
||||
def test_per_thread_sign_actually_varies():
|
||||
"""Guard against a 'constant +1 stride' bug: for the nvfp4-like case,
|
||||
different tids MUST produce different signed_strides[0] (since the
|
||||
formula is sigma_0 = 1 - 2 * bit_(at)(M0/C) and that bit toggles with
|
||||
tid). If a buggy impl always returned +stride, this test would catch it."""
|
||||
p, sw, at = 3, 3, 3
|
||||
row_stride = 64 # M0/C = tid * 8 → bit (at=3) = bit 0 of tid
|
||||
# tid=0 → M0/C bit 3 = 0 → sigma = +1; tid=1 → bit 3 = 1 → sigma = -1
|
||||
ss_even = py_signed_strides(0 * row_stride, p, sw, at, [0], [8])
|
||||
ss_odd = py_signed_strides(1 * row_stride, p, sw, at, [0], [8])
|
||||
assert ss_even != ss_odd, "per-thread sign formula degenerated to constant — proof / impl bug"
|
||||
assert ss_even == [8]
|
||||
assert ss_odd == [-8]
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Fallback path — when recognizer rejects, per-iter swizzle.apply gives the
|
||||
# right answer. This is trivial (we delegate to layout.apply) but documents
|
||||
# the contract.
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_recognize_linear_iter_pure_case_1d():
|
||||
"""Outer iter with non-pow2 ext is accepted IF its stride is a multiple
|
||||
of the swizzle period 2^(p+at+sw) (pure Case 1.D, swizzle has no XOR
|
||||
effect). The iter is stored as a LinearIter (no bit decomposition).
|
||||
"""
|
||||
from tvm.tirx.cuda.operator.tile_primitive.copy._swizzle_iter import (
|
||||
_BitIter,
|
||||
_LinearIter,
|
||||
)
|
||||
|
||||
p, sw, at = 3, 3, 3
|
||||
swizzle = SwizzleLayout(p, sw, at)
|
||||
period = 1 << (p + at + sw) # 512
|
||||
# Outer iter (ext=3, stride=period) — non-pow2 but pure Case 1.D.
|
||||
# Inner iter (ext=2, stride=8) — pow2, Case 1.A (bj=0).
|
||||
pat = try_recognize(swizzle, [3, 2], [period, 8], _IntImm("int32", 0))
|
||||
assert pat is not None
|
||||
assert len(pat.outer_iters) == 2
|
||||
# Outermost (index 0) corresponds to first input iter = the linear one.
|
||||
assert isinstance(pat.outer_iters[0], _LinearIter)
|
||||
assert pat.outer_iters[0].ext == 3
|
||||
assert pat.outer_iters[0].stride == period
|
||||
# Innermost (index 1) is the binary-split iter.
|
||||
assert isinstance(pat.outer_iters[1], _BitIter)
|
||||
assert pat.outer_iters[1].ext == 2
|
||||
assert pat.outer_iters[1].n_bits == 1
|
||||
assert pat.outer_iters[1].slot_start == 0
|
||||
# bit_positions / iter_strides_elems only contain the binary iter's bit.
|
||||
assert pat.bit_positions == [0] # 8/8 = 2^0
|
||||
assert pat.iter_strides_elems == [8]
|
||||
|
||||
|
||||
def test_reject_non_pow2_ext_not_case_1d():
|
||||
"""Non-pow2 ext where stride is NOT in pure Case 1.D regime — reject.
|
||||
stride=64 = 2^(p+at) = one atom row, which is Case 1.C (in [at, at+sw))
|
||||
territory and the XOR depends on M0, so the linear path is unsafe."""
|
||||
swizzle = SwizzleLayout(3, 3, 3)
|
||||
pat = try_recognize(swizzle, [3], [64], _IntImm("int32", 0))
|
||||
assert pat is None
|
||||
|
||||
|
||||
def test_emit_mixed_linear_bit_correctness():
|
||||
"""Brute-force: for a mixed (LinearIter outer, BitIter inner) pattern,
|
||||
emit_iter_offset's prediction must equal the actual swizzle output for
|
||||
every (tid, k) — including the non-pow2 outer extent's coord 2."""
|
||||
from tvm.tirx.cuda.operator.tile_primitive.copy._swizzle_iter import (
|
||||
_LinearIter,
|
||||
)
|
||||
|
||||
p, sw, at = 3, 3, 3
|
||||
swizzle = SwizzleLayout(p, sw, at)
|
||||
period = 1 << (p + at + sw) # 512
|
||||
iter_extents, iter_strides = [3, 2], [period, 8]
|
||||
tid = _TirVar("tid", "int32")
|
||||
# Inner iter bj=0 in [0, sw); (C1) needs bit_0(M0/C) = 0 ⇒ M0/C even
|
||||
# ⇒ M0 multiple of 16. So row_stride = 16.
|
||||
M0_template = tid * _IntImm("int32", 16)
|
||||
pat = try_recognize(swizzle, iter_extents, iter_strides, M0_template)
|
||||
assert pat is not None
|
||||
|
||||
def py_emit(pattern, signed_strides, base_off, k):
|
||||
off = base_off
|
||||
remaining = k
|
||||
for it in reversed(pattern.outer_iters):
|
||||
c = remaining % it.ext
|
||||
remaining = remaining // it.ext
|
||||
if isinstance(it, _LinearIter):
|
||||
off += c * it.stride
|
||||
else:
|
||||
for b in range(it.n_bits):
|
||||
bit_pos = it.n_bits - 1 - b
|
||||
slot = it.slot_start + b
|
||||
if (c >> bit_pos) & 1:
|
||||
off += signed_strides[slot]
|
||||
return off
|
||||
|
||||
total_k = iter_extents[0] * iter_extents[1]
|
||||
for tid_val in [0, 1, 5, 7, 13]:
|
||||
M0 = tid_val * 16
|
||||
base_off = py_swizzle_apply(M0, p, sw, at)
|
||||
ss = py_signed_strides(
|
||||
M0,
|
||||
p,
|
||||
sw,
|
||||
at,
|
||||
pat.bit_positions,
|
||||
pat.iter_strides_elems,
|
||||
)
|
||||
for k in range(total_k):
|
||||
ds_k = py_outer_ds(k, iter_extents, iter_strides)
|
||||
ground_truth = py_swizzle_apply(M0 + ds_k, p, sw, at)
|
||||
formula = py_emit(pat, ss, base_off, k)
|
||||
assert formula == ground_truth, (
|
||||
f"mixed mismatch: tid={tid_val} M0={M0} k={k} ds_k={ds_k} "
|
||||
f"truth={ground_truth} formula={formula} ss={ss}"
|
||||
)
|
||||
|
||||
|
||||
def test_fallback_path_when_recognizer_rejects():
|
||||
"""The recognizer should reject when (c) fails, and the resulting
|
||||
fallback emit (swizzle.apply per iter) is the correct path. This test
|
||||
proves the rejection and demonstrates that the swizzled offset really
|
||||
differs from the linear offset for the rejected case — so a buggy
|
||||
`linear-offset-without-XOR` emit (the pre-fix behavior) would give the
|
||||
wrong answer on at least one (tid, k) sample. The fallback emit, by
|
||||
construction, delegates to swizzle.apply and is thus correct."""
|
||||
p, sw, at = 3, 3, 3
|
||||
swizzle = SwizzleLayout(p, sw, at)
|
||||
tid = _TirVar("tid", "int32")
|
||||
M0_template = tid * _IntImm("int32", 8) # (c) fails: bit 0 of M0/C = bit 0 of tid
|
||||
pat = try_recognize(swizzle, [2], [8], M0_template)
|
||||
assert pat is None, "recognizer must reject when (c) fails"
|
||||
|
||||
# Demonstrate the swizzled offset differs from linear for at least one
|
||||
# (tid, k) — proves the swizzle is actually non-trivial here and the
|
||||
# broken linear-offset emit would give the wrong physical address.
|
||||
iter_extents, iter_strides = [2], [8]
|
||||
diverging_samples = 0
|
||||
for tid_val in range(16):
|
||||
M0 = tid_val * 8
|
||||
for k in range(2):
|
||||
ds_k = py_outer_ds(k, iter_extents, iter_strides)
|
||||
linear = M0 + ds_k
|
||||
swizzled = py_swizzle_apply(linear, p, sw, at)
|
||||
if swizzled != linear:
|
||||
diverging_samples += 1
|
||||
assert diverging_samples > 0, (
|
||||
"no (tid, k) sample shows swizzled != linear — the swizzle is a "
|
||||
"no-op for this layout, so the test isn't catching anything"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,251 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
"""Tests for the DSMEM (shared::cta → shared::cluster) copy_async variant.
|
||||
|
||||
Split out from ``test_copy_async.py`` so the TMA-focused file stays focused
|
||||
on the g2s/s2g TMA family. Any cross-cutting copy_async helper that both
|
||||
files need should live in a shared module, not be duplicated.
|
||||
"""
|
||||
|
||||
import functools
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx import IntImm, Var
|
||||
from tvm.tirx.cuda.operator.tile_primitive.copy_async.dsmem import copy_dsmem_impl
|
||||
from tvm.tirx.exec_scope import ExecScope
|
||||
from tvm.tirx.layout import S, TileLayout
|
||||
from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import DispatchFail
|
||||
from tvm.tirx.operator.tile_primitive.ops import CopyAsync
|
||||
from tvm.tirx.stmt_functor import StmtExprVisitor
|
||||
|
||||
|
||||
def _make_dsmem_dispatch_call(shape, dtype, src_layout, dst_layout):
|
||||
"""Call copy_dsmem_impl directly. Returns impl or raises DispatchFail."""
|
||||
from tvm.ir import Range
|
||||
from tvm.tirx.stmt import BufferRegion
|
||||
|
||||
src_buf = tvm.tirx.decl_buffer(shape, dtype, "A", scope="shared.dyn", layout=src_layout)
|
||||
dst_buf = tvm.tirx.decl_buffer(shape, dtype, "B", scope="shared.dyn", layout=dst_layout)
|
||||
ranges = [Range.from_min_extent(0, s) for s in shape]
|
||||
config = {"mbar": Var("mbar", "handle"), "remote_cta_id": IntImm("int32", 1)}
|
||||
op_call = CopyAsync(BufferRegion(dst_buf, ranges), BufferRegion(src_buf, ranges), config=config)
|
||||
target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
|
||||
sctx = DispatchContext(target, ExecScope("thread"), {}, {})
|
||||
return copy_dsmem_impl(op_call, sctx)
|
||||
|
||||
|
||||
class _S2CCounter(StmtExprVisitor):
|
||||
"""Count cp.async.bulk.shared_to_cluster calls including loop iterations."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._loop_extents = []
|
||||
self.total = 0
|
||||
|
||||
def visit_for_(self, op):
|
||||
self._loop_extents.append(op.extent)
|
||||
self.visit_stmt(op.body)
|
||||
self._loop_extents.pop()
|
||||
|
||||
def visit_evaluate_(self, op):
|
||||
if isinstance(op.value, tvm.ir.Call):
|
||||
if op.value.op.name == "tirx.ptx.cp_async_bulk_shared_to_cluster":
|
||||
n = 1
|
||||
for e in self._loop_extents:
|
||||
n *= e
|
||||
self.total += n
|
||||
|
||||
|
||||
def _count_s2c_ops(impl):
|
||||
c = _S2CCounter()
|
||||
c.visit_stmt(impl.body)
|
||||
return c.total
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Parametrized DSMEM test: dispatch assertion + GPU correctness
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# (shape, dtype, src_spec, dst_spec, expected_s2c_ops | "fail")
|
||||
# Dispatch assertion uses src_spec/dst_spec as given.
|
||||
# GPU correctness (all non-fail cases) uses src_spec as the layout for both CTAs.
|
||||
DSMEM_CONFIGS = [
|
||||
pytest.param((128, 64), "float16", S[128, 64], S[128, 64], 1, id="contiguous-2d"),
|
||||
pytest.param((256,), "float16", S[256], S[256], 1, id="contiguous-1d"),
|
||||
# Stride gap: inner 128 contiguous, outer stride=256 (gap) → 8 bulk copies
|
||||
pytest.param(
|
||||
(8, 128), "float16", S[(8, 128) : (256, 1)], S[(8, 128) : (256, 1)], 8, id="stride-gap"
|
||||
),
|
||||
# Different outer strides → 8 bulk copies in dispatch
|
||||
pytest.param(
|
||||
(8, 128),
|
||||
"float16",
|
||||
S[(8, 128) : (256, 1)],
|
||||
S[(8, 128) : (512, 1)],
|
||||
8,
|
||||
id="partial-contiguity-diff-stride",
|
||||
),
|
||||
# Incompatible: row-major vs column-major → DispatchFail
|
||||
pytest.param(
|
||||
(4, 64), "float16", S[4, 64], S[(4, 64) : (1, 4)], "fail", id="incompatible-row-vs-col"
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _layout_physical_elements(layout):
|
||||
"""Compute number of physical elements needed for a TileLayout."""
|
||||
max_offset = 0
|
||||
for shard in layout.shard:
|
||||
if shard.axis.is_memory():
|
||||
max_offset += int(shard.stride) * (int(shard.extent) - 1)
|
||||
return max_offset + 1
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
@pytest.mark.parametrize("shape,dtype,src_spec,dst_spec,expected", DSMEM_CONFIGS)
|
||||
def test_dsmem(shape, dtype, src_spec, dst_spec, expected):
|
||||
"""Dispatch assertion + GPU correctness for DSMEM copy.
|
||||
|
||||
Always tests dispatch (s2c op count or DispatchFail).
|
||||
For non-fail cases: also runs a 2-CTA cluster kernel via T.copy_async
|
||||
dispatch (using src_spec as layout for both CTAs) and verifies correctness.
|
||||
"""
|
||||
from tvm.tirx.lang.pipeline import MBarrier
|
||||
|
||||
src_layout = TileLayout(src_spec)
|
||||
dst_layout = TileLayout(dst_spec)
|
||||
|
||||
# --- Dispatch assertion ---
|
||||
if expected == "fail":
|
||||
with pytest.raises(DispatchFail):
|
||||
_make_dsmem_dispatch_call(shape, dtype, src_layout, dst_layout)
|
||||
return
|
||||
|
||||
impl = _make_dsmem_dispatch_call(shape, dtype, src_layout, dst_layout)
|
||||
assert _count_s2c_ops(impl) == expected
|
||||
|
||||
# --- GPU correctness ---
|
||||
# Allocate two separate smem buffers: src_smem (src_layout) and dst_smem
|
||||
# (dst_layout). CTA 0 loads global→src_smem, copy_async copies src_smem→
|
||||
# dst_smem on CTA 1. CTA 1 reads dst_smem and writes to global output.
|
||||
|
||||
CLUSTER_N = 2
|
||||
n_elements = functools.reduce(lambda a, b: a * b, shape, 1)
|
||||
copy_bytes = n_elements * tvm.DataType(dtype).bits // 8
|
||||
src_phys = _layout_physical_elements(src_layout)
|
||||
dst_phys = _layout_physical_elements(dst_layout)
|
||||
r = tuple(slice(0, s) for s in shape)
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def dsmem_copy(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, dtype)
|
||||
B = T.match_buffer(B_ptr, shape, dtype)
|
||||
|
||||
T.device_entry()
|
||||
cbx = T.cta_id_in_cluster([CLUSTER_N])
|
||||
T.cta_id([CLUSTER_N])
|
||||
tid = T.thread_id([1])
|
||||
pool = T.SMEMPool()
|
||||
# src_smem: CTA 0 writes here, dispatch reads from here
|
||||
src_raw = pool.alloc([src_phys], dtype, align=128)
|
||||
src_smem = T.decl_buffer(
|
||||
list(shape), dtype, src_raw.data,
|
||||
elem_offset=0, scope="shared.dyn", layout=src_layout,
|
||||
)
|
||||
# dst_smem: dispatch writes here (on remote CTA), CTA 1 reads
|
||||
dst_raw = pool.alloc([dst_phys], dtype, align=128)
|
||||
dst_smem = T.decl_buffer(
|
||||
list(shape), dtype, dst_raw.data,
|
||||
elem_offset=0, scope="shared.dyn", layout=dst_layout,
|
||||
)
|
||||
mbar = MBarrier(pool, 1)
|
||||
pool.commit()
|
||||
|
||||
mbar.init(1)
|
||||
T.ptx.fence.mbarrier_init()
|
||||
T.cuda.cluster_sync()
|
||||
|
||||
if tid == 0:
|
||||
if cbx == 0:
|
||||
Tx.copy(src_smem[r], A[r])
|
||||
T.ptx.fence.proxy_async("shared::cta")
|
||||
|
||||
Tx.copy_async(
|
||||
dst_smem[r], src_smem[r],
|
||||
dispatch="dsmem",
|
||||
mbar=mbar.ptr_to([0]),
|
||||
remote_cta_id=T.int32(1),
|
||||
)
|
||||
else:
|
||||
T.ptx.mbarrier.arrive.expect_tx(mbar.ptr_to([0]), copy_bytes)
|
||||
mbar.wait(0, 0)
|
||||
|
||||
Tx.copy(B[r], dst_smem[r])
|
||||
# fmt: on
|
||||
|
||||
np_dtype = tvm.testing.np_dtype_from_str(dtype)
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": dsmem_copy})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
cuda_src = mod.mod.imports[0].inspect_source()
|
||||
assert "cp.async.bulk.shared::cluster.shared::cta" in cuda_src
|
||||
|
||||
np.random.seed(0)
|
||||
A_np = tvm.testing.generate_random_array(dtype, shape)
|
||||
B_np = np.zeros(shape, dtype=np_dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_tvm = tvm.runtime.tensor(A_np, dev)
|
||||
B_tvm = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A_tvm, B_tvm)
|
||||
np.testing.assert_allclose(A_np, B_tvm.numpy())
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
def test_dsmem_dispatch_missing_config():
|
||||
"""Dispatch fails when required config keys are missing."""
|
||||
from tvm.ir import Range
|
||||
from tvm.tirx.stmt import BufferRegion
|
||||
|
||||
layout = TileLayout(S[64])
|
||||
buf = tvm.tirx.decl_buffer((64,), "float16", "A", scope="shared.dyn", layout=layout)
|
||||
br = BufferRegion(buf, [Range.from_min_extent(0, 64)])
|
||||
target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
|
||||
sctx = DispatchContext(target, ExecScope("thread"), {}, {})
|
||||
|
||||
with pytest.raises(DispatchFail, match="remote_cta_id"):
|
||||
copy_dsmem_impl(CopyAsync(br, br, config={"mbar": Var("m", "handle")}), sctx)
|
||||
with pytest.raises(DispatchFail, match="mbar"):
|
||||
copy_dsmem_impl(CopyAsync(br, br, config={"remote_cta_id": IntImm("int32", 1)}), sctx)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,123 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=invalid-name, missing-function-docstring
|
||||
"""Tests for the non-bulk CTA-level copy_async dispatch (vectorized load)."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx.layout import S, TileLayout
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"task",
|
||||
[
|
||||
################ A[0:128, 0:32] -> A_smem[0:128, 0:32] -> B[0:128, 0:32] ################
|
||||
(
|
||||
(128, 32), # g_shape
|
||||
(128, 32), # s_shape
|
||||
(0, 0), # g_st
|
||||
(128, 32), # g_extent
|
||||
32, # thread_cnt
|
||||
TileLayout(S[128, 32]), # layoutA
|
||||
TileLayout(S[128, 32]), # layoutB
|
||||
TileLayout(S[128, 32]), # layoutS
|
||||
),
|
||||
################ A[32:64, 32:64] -> A_smem[0:32, 0:32] -> B[32:64, 32:64] ################
|
||||
(
|
||||
(64, 64), # g_shape
|
||||
(32, 32), # s_shape
|
||||
(32, 0), # g_st
|
||||
(32, 32), # g_extent
|
||||
32, # thread_cnt
|
||||
TileLayout(S[64, 64]), # layoutA
|
||||
TileLayout(S[64, 64]), # layoutB
|
||||
TileLayout(S[32, 32]), # layoutS
|
||||
),
|
||||
################ A[0:1, 0:32, 0:32] -> A_smem[0:32, 0:32] -> B[0:1, 0:32, 0:32] ################ # noqa: E501
|
||||
(
|
||||
(4, 32, 32), # g_shape
|
||||
(32, 32), # s_shape
|
||||
(0, 0, 0), # g_st
|
||||
(1, 32, 32), # g_extent
|
||||
32, # thread_cnt
|
||||
TileLayout(S[4, 32, 32]), # layoutA
|
||||
TileLayout(S[4, 32, 32]), # layoutB
|
||||
TileLayout(S[32, 32]), # layoutS
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize(
|
||||
"dtype", ["int8", "float8_e4m3fn", "float8_e5m2", "float16", "bfloat16", "float32"]
|
||||
)
|
||||
def test_copy_g2s_s2g_cta_vec_load(task, dtype):
|
||||
g_shape, s_shape, g_st, g_extent, thread_cnt, layoutA, layoutB, layoutS = task
|
||||
|
||||
r_smem = list(slice(None) for i in range(len(s_shape)))
|
||||
r_gmem = list(slice(g_st[i], g_st[i] + g_extent[i]) for i in range(len(g_shape)))
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy_async(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, g_shape, dtype, layout=layoutA)
|
||||
B = T.match_buffer(B_ptr, g_shape, dtype, layout=layoutB)
|
||||
|
||||
T.device_entry()
|
||||
cta_id = T.cta_id([1])
|
||||
tid = T.thread_id([thread_cnt])
|
||||
A_smem = T.alloc_buffer(s_shape, dtype, scope="shared", layout=layoutS)
|
||||
|
||||
Tx.cta.copy_async(A_smem[tuple(r_smem)], A[tuple(r_gmem)], dispatch="ldgsts")
|
||||
T.ptx.cp_async.commit_group()
|
||||
T.ptx.cp_async.wait_group()
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(B[tuple(r_gmem)], A_smem[tuple(r_smem)])
|
||||
# fmt: on
|
||||
|
||||
np_dtype = tvm.testing.np_dtype_from_str(dtype)
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy_async})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
np.random.seed(0)
|
||||
A_np = np.random.rand(*g_shape).astype(np_dtype)
|
||||
B_np = np.zeros(g_shape, dtype=np_dtype)
|
||||
|
||||
B_ref = B_np.copy()
|
||||
B_ref[tuple(r_gmem)] = A_np[tuple(r_gmem)]
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
np.testing.assert_allclose(B_ref, B.numpy())
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,480 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=invalid-name, missing-function-docstring
|
||||
"""End-to-end tests for the smem->tmem (tcgen05.cp.32x128b.warpx4) dispatch.
|
||||
|
||||
The new dispatch requires the user to declare the t buffer with an
|
||||
explicit ``R[4 : 32@TLane]`` indicating warpx4 broadcast — i.e., t.shape[lane] = 32
|
||||
with replica 4 → 128 physical lanes.
|
||||
|
||||
Run with: pytest test_smem_tmem_dispatch.py -n 8 -v
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx.cuda.operator.tile_primitive.tma_utils import SwizzleMode, mma_shared_layout
|
||||
from tvm.tirx.layout import R, S, TCol, TileLayout, TLane
|
||||
|
||||
T_LAY_BASIC = TileLayout(S[(32, 16) : (1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane])
|
||||
|
||||
|
||||
def _make_2d_kernel(
|
||||
s_full,
|
||||
t_full,
|
||||
s_full_shape,
|
||||
t_full_shape,
|
||||
s_r0,
|
||||
s_r1,
|
||||
s_c0,
|
||||
s_c1,
|
||||
t_r0,
|
||||
t_r1,
|
||||
t_c0,
|
||||
t_c1,
|
||||
dtype,
|
||||
cta_group=1,
|
||||
):
|
||||
"""2D variant: SMEM/TMEM are both 2D; copy a rectangular sub-region."""
|
||||
n_tmem_cols_total = max(32, t_full_shape[-1])
|
||||
OUT_LANES = 32
|
||||
OUT_BYTES = 16
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle):
|
||||
A = T.match_buffer(A_ptr, s_full_shape, dtype)
|
||||
B = T.match_buffer(B_ptr, (OUT_LANES, OUT_BYTES), dtype)
|
||||
T.device_entry()
|
||||
warp_id = T.warp_id([4])
|
||||
wg_id = T.warpgroup_id([1])
|
||||
tid_in_wg = T.thread_id_in_wg([128])
|
||||
lane_id = T.lane_id([32])
|
||||
A_smem = T.alloc_buffer(s_full_shape, dtype, scope="shared", layout=s_full, align=1024)
|
||||
tmem_addr = T.alloc_shared([1], "uint32")
|
||||
cp_mbar = T.alloc_shared([1], "uint64")
|
||||
if wg_id == 0:
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.alloc(
|
||||
T.address_of(tmem_addr),
|
||||
n_cols=n_tmem_cols_total,
|
||||
cta_group=cta_group,
|
||||
)
|
||||
if tid_in_wg == 0:
|
||||
T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1)
|
||||
T.ptx.fence.proxy_async("shared::cta")
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(A_smem[:, :], A[:, :])
|
||||
T.cuda.cta_sync()
|
||||
tmem = T.decl_buffer(
|
||||
t_full_shape,
|
||||
dtype,
|
||||
scope="tmem",
|
||||
allocated_addr=tmem_addr[0],
|
||||
layout=t_full,
|
||||
)
|
||||
if tid_in_wg == 0:
|
||||
Tx.copy_async(
|
||||
tmem[t_r0:t_r1, t_c0:t_c1],
|
||||
A_smem[s_r0:s_r1, s_c0:s_c1],
|
||||
cta_group=cta_group,
|
||||
)
|
||||
T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=cta_group)
|
||||
T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0)
|
||||
T.cuda.cta_sync()
|
||||
T.ptx.tcgen05.fence.after_thread_sync()
|
||||
if warp_id == 0:
|
||||
reg = T.alloc_buffer((4,), "uint32", scope="local")
|
||||
for i in range(4):
|
||||
T.ptx.tcgen05.ld(
|
||||
tmem.allocated_addr[0],
|
||||
reg[i],
|
||||
shape="32x32b",
|
||||
num=1,
|
||||
row=0,
|
||||
col=i,
|
||||
)
|
||||
T.ptx.tcgen05.wait.ld()
|
||||
B_bytes = reg.view(dtype)
|
||||
for i in range(OUT_BYTES):
|
||||
B[lane_id, i] = B_bytes[i]
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
|
||||
T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=cta_group)
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
def _make_3d_4tile_kernel(s_full, t_full, s_full_shape, t_full_shape, dtype, cta_group=1):
|
||||
"""3D variant: 4 stacked tiles (NVFP4-style multi-cp test)."""
|
||||
n_tmem_cols_total = max(32, t_full_shape[-1])
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle):
|
||||
A = T.match_buffer(A_ptr, s_full_shape, dtype)
|
||||
B = T.match_buffer(B_ptr, (32, 16), dtype)
|
||||
T.device_entry()
|
||||
warp_id = T.warp_id([4])
|
||||
wg_id = T.warpgroup_id([1])
|
||||
tid_in_wg = T.thread_id_in_wg([128])
|
||||
lane_id = T.lane_id([32])
|
||||
A_smem = T.alloc_buffer(s_full_shape, dtype, scope="shared", layout=s_full, align=1024)
|
||||
tmem_addr = T.alloc_shared([1], "uint32")
|
||||
cp_mbar = T.alloc_shared([1], "uint64")
|
||||
if wg_id == 0:
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.alloc(
|
||||
T.address_of(tmem_addr),
|
||||
n_cols=n_tmem_cols_total,
|
||||
cta_group=cta_group,
|
||||
)
|
||||
if tid_in_wg == 0:
|
||||
T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1)
|
||||
T.ptx.fence.proxy_async("shared::cta")
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(A_smem[:, :, :], A[:, :, :])
|
||||
T.cuda.cta_sync()
|
||||
tmem = T.decl_buffer(
|
||||
t_full_shape,
|
||||
dtype,
|
||||
scope="tmem",
|
||||
allocated_addr=tmem_addr[0],
|
||||
layout=t_full,
|
||||
)
|
||||
if tid_in_wg == 0:
|
||||
Tx.copy_async(
|
||||
tmem[:, :, :],
|
||||
A_smem[:, :, :],
|
||||
cta_group=cta_group,
|
||||
)
|
||||
T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=cta_group)
|
||||
T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0)
|
||||
T.cuda.cta_sync()
|
||||
T.ptx.tcgen05.fence.after_thread_sync()
|
||||
if warp_id == 0:
|
||||
reg = T.alloc_buffer((4,), "uint32", scope="local")
|
||||
for i in range(4):
|
||||
T.ptx.tcgen05.ld(
|
||||
tmem.allocated_addr[0],
|
||||
reg[i],
|
||||
shape="32x32b",
|
||||
num=1,
|
||||
row=0,
|
||||
col=i,
|
||||
)
|
||||
T.ptx.tcgen05.wait.ld()
|
||||
B_bytes = reg.view(dtype)
|
||||
for i in range(16):
|
||||
B[lane_id, i] = B_bytes[i]
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
|
||||
T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=cta_group)
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
def _run_2d(s_full, t_full, s_full_shape, s_region, dtype, A_init, expected):
|
||||
s_r0, s_r1 = s_region[0]
|
||||
s_c0, s_c1 = s_region[1]
|
||||
kernel = _make_2d_kernel(
|
||||
s_full, t_full, s_full_shape, [32, 16], s_r0, s_r1, s_c0, s_c1, 0, 32, 0, 16, dtype
|
||||
)
|
||||
return _execute(kernel, A_init, expected)
|
||||
|
||||
|
||||
def _run_3d_4tile(s_full, t_full, s_full_shape, dtype, A_init, expected):
|
||||
kernel = _make_3d_4tile_kernel(s_full, t_full, s_full_shape, s_full_shape, dtype)
|
||||
return _execute(kernel, A_init, expected)
|
||||
|
||||
|
||||
def _execute(kernel, A_init, expected):
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
|
||||
B_np = np.zeros((32, 16), dtype=A_init.dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_init, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
B_out = B.numpy()
|
||||
assert np.array_equal(B_out, expected), (
|
||||
f"mismatch:\nlane 0 expected={expected[0].tolist()}\n"
|
||||
f" got ={B_out[0].tolist()}"
|
||||
)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
|
||||
@pytest.mark.parametrize(
|
||||
"name,s_full,s_full_shape,s_region",
|
||||
[
|
||||
("sw0_plain_atom_aligned", TileLayout(S[(32, 16) : (16, 1)]), [32, 16], [(0, 32), (0, 16)]),
|
||||
(
|
||||
"sw1_32B_atom",
|
||||
mma_shared_layout("uint8", SwizzleMode.SWIZZLE_32B_ATOM, [32, 32]),
|
||||
[32, 32],
|
||||
[(0, 32), (0, 16)],
|
||||
),
|
||||
(
|
||||
"sw2_64B_atom",
|
||||
mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 64]),
|
||||
[32, 64],
|
||||
[(0, 32), (0, 16)],
|
||||
),
|
||||
(
|
||||
"sw3_128B_atom",
|
||||
mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [32, 128]),
|
||||
[32, 128],
|
||||
[(0, 32), (0, 16)],
|
||||
),
|
||||
(
|
||||
"sw3_64x128_corner",
|
||||
mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]),
|
||||
[64, 128],
|
||||
[(0, 32), (0, 16)],
|
||||
),
|
||||
(
|
||||
"sw3_64x128_atom_row_8",
|
||||
mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]),
|
||||
[64, 128],
|
||||
[(8, 40), (0, 16)],
|
||||
),
|
||||
(
|
||||
"sw2_32x256_col_64",
|
||||
mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 256]),
|
||||
[32, 256],
|
||||
[(0, 32), (64, 80)],
|
||||
),
|
||||
(
|
||||
"sw0_M_atom_major_4_0",
|
||||
TileLayout(S[(8, 8, 2, 16) : (128, 16, 1024, 1)]),
|
||||
[64, 32],
|
||||
[(4, 36), (0, 16)],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_single_cp(name, s_full, s_full_shape, s_region):
|
||||
A_np = np.arange(int(np.prod(s_full_shape)), dtype=np.uint8).reshape(s_full_shape)
|
||||
r0, r1 = s_region[0]
|
||||
c0, c1 = s_region[1]
|
||||
expected = A_np[r0:r1, c0:c1]
|
||||
_run_2d(s_full, T_LAY_BASIC, s_full_shape, s_region, "uint8", A_np, expected)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
|
||||
def test_multi_cp_sw0_4tiles():
|
||||
s_full = TileLayout(S[(4, 32, 16) : (512, 16, 1)])
|
||||
t_full = TileLayout(S[(4, 32, 16) : (16 @ TCol, 1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane])
|
||||
A_np = (np.arange(4 * 32 * 16, dtype=np.int32) & 0xFF).astype(np.uint8).reshape(4, 32, 16)
|
||||
expected = A_np[0]
|
||||
_run_3d_4tile(s_full, t_full, [4, 32, 16], "uint8", A_np, expected)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
|
||||
def test_align_middle_2_to_1_nvfp4_sfb():
|
||||
"""SFB-style nvfp4 case: TMEM mid canonicalizes to single iter
|
||||
(16@TCol + 4@TCol merge), but SMEM mid stays as 2 iters
|
||||
(stride 512 + stride 2048 — outer/inner reversed so canon can't merge).
|
||||
Exercises ``_align_middles`` union-cut algorithm.
|
||||
|
||||
Layout shapes mirror SFB nvfp4 with PIPE=1, SFB_n_chunks=2,
|
||||
MMA_K_BLOCKS=4, sf_mma_k=4.
|
||||
"""
|
||||
# SMEM: (2, 4, 32, 4, 4) extents, strides (2048, 4, 16, 512, 1)
|
||||
# — N_chunk outer (stride 2048), then sub-warp tile (4, stride 4), lane
|
||||
# (32, stride 16), K_block (4, stride 512), sf_mma_k (4, stride 1).
|
||||
# Mid post-canon = [(4, 512), (2, 2048)] — non-mergeable in this order.
|
||||
s_full = TileLayout(S[(2, 4, 32, 4, 4) : (2048, 4, 16, 512, 1)])
|
||||
# TMEM: SFB-style 5-axis layout. K_outer (4, 4@TCol) and N_chunk
|
||||
# (2, 16@TCol) merge into single mid iter (8, 4@TCol).
|
||||
t_full = TileLayout(
|
||||
S[(2, 4, 32, 4, 4) : (16 @ TCol, 4 @ TCol, 1 @ TLane, 32 @ TCol, 1 @ TCol)]
|
||||
+ R[4 : 32 @ TLane]
|
||||
)
|
||||
s_full_shape = [256, 16]
|
||||
t_full_shape = [256, 16]
|
||||
n_tmem_cols_total = max(32, 32) # SFB occupies 32 cols total (8*4 elements / 4 epc)
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle):
|
||||
A = T.match_buffer(A_ptr, s_full_shape, "uint8")
|
||||
B = T.match_buffer(B_ptr, (32, 16), "uint8")
|
||||
T.device_entry()
|
||||
warp_id = T.warp_id([4])
|
||||
wg_id = T.warpgroup_id([1])
|
||||
tid_in_wg = T.thread_id_in_wg([128])
|
||||
lane_id = T.lane_id([32])
|
||||
A_smem = T.alloc_buffer(s_full_shape, "uint8", scope="shared", layout=s_full, align=1024)
|
||||
tmem_addr = T.alloc_shared([1], "uint32")
|
||||
cp_mbar = T.alloc_shared([1], "uint64")
|
||||
if wg_id == 0:
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=n_tmem_cols_total, cta_group=1)
|
||||
if tid_in_wg == 0:
|
||||
T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1)
|
||||
T.ptx.fence.proxy_async("shared::cta")
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(A_smem[:, :], A[:, :])
|
||||
T.cuda.cta_sync()
|
||||
tmem = T.decl_buffer(
|
||||
t_full_shape,
|
||||
"uint8",
|
||||
scope="tmem",
|
||||
allocated_addr=tmem_addr[0],
|
||||
layout=t_full,
|
||||
)
|
||||
if tid_in_wg == 0:
|
||||
Tx.copy_async(tmem[:, :], A_smem[:, :], cta_group=1)
|
||||
T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=1)
|
||||
T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0)
|
||||
T.cuda.cta_sync()
|
||||
T.ptx.tcgen05.fence.after_thread_sync()
|
||||
if warp_id == 0:
|
||||
reg = T.alloc_buffer((4,), "uint32", scope="local")
|
||||
for i in range(4):
|
||||
T.ptx.tcgen05.ld(
|
||||
tmem.allocated_addr[0],
|
||||
reg[i],
|
||||
shape="32x32b",
|
||||
num=1,
|
||||
row=0,
|
||||
col=i,
|
||||
)
|
||||
T.ptx.tcgen05.wait.ld()
|
||||
B_bytes = reg.view("uint8")
|
||||
for i in range(16):
|
||||
B[lane_id, i] = B_bytes[i]
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
|
||||
T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=1)
|
||||
|
||||
A_np = (np.arange(256 * 16, dtype=np.int32) & 0xFF).astype(np.uint8).reshape(256, 16)
|
||||
|
||||
# Compute expected: for each (lane=L in 0..32, byte b in 0..15), the
|
||||
# tcgen05.ld reads physical (TLane=L, TCol=b). We must invert the TMEM
|
||||
# layout to find which logical (m, k) is at that physical position, then
|
||||
# expected[L, b] = A[m, k].
|
||||
# Layout shard iters (i0..i4) with extents (2, 4, 32, 4, 4) and TMEM
|
||||
# strides (16, 4, 1@TLane, 32, 1) — only TLane and TCol contribute.
|
||||
# For (TLane=L, TCol=p) with L in 0..32, replica r=0:
|
||||
# i2 = L; remaining iters (i0, i1, i3, i4) contribute to TCol:
|
||||
# p = 16*i0 + 4*i1 + 32*i3 + i4
|
||||
# For p in 0..15 only i1 and i4 vary (i0 = i3 = 0):
|
||||
# i1 = p // 4, i4 = p % 4
|
||||
# Logical buffer index: rev row-major over iter coords following shard order.
|
||||
# Shard order outer→inner: (i0, i1, i2, i3, i4) with extents (2, 4, 32, 4, 4).
|
||||
# Logical buffer index = i0*(4*32*4*4) + i1*(32*4*4) + i2*(4*4) + i3*4 + i4
|
||||
expected = np.zeros((32, 16), dtype=np.uint8)
|
||||
for L in range(32):
|
||||
for p in range(16):
|
||||
i0 = 0
|
||||
i3 = 0
|
||||
i1 = p // 4
|
||||
i4 = p % 4
|
||||
i2 = L
|
||||
logical = i0 * (4 * 32 * 4 * 4) + i1 * (32 * 4 * 4) + i2 * (4 * 4) + i3 * 4 + i4
|
||||
m, k = divmod(logical, 16)
|
||||
expected[L, p] = A_np[m, k]
|
||||
|
||||
_execute(kernel, A_np, expected)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
|
||||
@pytest.mark.parametrize(
|
||||
"bad",
|
||||
[
|
||||
pytest.param(
|
||||
(
|
||||
"sw3_mid_atom_row",
|
||||
mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]),
|
||||
[64, 128],
|
||||
[(4, 36), (0, 16)],
|
||||
),
|
||||
id="sw3_mid_atom_row",
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
"sw2_mid_atom_col",
|
||||
mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 128]),
|
||||
[32, 128],
|
||||
[(0, 32), (32, 48)],
|
||||
),
|
||||
id="sw2_mid_atom_col",
|
||||
),
|
||||
pytest.param(
|
||||
("sw0_row_stride_64", TileLayout(S[(64, 64) : (64, 1)]), [64, 64], [(4, 36), (0, 16)]),
|
||||
id="sw0_row_stride_64",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_dispatch_rejects_bad_inputs(bad):
|
||||
"""Configurations where cp 32x128b cannot read the user's intended sub-tile.
|
||||
Compilation should fail with a clear ValueError from the dispatch."""
|
||||
name, s_full, s_full_shape, s_region = bad
|
||||
s_r0, s_r1 = s_region[0]
|
||||
s_c0, s_c1 = s_region[1]
|
||||
kernel = _make_2d_kernel(
|
||||
s_full, T_LAY_BASIC, s_full_shape, [32, 16], s_r0, s_r1, s_c0, s_c1, 0, 32, 0, 16, "uint8"
|
||||
)
|
||||
with pytest.raises(Exception):
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
|
||||
|
||||
|
||||
def test_multi_cp_encodes_descriptor_once_and_patches_addr():
|
||||
"""Compile-only regression for the shared-descriptor cp path.
|
||||
|
||||
A multi-tile smem->tmem copy encodes ONE SMEM matrix descriptor template at
|
||||
SMEM base 0 (so the cache key no longer depends on the buffer identity) and
|
||||
patches its 14-bit address field per cp via ``cvta(addr) >> 4 & 0x3FFF``,
|
||||
instead of re-encoding a descriptor per tile. Verifies the 4-tile copy emits
|
||||
a single ``encode_matrix_descriptor`` reused across four
|
||||
``tcgen05.cp.32x128b.warpx4`` issues, each with the address-field patch.
|
||||
"""
|
||||
s_full = TileLayout(S[(4, 32, 16) : (512, 16, 1)])
|
||||
t_full = TileLayout(S[(4, 32, 16) : (16 @ TCol, 1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane])
|
||||
kernel = _make_3d_4tile_kernel(s_full, t_full, [4, 32, 16], [4, 32, 16], "uint8")
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
|
||||
src = mod.mod.imports[0].inspect_source()
|
||||
|
||||
assert "tcgen05.cp.cta_group::1.32x128b.warpx4" in src, f"cp not emitted; src=\n{src}"
|
||||
# Descriptor encoded once (single matrix-descriptor encode call), then
|
||||
# reused with a per-cp 14-bit SMEM address patch (0x3FFF == 16383 mask).
|
||||
assert "16383" in src, "expected 14-bit SMEM address-field patch (0x3FFF mask)"
|
||||
assert src.count("cp_desc[0] &") == 4, (
|
||||
f"expected 4 address-patched cp's reusing one cp_desc; got "
|
||||
f"{src.count('cp_desc[0] &')}\nsrc=\n{src}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,344 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=invalid-name, missing-function-docstring
|
||||
"""Tests for the TMEM copy_async dispatch (tcgen05-based tmem<->reg and smem<->tmem)."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx.layout import S, TCol, TileLayout, TLane
|
||||
from tvm.tirx.layout import tid_in_wg as axis_tid_in_wg
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
|
||||
@pytest.mark.parametrize("dtype", ["float16", "float32"])
|
||||
@pytest.mark.parametrize("width_32b", [4, 8, 16, 32])
|
||||
def test_copy_tmem2reg_async(dtype, width_32b):
|
||||
"""Test async tmem<->local copy using copy_async instead of copy.
|
||||
|
||||
This tests the new copy_async dispatch for tmem<->local that doesn't
|
||||
immediately wait after the operation, allowing for pipelining.
|
||||
"""
|
||||
|
||||
def next_power_of_2(x):
|
||||
"""Return the smallest power of 2 greater than or equal to x."""
|
||||
if x <= 1:
|
||||
return 1
|
||||
return 1 << (x - 1).bit_length()
|
||||
|
||||
bits = tvm.runtime.DataType(dtype).bits
|
||||
if 128 % bits != 0 or 32 % bits != 0:
|
||||
pytest.skip(f"dtype {dtype} is not supported")
|
||||
|
||||
WIDTH = width_32b * (32 // bits)
|
||||
VEC_LEN = 128 // bits
|
||||
if WIDTH % VEC_LEN != 0:
|
||||
pytest.skip(f"dtype {dtype} + width {width_32b} is not supported")
|
||||
|
||||
g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)])
|
||||
local_view = TileLayout(S[(128, WIDTH) : (1 @ axis_tid_in_wg, 1)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy_async_test(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (128, WIDTH), dtype)
|
||||
B = T.match_buffer(B_ptr, (128, WIDTH), dtype)
|
||||
|
||||
A_flat = A.view(-1)
|
||||
B_flat = B.view(-1)
|
||||
|
||||
T.device_entry()
|
||||
warp_id = T.warp_id([(128) // 32])
|
||||
cta_id = T.cta_id([2])
|
||||
wg_id = T.warpgroup_id([1])
|
||||
warp_id_in_wg = T.warp_id_in_wg([4])
|
||||
lane_id = T.lane_id([32])
|
||||
tid_in_wg = T.thread_id([128])
|
||||
|
||||
tmem_addr = T.alloc_shared([1], "uint32")
|
||||
|
||||
if wg_id == 0:
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501
|
||||
|
||||
T.tvm_storage_sync("shared")
|
||||
|
||||
tmem = T.decl_buffer((128, WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0],
|
||||
layout=TileLayout(S[(128, WIDTH) : (1 @ TLane, 1 @ TCol)]))
|
||||
|
||||
A_reg = T.alloc_local((WIDTH), dtype)
|
||||
B_reg = T.alloc_local((WIDTH), dtype)
|
||||
A_local = A_reg.view(128, WIDTH, layout=local_view)
|
||||
B_local = B_reg.view(128, WIDTH, layout=local_view)
|
||||
for i in range(WIDTH // VEC_LEN):
|
||||
g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
|
||||
Tx.copy(A_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN]) # noqa: E501
|
||||
for i in range(WIDTH):
|
||||
B_reg[i] = T.cast(0, dtype)
|
||||
T.cuda.cta_sync()
|
||||
|
||||
# A_local -> tmem (async)
|
||||
Tx.wg.copy_async(tmem[:, :], A_local[:, :])
|
||||
T.ptx.tcgen05.wait.st() # explicit wait
|
||||
T.cuda.cta_sync()
|
||||
|
||||
# tmem -> B_local (async)
|
||||
Tx.wg.copy_async(B_local[:, :], tmem[:, :])
|
||||
T.ptx.tcgen05.wait.ld() # explicit wait
|
||||
T.cuda.cta_sync()
|
||||
for i in range(WIDTH // VEC_LEN):
|
||||
g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
|
||||
Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN]) # noqa: E501
|
||||
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
|
||||
T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy_async_test})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH))
|
||||
B_np = np.zeros((128, WIDTH), dtype=dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
np.testing.assert_allclose(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Migrated from test_copy_sync.py: tmem<->reg round-trip via T.copy_async
|
||||
# (the kernels themselves are the actual async tmem dispatch tests; the
|
||||
# G↔L copies bookending them just stage data).
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
|
||||
@pytest.mark.parametrize("dtype", ["uint8", "float16", "float32"])
|
||||
@pytest.mark.parametrize("width_32b", [2, 4, 8, 16, 32, 64, 128])
|
||||
@pytest.mark.parametrize("offset_32b", [0, 3, 10])
|
||||
def test_copy_tmem2reg(dtype, width_32b, offset_32b):
|
||||
def next_power_of_2(x):
|
||||
if x <= 1:
|
||||
return 1
|
||||
return 1 << (x - 1).bit_length()
|
||||
|
||||
bits = tvm.runtime.DataType(dtype).bits
|
||||
if 128 % bits != 0 or 32 % bits != 0:
|
||||
pytest.skip(f"dtype {dtype} is not supported")
|
||||
|
||||
WIDTH = width_32b * (32 // bits)
|
||||
OFFSET = offset_32b * (32 // bits)
|
||||
VEC_LEN = 128 // bits
|
||||
if WIDTH % VEC_LEN != 0:
|
||||
pytest.skip(f"dtype {dtype} + width {width_32b} is not supported")
|
||||
|
||||
g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)])
|
||||
local_view = TileLayout(S[(128, WIDTH) : (1 @ axis_tid_in_wg, 1)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (128, WIDTH), dtype)
|
||||
B = T.match_buffer(B_ptr, (128, WIDTH), dtype)
|
||||
|
||||
A_flat = A.view(-1)
|
||||
B_flat = B.view(-1)
|
||||
|
||||
T.device_entry()
|
||||
warp_id = T.warp_id([(128) // 32])
|
||||
T.cta_id([2])
|
||||
wg_id = T.warpgroup_id([1])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
tid_in_wg = T.thread_id([128])
|
||||
|
||||
tmem_addr = T.alloc_shared([1], "uint32")
|
||||
|
||||
if wg_id == 0:
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(offset_32b + width_32b)), cta_group=1) # noqa: E501
|
||||
|
||||
T.tvm_storage_sync("shared")
|
||||
|
||||
tmem = T.decl_buffer((128, OFFSET + WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0], # noqa: E501
|
||||
layout=TileLayout(S[(128, OFFSET + WIDTH) : (1 @ TLane, 1 @ TCol)])) # noqa: E501
|
||||
|
||||
A_reg = T.alloc_local((WIDTH), dtype)
|
||||
B_reg = T.alloc_local((WIDTH), dtype)
|
||||
A_local = A_reg.view(128, WIDTH, layout=local_view)
|
||||
B_local = B_reg.view(128, WIDTH, layout=local_view)
|
||||
for i in range(WIDTH // VEC_LEN):
|
||||
g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
|
||||
Tx.copy(A_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN]) # noqa: E501
|
||||
for i in range(WIDTH):
|
||||
B_reg[i] = T.cast(0, dtype)
|
||||
T.cuda.cta_sync()
|
||||
|
||||
# A_local -> tmem
|
||||
Tx.wg.copy_async(tmem[:, OFFSET: OFFSET + WIDTH], A_local[:, :])
|
||||
T.ptx.tcgen05.wait.st()
|
||||
T.cuda.cta_sync()
|
||||
|
||||
# tmem -> B_local
|
||||
Tx.wg.copy_async(B_local[:, :], tmem[:, OFFSET: OFFSET + WIDTH])
|
||||
T.ptx.tcgen05.wait.ld()
|
||||
T.cuda.cta_sync()
|
||||
for i in range(WIDTH // VEC_LEN):
|
||||
g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
|
||||
Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN]) # noqa: E501
|
||||
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
|
||||
T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(offset_32b + width_32b)), cta_group=1) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy_sync})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH))
|
||||
B_np = np.zeros((128, WIDTH), dtype=dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
np.testing.assert_allclose(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
|
||||
@pytest.mark.parametrize("dtype", ["float16", "float32"])
|
||||
@pytest.mark.parametrize("width_32b", [4, 8, 16, 32])
|
||||
@pytest.mark.parametrize("local_offset_32b", [0, 2, 4])
|
||||
def test_copy_tmem2reg_sliced_local(dtype, width_32b, local_offset_32b):
|
||||
"""tmem<->local copy with a sliced local buffer region."""
|
||||
|
||||
def next_power_of_2(x):
|
||||
if x <= 1:
|
||||
return 1
|
||||
return 1 << (x - 1).bit_length()
|
||||
|
||||
bits = tvm.runtime.DataType(dtype).bits
|
||||
if 128 % bits != 0 or 32 % bits != 0:
|
||||
pytest.skip(f"dtype {dtype} is not supported")
|
||||
|
||||
WIDTH = width_32b * (32 // bits)
|
||||
LOCAL_OFFSET = local_offset_32b * (32 // bits)
|
||||
TOTAL_LOCAL_WIDTH = WIDTH + LOCAL_OFFSET
|
||||
VEC_LEN = 128 // bits
|
||||
if WIDTH % VEC_LEN != 0 or TOTAL_LOCAL_WIDTH % VEC_LEN != 0:
|
||||
pytest.skip(
|
||||
f"dtype {dtype} + width {width_32b} + offset {local_offset_32b} is not supported"
|
||||
)
|
||||
|
||||
g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)])
|
||||
local_view = TileLayout(S[(128, TOTAL_LOCAL_WIDTH) : (1 @ axis_tid_in_wg, 1)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (128, WIDTH), dtype)
|
||||
B = T.match_buffer(B_ptr, (128, WIDTH), dtype)
|
||||
|
||||
A_flat = A.view(-1)
|
||||
B_flat = B.view(-1)
|
||||
|
||||
T.device_entry()
|
||||
warp_id = T.warp_id([(128) // 32])
|
||||
T.cta_id([2])
|
||||
wg_id = T.warpgroup_id([1])
|
||||
T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
tid_in_wg = T.thread_id([128])
|
||||
|
||||
tmem_addr = T.alloc_shared([1], "uint32")
|
||||
|
||||
if wg_id == 0:
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501
|
||||
|
||||
T.tvm_storage_sync("shared")
|
||||
|
||||
tmem = T.decl_buffer((128, WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0],
|
||||
layout=TileLayout(S[(128, WIDTH) : (1 @ TLane, 1 @ TCol)]))
|
||||
|
||||
A_reg = T.alloc_local((TOTAL_LOCAL_WIDTH), dtype)
|
||||
B_reg = T.alloc_local((TOTAL_LOCAL_WIDTH), dtype)
|
||||
A_local = A_reg.view(128, TOTAL_LOCAL_WIDTH, layout=local_view)
|
||||
B_local = B_reg.view(128, TOTAL_LOCAL_WIDTH, layout=local_view)
|
||||
for i in range(WIDTH // VEC_LEN):
|
||||
g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
|
||||
Tx.copy(A_reg[LOCAL_OFFSET + i * VEC_LEN: LOCAL_OFFSET + i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN]) # noqa: E501
|
||||
for i in range(TOTAL_LOCAL_WIDTH):
|
||||
B_reg[i] = T.cast(0, dtype)
|
||||
T.cuda.cta_sync()
|
||||
|
||||
# A_local[sliced] -> tmem (use sliced region)
|
||||
Tx.wg.copy_async(tmem[:, 0:WIDTH], A_local[:, LOCAL_OFFSET:LOCAL_OFFSET + WIDTH])
|
||||
T.ptx.tcgen05.wait.st()
|
||||
T.cuda.cta_sync()
|
||||
|
||||
# tmem -> B_local[sliced] (use sliced region)
|
||||
Tx.wg.copy_async(B_local[:, LOCAL_OFFSET:LOCAL_OFFSET + WIDTH], tmem[:, 0:WIDTH])
|
||||
T.ptx.tcgen05.wait.ld()
|
||||
T.cuda.cta_sync()
|
||||
for i in range(WIDTH // VEC_LEN):
|
||||
g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
|
||||
Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[LOCAL_OFFSET + i * VEC_LEN: LOCAL_OFFSET + i * VEC_LEN + VEC_LEN]) # noqa: E501
|
||||
|
||||
if warp_id == 0:
|
||||
T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
|
||||
T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy_sync})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH))
|
||||
B_np = np.zeros((128, WIDTH), dtype=dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
np.testing.assert_allclose(B.numpy(), A_np)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,870 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
import re
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx.layout import S, TileLayout, wg_local_layout
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input",
|
||||
[
|
||||
######### basic test #########
|
||||
(
|
||||
(32, 32), # g_shape
|
||||
(0, 0), # st_a
|
||||
(0, 0), # st_b
|
||||
(0, 0), # st_res
|
||||
(32, 32), # extent_a
|
||||
(32, 32), # extent_b
|
||||
(32, 32), # extent_res
|
||||
64, # thread_cnt
|
||||
0, # device_id
|
||||
),
|
||||
######### offset test #########
|
||||
(
|
||||
(32, 8, 12), # g_shape
|
||||
(10, 0, 3), # st_a
|
||||
(14, 1, 4), # st_b
|
||||
(20, 0, 2), # st_res
|
||||
(5, 6, 7), # extent_a
|
||||
(5, 6, 7), # extent_b
|
||||
(5, 6, 7), # extent_res
|
||||
64, # thread_cnt
|
||||
0, # device_id
|
||||
),
|
||||
######### broadcast test #########
|
||||
(
|
||||
(32, 8, 12), # g_shape
|
||||
(10, 0, 3), # st_a
|
||||
(14, 1, 4), # st_b
|
||||
(20, 0, 2), # st_res
|
||||
(5, 6, 7), # extent_a
|
||||
(1, 6, 1), # extent_b
|
||||
(5, 6, 7), # extent_res
|
||||
64, # thread_cnt
|
||||
0, # device_id
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
|
||||
@pytest.mark.parametrize("operands_type", ["region_region", "region_const", "const_region"])
|
||||
@pytest.mark.parametrize("dtype", ["float16"])
|
||||
def test_binary_op_shared(input, op_type, operands_type, dtype):
|
||||
# skip test
|
||||
if op_type in ["sub", "fdiv"] and operands_type == "const_region":
|
||||
return
|
||||
|
||||
g_shape, st_a, st_b, st_res, ext_a, ext_b, ext_res, thread_cnt, device_id = input
|
||||
g_layout = s_layout = TileLayout(S[g_shape])
|
||||
|
||||
copy_slice = list(slice(None) for i in range(len(g_shape)))
|
||||
map_slice_a = list(slice(st_a[i], st_a[i] + ext_a[i]) for i in range(len(g_shape)))
|
||||
map_slice_b = list(slice(st_b[i], st_b[i] + ext_b[i]) for i in range(len(g_shape)))
|
||||
map_slice_res = list(slice(st_res[i], st_res[i] + ext_res[i]) for i in range(len(g_shape)))
|
||||
|
||||
const = T.float16(3.0) if dtype == "float16" else T.float32(3.0)
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary_op_region_region(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, g_shape, dtype, layout=g_layout)
|
||||
B = T.match_buffer(B_ptr, g_shape, dtype, layout=g_layout)
|
||||
|
||||
T.device_entry()
|
||||
cta_id = T.cta_id([1])
|
||||
tid = T.thread_id([thread_cnt])
|
||||
A_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=s_layout)
|
||||
B_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=s_layout)
|
||||
|
||||
Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
|
||||
Tx.cta.copy(B_smem[tuple(copy_slice)], B[tuple(copy_slice)])
|
||||
T.cuda.cta_sync()
|
||||
if op_type == "add":
|
||||
Tx.cta.add(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)]) # noqa: E501
|
||||
elif op_type == "sub":
|
||||
Tx.cta.sub(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)]) # noqa: E501
|
||||
elif op_type == "mul":
|
||||
Tx.cta.mul(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)]) # noqa: E501
|
||||
elif op_type == "fdiv":
|
||||
Tx.cta.fdiv(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)]) # noqa: E501
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(A[tuple(copy_slice)], A_smem[tuple(copy_slice)])
|
||||
|
||||
@T.prim_func
|
||||
def binary_op_const_region_or_region_const(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, g_shape, dtype, layout=g_layout)
|
||||
_B = T.match_buffer(B_ptr, g_shape, dtype, layout=g_layout)
|
||||
|
||||
T.device_entry()
|
||||
cta_id = T.cta_id([1])
|
||||
tid = T.thread_id([thread_cnt])
|
||||
A_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=s_layout)
|
||||
|
||||
Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
|
||||
T.cuda.cta_sync()
|
||||
if op_type == "add":
|
||||
if operands_type == "const_region":
|
||||
Tx.cta.add(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
|
||||
elif operands_type == "region_const":
|
||||
Tx.cta.add(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
|
||||
elif op_type == "sub":
|
||||
if operands_type == "const_region":
|
||||
Tx.cta.sub(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
|
||||
elif operands_type == "region_const":
|
||||
Tx.cta.sub(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
|
||||
elif op_type == "mul":
|
||||
if operands_type == "const_region":
|
||||
Tx.cta.mul(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
|
||||
elif operands_type == "region_const":
|
||||
Tx.cta.mul(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
|
||||
elif op_type == "fdiv":
|
||||
if operands_type == "const_region":
|
||||
Tx.cta.fdiv(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
|
||||
elif operands_type == "region_const":
|
||||
Tx.cta.fdiv(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(A[tuple(copy_slice)], A_smem[tuple(copy_slice)])
|
||||
# fmt: on
|
||||
|
||||
def get_prim_func(operands_type):
|
||||
if operands_type == "region_region":
|
||||
return binary_op_region_region
|
||||
elif operands_type in ["const_region", "region_const"]:
|
||||
return binary_op_const_region_or_region_const
|
||||
raise ValueError(f"operands_type={operands_type} is not supported")
|
||||
|
||||
def get_ref(A_np, B_np):
|
||||
A_ref = A_np.copy()
|
||||
if op_type == "add":
|
||||
if operands_type == "region_region":
|
||||
A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] + B_np[tuple(map_slice_b)]
|
||||
elif operands_type in ["const_region", "region_const"]:
|
||||
A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] + 3.0
|
||||
elif op_type == "sub":
|
||||
if operands_type == "region_region":
|
||||
A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] - B_np[tuple(map_slice_b)]
|
||||
elif operands_type in ["const_region", "region_const"]:
|
||||
A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] - 3.0
|
||||
elif op_type == "mul":
|
||||
if operands_type == "region_region":
|
||||
A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] * B_np[tuple(map_slice_b)]
|
||||
elif operands_type in ["const_region", "region_const"]:
|
||||
A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] * 3.0
|
||||
elif op_type == "fdiv":
|
||||
if operands_type == "region_region":
|
||||
A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] / B_np[tuple(map_slice_b)]
|
||||
elif operands_type in ["const_region", "region_const"]:
|
||||
A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] / 3.0
|
||||
|
||||
return A_ref
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
np.random.seed(0)
|
||||
A_np = np.random.rand(*g_shape).astype(dtype)
|
||||
B_np = np.random.rand(*g_shape).astype(dtype)
|
||||
mod = tvm.IRModule({"main": get_prim_func(operands_type)})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
print(f"compiled source code: {mod.mod.imports[0].inspect_source()}")
|
||||
A_ref = get_ref(A_np, B_np)
|
||||
atol = 1e-3
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(device_id)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
tvm.testing.assert_allclose(A_ref, A.numpy(), atol=atol)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_type", ["sub", "fdiv"])
|
||||
def test_binary_non_commutative_const_lhs_rejected(op_type):
|
||||
dtype = "float16"
|
||||
shape = (16, 16)
|
||||
layout = TileLayout(S[shape])
|
||||
const = T.float16(3.0)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
|
||||
@T.prim_func
|
||||
def bad_kernel() -> None:
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
_tid = T.thread_id([64])
|
||||
A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=layout)
|
||||
if op_type == "sub":
|
||||
Tx.cta.sub(A_smem, const, A_smem)
|
||||
elif op_type == "fdiv":
|
||||
Tx.cta.fdiv(A_smem, const, A_smem)
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": bad_kernel})
|
||||
tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("exec_scope", ["warp", "warpgroup"])
|
||||
@pytest.mark.parametrize("op_type", ["add", "mul"])
|
||||
def test_binary_op_shared_subcta_scope(exec_scope, op_type):
|
||||
"""Test binary ops in warp/warpgroup scope with shared memory."""
|
||||
dtype = "float16"
|
||||
n_warps = 4 if exec_scope == "warpgroup" else 1
|
||||
g_shape = (n_warps * 32, 8)
|
||||
tx_op = {
|
||||
("warp", "add"): Tx.warp.add,
|
||||
("warp", "mul"): Tx.warp.mul,
|
||||
("warpgroup", "add"): Tx.wg.add,
|
||||
("warpgroup", "mul"): Tx.wg.mul,
|
||||
}[(exec_scope, op_type)]
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, g_shape, dtype, layout=TileLayout(S[g_shape]))
|
||||
B = T.match_buffer(B_ptr, g_shape, dtype, layout=TileLayout(S[g_shape]))
|
||||
T.device_entry()
|
||||
warp_id = T.warp_id([(256) // 32])
|
||||
wg_id = T.warpgroup_id([(256) // 128])
|
||||
_bx = T.cta_id([1])
|
||||
_tid = T.thread_id([256])
|
||||
A_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=TileLayout(S[g_shape]))
|
||||
B_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=TileLayout(S[g_shape]))
|
||||
Tx.cta.copy(A_smem, A)
|
||||
Tx.cta.copy(B_smem, B)
|
||||
T.cuda.cta_sync()
|
||||
if exec_scope == "warp":
|
||||
if warp_id == 5:
|
||||
tx_op(A_smem, A_smem, B_smem)
|
||||
elif exec_scope == "warpgroup":
|
||||
if wg_id == 1:
|
||||
tx_op(A_smem, A_smem, B_smem)
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(A, A_smem)
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
np.random.seed(0)
|
||||
A_np = np.random.rand(*g_shape).astype(dtype)
|
||||
B_np = np.random.rand(*g_shape).astype(dtype)
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
np_op = {"add": np.add, "mul": np.multiply}[op_type]
|
||||
A_ref = np_op(A_np, B_np).astype(dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
tvm.testing.assert_allclose(A_ref, A.numpy(), atol=1e-3)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("exec_scope", ["cta", "warpgroup", "warp"])
|
||||
@pytest.mark.parametrize("rhs_kind", ["region", "broadcast", "const"])
|
||||
@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
|
||||
def test_binary_op_local_subcta_trivial(exec_scope, rhs_kind, op_type):
|
||||
dtype = "float16"
|
||||
m, n = 4, 8
|
||||
n_threads = 256 if exec_scope == "cta" else (128 if exec_scope == "warpgroup" else 32)
|
||||
# in this test, use warp3/warpgroup1 to test
|
||||
thr_str = 0 if exec_scope == "cta" else (128 if exec_scope == "warpgroup" else 32 * 3)
|
||||
a_shape = (n_threads, m, n)
|
||||
b_shape = (n_threads, m, n if rhs_kind == "region" else 1)
|
||||
c_shape = a_shape
|
||||
const = T.float16(1.25)
|
||||
tx_op = {"add": Tx.add, "sub": Tx.sub, "mul": Tx.mul, "fdiv": Tx.fdiv}[op_type]
|
||||
tid_in_scope_fn = {"cta": T.thread_id, "warpgroup": T.thread_id_in_wg, "warp": T.lane_id}[
|
||||
exec_scope
|
||||
]
|
||||
|
||||
@T.prim_func
|
||||
def kernel(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
|
||||
B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
|
||||
C = T.match_buffer(C_ptr, c_shape, dtype, layout=TileLayout(S[c_shape]))
|
||||
|
||||
T.device_entry()
|
||||
wg_id = T.warpgroup_id([(256) // 128])
|
||||
warp_id = T.warp_id([(256) // 32])
|
||||
_bx = T.cta_id([1])
|
||||
_tid = T.thread_id([256])
|
||||
tid_in_scope = tid_in_scope_fn([n_threads])
|
||||
b_n = T.meta_var(n if rhs_kind == "region" else 1)
|
||||
A_local = T.alloc_buffer((m, n), dtype, scope="local", layout=TileLayout(S[(m, n)]))
|
||||
C_local = T.alloc_buffer((m, n), dtype, scope="local", layout=TileLayout(S[(m, n)]))
|
||||
B_local = T.alloc_buffer((m, b_n), dtype, scope="local", layout=TileLayout(S[(m, b_n)]))
|
||||
|
||||
if thr_str <= _tid and _tid < thr_str + n_threads:
|
||||
for i in T.serial(m):
|
||||
for j in T.serial(n):
|
||||
A_local[i, j] = A[tid_in_scope, i, j]
|
||||
if rhs_kind != "const":
|
||||
for i in T.serial(m):
|
||||
for j in T.serial(b_n):
|
||||
B_local[i, j] = B[tid_in_scope, i, j]
|
||||
|
||||
if exec_scope == "cta":
|
||||
if rhs_kind == "const":
|
||||
tx_op(C_local, A_local, const)
|
||||
else:
|
||||
tx_op(C_local, A_local, B_local)
|
||||
elif exec_scope == "warpgroup":
|
||||
if wg_id == 1:
|
||||
if rhs_kind == "const":
|
||||
tx_op(C_local, A_local, const)
|
||||
else:
|
||||
tx_op(C_local, A_local, B_local)
|
||||
else:
|
||||
if warp_id == 3:
|
||||
if rhs_kind == "const":
|
||||
tx_op(C_local, A_local, const)
|
||||
else:
|
||||
tx_op(C_local, A_local, B_local)
|
||||
# T.cuda.cta_sync()
|
||||
|
||||
if thr_str <= _tid and _tid < thr_str + n_threads:
|
||||
for i in T.serial(m):
|
||||
for j in T.serial(n):
|
||||
C[tid_in_scope, i, j] = C_local[i, j]
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
np.random.seed(0)
|
||||
A_np = np.random.rand(*a_shape).astype(dtype)
|
||||
B_np = np.random.rand(*b_shape).astype(dtype)
|
||||
C_np = np.zeros(c_shape, dtype=dtype)
|
||||
mod = tvm.IRModule({"main": kernel})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
print(f"compiled source code: {mod.mod.imports[0].inspect_source()}")
|
||||
np_op = {"add": np.add, "sub": np.subtract, "mul": np.multiply, "fdiv": np.divide}[op_type]
|
||||
if rhs_kind == "region":
|
||||
C_ref = np_op(A_np, B_np)
|
||||
elif rhs_kind == "broadcast":
|
||||
C_ref = np_op(A_np, np.repeat(B_np, n, axis=2))
|
||||
else:
|
||||
C_ref = np_op(A_np, const.value)
|
||||
atol = 1e-2 if op_type == "fdiv" else 1e-3
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
C = tvm.runtime.tensor(C_np, dev)
|
||||
mod(A, B, C)
|
||||
tvm.testing.assert_allclose(C_ref, C.numpy(), atol=atol)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input",
|
||||
[
|
||||
######### basic test #########
|
||||
(
|
||||
(64, 32), # a_shape
|
||||
(64, 32), # b_shape
|
||||
(64, 32), # res_shape
|
||||
64, # thread_cnt
|
||||
0, # device_id
|
||||
),
|
||||
######### broadcast test #########
|
||||
(
|
||||
(32, 5, 4), # a_shape
|
||||
(32, 1, 4), # b_shape
|
||||
(32, 5, 4), # res_shape
|
||||
32, # thread_cnt (≥ warp size so sctx.intra at cta scope models cleanly)
|
||||
0, # device_id
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("storage_scope", ["shared", "local"])
|
||||
@pytest.mark.parametrize("exec_scope", ["cta", "thread"])
|
||||
@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
|
||||
@pytest.mark.parametrize("dtype", ["float16"])
|
||||
def test_binary_op_vectorized(input, storage_scope, exec_scope, op_type, dtype):
|
||||
a_shape, b_shape, res_shape, thread_cnt, device_id = input
|
||||
tx_op = {"add": Tx.add, "sub": Tx.sub, "mul": Tx.mul, "fdiv": Tx.fdiv}[op_type]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def test_binary_cta(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
|
||||
B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
|
||||
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
tx = T.thread_id([thread_cnt])
|
||||
if storage_scope == "shared":
|
||||
A_smem = T.alloc_buffer(
|
||||
a_shape, dtype, scope="shared", layout=TileLayout(S[a_shape])
|
||||
)
|
||||
B_smem = T.alloc_buffer(
|
||||
b_shape, dtype, scope="shared", layout=TileLayout(S[b_shape])
|
||||
)
|
||||
Tx.cta.copy(A_smem, A)
|
||||
Tx.cta.copy(B_smem, B)
|
||||
T.cuda.cta_sync()
|
||||
tx_op(A_smem, A_smem, B_smem)
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(A, A_smem)
|
||||
if storage_scope == "local":
|
||||
A_local = T.alloc_buffer(
|
||||
a_shape[1:], dtype, scope="local", layout=TileLayout(S[a_shape[1:]])
|
||||
)
|
||||
B_local = T.alloc_buffer(
|
||||
b_shape[1:], dtype, scope="local", layout=TileLayout(S[b_shape[1:]])
|
||||
)
|
||||
Tx.copy(A_local, A[tx])
|
||||
Tx.copy(B_local, B[tx])
|
||||
tx_op(A_local, A_local, B_local)
|
||||
Tx.copy(A[tx], A_local)
|
||||
|
||||
@T.prim_func
|
||||
def test_binary_thread(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
|
||||
B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
|
||||
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
tx = T.thread_id([thread_cnt])
|
||||
if storage_scope == "shared":
|
||||
A_smem = T.alloc_buffer(
|
||||
a_shape, dtype, scope="shared", layout=TileLayout(S[a_shape])
|
||||
)
|
||||
B_smem = T.alloc_buffer(
|
||||
b_shape, dtype, scope="shared", layout=TileLayout(S[b_shape])
|
||||
)
|
||||
Tx.copy(A_smem, A)
|
||||
Tx.copy(B_smem, B)
|
||||
T.cuda.cta_sync()
|
||||
tx_op(A_smem, A_smem, B_smem)
|
||||
T.cuda.cta_sync()
|
||||
Tx.copy(A, A_smem)
|
||||
elif storage_scope == "local":
|
||||
A_local = T.alloc_buffer(
|
||||
a_shape[1:], dtype, scope="local", layout=TileLayout(S[a_shape[1:]])
|
||||
)
|
||||
B_local = T.alloc_buffer(
|
||||
b_shape[1:], dtype, scope="local", layout=TileLayout(S[b_shape[1:]])
|
||||
)
|
||||
Tx.copy(A_local, A[tx])
|
||||
Tx.copy(B_local, B[tx])
|
||||
tx_op(A_local, A_local, B_local)
|
||||
Tx.copy(A[tx], A_local)
|
||||
# fmt: on
|
||||
|
||||
def get_prim_func():
|
||||
if exec_scope == "cta":
|
||||
return test_binary_cta
|
||||
elif exec_scope == "thread":
|
||||
return test_binary_thread
|
||||
else:
|
||||
raise ValueError(f"exec_scope={exec_scope} is not supported")
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
np.random.seed(0)
|
||||
A_np = np.random.rand(*a_shape).astype(dtype)
|
||||
B_np = np.random.rand(*b_shape).astype(dtype)
|
||||
mod = tvm.IRModule({"main": get_prim_func()})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
print(f"compiled source code: {mod.mod.imports[0].inspect_source()}")
|
||||
np_op = {"add": np.add, "sub": np.subtract, "mul": np.multiply, "fdiv": np.divide}[op_type]
|
||||
A_ref = np_op(A_np, B_np)
|
||||
atol = 1e-2 if op_type == "fdiv" else 1e-3
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(device_id)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
tvm.testing.assert_allclose(A_ref, A.numpy(), atol=atol)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("op_type", ["add", "sub", "mul"])
|
||||
def test_binary_op_packed_f32x2_auto_dispatch(op_type):
|
||||
target = tvm.target.Target("cuda")
|
||||
arch = target.arch if hasattr(target, "arch") else ""
|
||||
if not arch.startswith("sm_"):
|
||||
pytest.skip(f"unknown target arch: {arch}")
|
||||
sm_digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
|
||||
if not sm_digits:
|
||||
pytest.skip(f"cannot parse target arch: {arch}")
|
||||
sm_version = int(sm_digits)
|
||||
if sm_version < 100:
|
||||
pytest.skip(f"packed_f32x2 auto-dispatch requires sm_100+, got {arch}")
|
||||
|
||||
a_shape, b_shape = (64, 32), (64, 32)
|
||||
dtype = "float32"
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
|
||||
B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
|
||||
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
tx = T.thread_id([64])
|
||||
A_local = T.alloc_buffer(
|
||||
a_shape[1:], dtype, scope="local", layout=TileLayout(S[a_shape[1:]])
|
||||
)
|
||||
B_local = T.alloc_buffer(
|
||||
b_shape[1:], dtype, scope="local", layout=TileLayout(S[b_shape[1:]])
|
||||
)
|
||||
Tx.copy(A_local, A[tx])
|
||||
Tx.copy(B_local, B[tx])
|
||||
if op_type == "add":
|
||||
Tx.add(A_local, A_local, B_local)
|
||||
elif op_type == "sub":
|
||||
Tx.sub(A_local, A_local, B_local)
|
||||
elif op_type == "mul":
|
||||
Tx.mul(A_local, A_local, B_local)
|
||||
Tx.copy(A[tx], A_local)
|
||||
|
||||
with target:
|
||||
np.random.seed(0)
|
||||
A_np = np.random.rand(*a_shape).astype(dtype)
|
||||
B_np = np.random.rand(*b_shape).astype(dtype)
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
src = mod.mod.imports[0].inspect_source()
|
||||
ptx_pat = {
|
||||
"add": r"add\.[a-z]+\.ftz\.f32x2",
|
||||
"sub": r"sub\.[a-z]+\.ftz\.f32x2",
|
||||
"mul": r"mul\.[a-z]+\.ftz\.f32x2",
|
||||
}[op_type]
|
||||
builtin_pat = {
|
||||
"add": r"tvm_builtin_ptx_add_packed_",
|
||||
"sub": r"tvm_builtin_ptx_sub_packed_",
|
||||
"mul": r"tvm_builtin_ptx_mul_packed_",
|
||||
}[op_type]
|
||||
assert re.search(ptx_pat, src) or re.search(builtin_pat, src), src
|
||||
if op_type == "add":
|
||||
A_ref = A_np + B_np
|
||||
elif op_type == "sub":
|
||||
A_ref = A_np - B_np
|
||||
elif op_type == "mul":
|
||||
A_ref = A_np * B_np
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
tvm.testing.assert_allclose(A_ref, A.numpy(), atol=1e-3)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("op_name", ["add", "sub", "mul"])
|
||||
def test_binary_op_warpgroup_wg_local_layout(op_name):
|
||||
dtype = "float32"
|
||||
rows, cols = 128, 16
|
||||
target = tvm.target.Target("cuda")
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
|
||||
B = T.match_buffer(B_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
|
||||
C = T.match_buffer(C_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
|
||||
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
wg_id = T.warpgroup_id([1])
|
||||
tid = T.thread_id_in_wg([rows])
|
||||
|
||||
lhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
|
||||
rhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
|
||||
out = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
|
||||
lhs_row = lhs.local(cols)
|
||||
rhs_row = rhs.local(cols)
|
||||
out_row = out.local(cols)
|
||||
for i in T.serial(cols):
|
||||
lhs_row[i] = A[tid, i]
|
||||
rhs_row[i] = B[tid, i]
|
||||
out_row[i] = T.float32(0)
|
||||
if op_name == "add":
|
||||
Tx.wg.add(out, lhs, rhs)
|
||||
elif op_name == "sub":
|
||||
Tx.wg.sub(out, lhs, rhs)
|
||||
elif op_name == "mul":
|
||||
Tx.wg.mul(out, lhs, rhs)
|
||||
out_row_1 = out.local(cols)
|
||||
for i in T.serial(cols):
|
||||
C[tid, i] = out_row_1[i]
|
||||
|
||||
with target:
|
||||
np.random.seed(0)
|
||||
A_np = np.random.rand(rows, cols).astype(dtype)
|
||||
B_np = np.random.rand(rows, cols).astype(dtype)
|
||||
C_np = np.zeros((rows, cols), dtype=dtype)
|
||||
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
if op_name == "add":
|
||||
C_ref = A_np + B_np
|
||||
elif op_name == "sub":
|
||||
C_ref = A_np - B_np
|
||||
else:
|
||||
C_ref = A_np * B_np
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
C = tvm.runtime.tensor(C_np, dev)
|
||||
mod(A, B, C)
|
||||
tvm.testing.assert_allclose(C_ref, C.numpy(), atol=1e-5)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_name,ptx_op", [("add", "add"), ("sub", "sub"), ("mul", "mul")])
|
||||
def test_binary_op_warpgroup_wg_local_emits_packed_f32x2(op_name, ptx_op):
|
||||
"""Warpgroup-scope binary on a wg-local fp32 view must lower to packed
|
||||
f32x2 PTX on SM100+, mirroring the thread-scope packed dispatch.
|
||||
|
||||
Regression test for the fa4 perf path: rescale-style ``T.{add,sub,mul}``
|
||||
calls in warpgroup scope used to fall through to scalar codegen because
|
||||
``_emit_binary_local_view`` only emitted ``op_func(...)`` per element.
|
||||
"""
|
||||
target = tvm.target.Target("cuda")
|
||||
arch = target.arch if hasattr(target, "arch") else ""
|
||||
if not arch.startswith("sm_"):
|
||||
pytest.skip(f"unknown target arch: {arch}")
|
||||
sm_digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
|
||||
if not sm_digits or int(sm_digits) < 100:
|
||||
pytest.skip(f"packed_f32x2 wg-local path requires sm_100+, got {arch}")
|
||||
|
||||
dtype = "float32"
|
||||
rows, cols = 128, 16
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
|
||||
B = T.match_buffer(B_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
|
||||
C = T.match_buffer(C_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
|
||||
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
_wg_id = T.warpgroup_id([1])
|
||||
tid = T.thread_id_in_wg([rows])
|
||||
|
||||
lhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
|
||||
rhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
|
||||
out = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
|
||||
lhs_row = lhs.local(cols)
|
||||
rhs_row = rhs.local(cols)
|
||||
out_row = out.local(cols)
|
||||
for i in T.serial(cols):
|
||||
lhs_row[i] = A[tid, i]
|
||||
rhs_row[i] = B[tid, i]
|
||||
out_row[i] = T.float32(0)
|
||||
if op_name == "add":
|
||||
Tx.wg.add(out, lhs, rhs)
|
||||
elif op_name == "sub":
|
||||
Tx.wg.sub(out, lhs, rhs)
|
||||
else:
|
||||
Tx.wg.mul(out, lhs, rhs)
|
||||
out_row_1 = out.local(cols)
|
||||
for i in T.serial(cols):
|
||||
C[tid, i] = out_row_1[i]
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
ex = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
src = ex.mod.imports[0].inspect_source()
|
||||
|
||||
# Codegen must use the packed f32x2 path, not scalar fallback.
|
||||
assert re.search(rf"{ptx_op}\.[a-z]+\.ftz\.f32x2", src) or re.search(
|
||||
rf"tvm_builtin_ptx_{ptx_op}_packed_[a-z]+_f32x2", src
|
||||
), f"expected packed f32x2 PTX for op={op_name}, source preview:\n{src[:2000]}"
|
||||
|
||||
|
||||
def test_fma_warpgroup_wg_local_emits_packed_f32x2():
|
||||
"""Same regression coverage as the binary case but for ``T.fma``."""
|
||||
target = tvm.target.Target("cuda")
|
||||
arch = target.arch if hasattr(target, "arch") else ""
|
||||
if not arch.startswith("sm_"):
|
||||
pytest.skip(f"unknown target arch: {arch}")
|
||||
sm_digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
|
||||
if not sm_digits or int(sm_digits) < 100:
|
||||
pytest.skip(f"packed_f32x2 wg-local path requires sm_100+, got {arch}")
|
||||
|
||||
dtype = "float32"
|
||||
rows, cols = 128, 16
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle, C_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
|
||||
C = T.match_buffer(C_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
|
||||
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
_wg_id = T.warpgroup_id([1])
|
||||
tid = T.thread_id_in_wg([rows])
|
||||
|
||||
buf = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
|
||||
buf_row = buf.local(cols)
|
||||
for i in T.serial(cols):
|
||||
buf_row[i] = A[tid, i]
|
||||
Tx.wg.fma(buf, buf, T.float32(2.0), T.float32(0.5))
|
||||
buf_row_1 = buf.local(cols)
|
||||
for i in T.serial(cols):
|
||||
C[tid, i] = buf_row_1[i]
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
ex = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
src = ex.mod.imports[0].inspect_source()
|
||||
|
||||
assert re.search(r"fma\.[a-z]+\.ftz\.f32x2", src) or re.search(
|
||||
r"tvm_builtin_ptx_fma_packed_[a-z]+_f32x2", src
|
||||
), f"expected packed f32x2 fma PTX, source preview:\n{src[:2000]}"
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Dispatch codegen checks (no GPU runtime — explicit target arch).
|
||||
# These complement the existing `*_warpgroup_wg_local_layout` / `*_auto_dispatch`
|
||||
# variants by forcing the arch in the Target dict, so the codegen path runs
|
||||
# even on hosts where ``Target("cuda")`` cannot detect the GPU.
|
||||
# -----------------------------------------------------------------------------
|
||||
def test_binary_add_f32_sm100_packed_f32x2_dispatch():
|
||||
"""add f32 + all-local → reg.py + add_f32x2 packed (no T.vectorized)."""
|
||||
shape = (64, 32)
|
||||
lay = TileLayout(S[shape])
|
||||
|
||||
@T.prim_func
|
||||
def k(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, "float32", layout=lay)
|
||||
B = T.match_buffer(B_ptr, shape, "float32", layout=lay)
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
tx = T.thread_id([64])
|
||||
ra = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
|
||||
rb = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
|
||||
Tx.copy(ra, A[tx])
|
||||
Tx.copy(rb, B[tx])
|
||||
Tx.add(ra, ra, rb)
|
||||
Tx.copy(A[tx], ra)
|
||||
|
||||
target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": k})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
src = mod.mod.imports[0].inspect_source()
|
||||
assert re.search(r"add\.[a-z]+\.ftz\.f32x2", src) or re.search(
|
||||
r"tvm_builtin_ptx_add_packed_", src
|
||||
), f"expected packed add_f32x2; got:\n{src[:2000]}"
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_binary_maximum_reg():
|
||||
N = 128
|
||||
|
||||
@T.prim_func
|
||||
def relu_max(a_ptr: T.handle, b_ptr: T.handle, c_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(a_ptr, (N,), "float32")
|
||||
B = T.match_buffer(b_ptr, (N,), "float32")
|
||||
C = T.match_buffer(c_ptr, (N,), "float32")
|
||||
T.device_entry()
|
||||
T.warp_id([4])
|
||||
T.cta_id([1])
|
||||
T.warpgroup_id([1])
|
||||
tid = T.thread_id_in_wg([N])
|
||||
a_reg = T.alloc_local((1,), "float32")
|
||||
b_reg = T.alloc_local((1,), "float32")
|
||||
Tx.copy(a_reg[:], A[tid : tid + 1])
|
||||
Tx.copy(b_reg[:], B[tid : tid + 1])
|
||||
Tx.maximum(a_reg[:], a_reg[:], b_reg[:])
|
||||
Tx.copy(C[tid : tid + 1], a_reg[:])
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.compile(tvm.IRModule({"main": relu_max}), target=target, tir_pipeline="tirx")
|
||||
np.random.seed(0)
|
||||
a = np.random.randn(N).astype("float32")
|
||||
b = np.random.randn(N).astype("float32")
|
||||
c = np.zeros(N, dtype="float32")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
a_t, b_t, c_t = (tvm.runtime.tensor(x, dev) for x in (a, b, c))
|
||||
mod["main"](a_t, b_t, c_t)
|
||||
np.testing.assert_allclose(c_t.numpy(), np.maximum(a, b), atol=0, rtol=0)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
def test_binary_add_f16_scalar_fallback_dispatch():
|
||||
"""add f16 has no packed VecImpl → reg.py scalar fallback (T.vectorized)."""
|
||||
shape = (64, 32)
|
||||
lay = TileLayout(S[shape])
|
||||
|
||||
@T.prim_func
|
||||
def k(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, "float16", layout=lay)
|
||||
B = T.match_buffer(B_ptr, shape, "float16", layout=lay)
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
tx = T.thread_id([64])
|
||||
ra = T.alloc_buffer(shape[1:], "float16", scope="local", layout=TileLayout(S[shape[1:]]))
|
||||
rb = T.alloc_buffer(shape[1:], "float16", scope="local", layout=TileLayout(S[shape[1:]]))
|
||||
Tx.copy(ra, A[tx])
|
||||
Tx.copy(rb, B[tx])
|
||||
Tx.add(ra, ra, rb)
|
||||
Tx.copy(A[tx], ra)
|
||||
|
||||
target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": k})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
src = mod.mod.imports[0].inspect_source()
|
||||
assert "half" in src or "__half" in src, f"expected scalar half add; got:\n{src[:2000]}"
|
||||
assert not re.search(r"add\.[a-z]+\.ftz\.f(32|16)x2", src), (
|
||||
f"unexpected packed f32x2/f16x2 add in scalar-fallback path; got:\n{src[:2000]}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,403 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
"""Tests for FMA op dispatch, layout=None local dispatch, scalar broadcast,
|
||||
and rounding mode support."""
|
||||
|
||||
import re
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx.layout import S, TileLayout, wg_local_layout
|
||||
|
||||
|
||||
def _get_sm_version():
|
||||
target = tvm.target.Target("cuda")
|
||||
arch = target.arch if hasattr(target, "arch") else ""
|
||||
if not arch.startswith("sm_"):
|
||||
return 0
|
||||
digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
|
||||
return int(digits) if digits else 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FMA op: scalar scale + scalar bias
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_fma_scalar_scalar():
|
||||
sm = _get_sm_version()
|
||||
if sm < 100:
|
||||
pytest.skip(f"packed fma requires sm_100+, got sm_{sm}")
|
||||
|
||||
N = 128
|
||||
dtype = "float32"
|
||||
target = tvm.target.Target("cuda")
|
||||
|
||||
scale_val = 0.5
|
||||
bias_val = -1.0
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
tx = T.thread_id([N])
|
||||
buf = T.alloc_buffer((1,), dtype, scope="local", layout=TileLayout(S[1]))
|
||||
Tx.copy(buf, A[tx : tx + 1])
|
||||
Tx.fma(buf, buf, T.float32(scale_val), T.float32(bias_val))
|
||||
Tx.copy(A[tx : tx + 1], buf)
|
||||
|
||||
with target:
|
||||
A_np = np.random.rand(N).astype(dtype)
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
mod(A)
|
||||
expected = A_np * scale_val + bias_val
|
||||
tvm.testing.assert_allclose(expected, A.numpy(), atol=1e-3)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FMA op: buffer scale + scalar bias (Horner pattern)
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_fma_buffer_scale_scalar_bias():
|
||||
sm = _get_sm_version()
|
||||
if sm < 100:
|
||||
pytest.skip(f"packed fma requires sm_100+, got sm_{sm}")
|
||||
|
||||
N = 2
|
||||
dtype = "float32"
|
||||
target = tvm.target.Target("cuda")
|
||||
|
||||
coeff = 0.695
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
|
||||
B = T.match_buffer(B_ptr, (N,), dtype, layout=TileLayout(S[N]))
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
_tx = T.thread_id([1])
|
||||
acc = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
|
||||
frac = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
|
||||
Tx.copy(acc, A[0:N])
|
||||
Tx.copy(frac, B[0:N])
|
||||
Tx.fma(acc, acc, frac, T.float32(coeff))
|
||||
Tx.copy(A[0:N], acc)
|
||||
|
||||
with target:
|
||||
A_np = np.random.rand(N).astype(dtype)
|
||||
B_np = np.random.rand(N).astype(dtype)
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A = tvm.runtime.tensor(A_np, dev)
|
||||
B = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A, B)
|
||||
expected = A_np * B_np + coeff
|
||||
tvm.testing.assert_allclose(expected, A.numpy(), atol=1e-3)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Binary op with scalar broadcast (Expr scalar, e.g. BufferLoad)
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_mul_scalar_broadcast():
|
||||
sm = _get_sm_version()
|
||||
if sm < 100:
|
||||
pytest.skip(f"packed mul requires sm_100+, got sm_{sm}")
|
||||
|
||||
N = 16
|
||||
dtype = "float32"
|
||||
target = tvm.target.Target("cuda")
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle, S_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
|
||||
Scale = T.match_buffer(S_ptr, (1,), dtype, layout=TileLayout(S[1]))
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
_tx = T.thread_id([1])
|
||||
a_local = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
|
||||
s_local = T.alloc_buffer((1,), dtype, scope="local", layout=TileLayout(S[1]))
|
||||
Tx.copy(a_local, A[0:N])
|
||||
Tx.copy(s_local, Scale[0:1])
|
||||
Tx.mul(a_local, a_local, s_local[0])
|
||||
Tx.copy(A[0:N], a_local)
|
||||
|
||||
with target:
|
||||
A_np = np.random.rand(N).astype(dtype)
|
||||
S_np = np.array([2.5], dtype=dtype)
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_np, dev)
|
||||
S_dev = tvm.runtime.tensor(S_np, dev)
|
||||
mod(A_dev, S_dev)
|
||||
expected = A_np * S_np[0]
|
||||
tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1e-3)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Binary add with rounding mode
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_add_rounding_mode():
|
||||
sm = _get_sm_version()
|
||||
if sm < 100:
|
||||
pytest.skip(f"packed add with rounding requires sm_100+, got sm_{sm}")
|
||||
|
||||
N = 2
|
||||
dtype = "float32"
|
||||
target = tvm.target.Target("cuda")
|
||||
|
||||
round_const = float(2**23 + 2**22)
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
_tx = T.thread_id([1])
|
||||
buf = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
|
||||
Tx.copy(buf, A[0:N])
|
||||
Tx.add(buf, buf, T.float32(round_const), rounding_mode="rm")
|
||||
Tx.copy(A[0:N], buf)
|
||||
|
||||
with target:
|
||||
A_np = np.array([1.3, 2.7], dtype=dtype)
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
# Check that the PTX uses the rounding mode
|
||||
src = mod.mod.imports[0].inspect_source()
|
||||
assert re.search(r"add\.rm\.ftz\.f32x2", src) or re.search(
|
||||
r"tvm_builtin_ptx_add_packed_", src
|
||||
), f"Expected packed add with rm rounding in PTX:\n{src}"
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_np, dev)
|
||||
mod(A_dev)
|
||||
expected = A_np + round_const
|
||||
tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1.0)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FMA op: layout=None local buffer (no TileLayout)
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_fma_no_layout():
|
||||
sm = _get_sm_version()
|
||||
if sm < 100:
|
||||
pytest.skip(f"packed fma requires sm_100+, got sm_{sm}")
|
||||
|
||||
N = 4
|
||||
dtype = "float32"
|
||||
target = tvm.target.Target("cuda")
|
||||
|
||||
scale_val = 2.0
|
||||
bias_val = 1.0
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
_tx = T.thread_id([1])
|
||||
buf = T.alloc_local([N], dtype)
|
||||
for i in T.serial(N):
|
||||
buf[i] = A[i]
|
||||
Tx.fma(buf[0:N], buf[0:N], T.float32(scale_val), T.float32(bias_val))
|
||||
for i in T.serial(N):
|
||||
A[i] = buf[i]
|
||||
|
||||
with target:
|
||||
A_np = np.array([1.0, 2.0, 3.0, 4.0], dtype=dtype)
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_np, dev)
|
||||
mod(A_dev)
|
||||
expected = A_np * scale_val + bias_val
|
||||
tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1e-3)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Binary sub with rounding mode (buffer-buffer)
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_sub_buffer_buffer_rounding():
|
||||
sm = _get_sm_version()
|
||||
if sm < 100:
|
||||
pytest.skip(f"packed sub with rounding requires sm_100+, got sm_{sm}")
|
||||
|
||||
N = 2
|
||||
dtype = "float32"
|
||||
target = tvm.target.Target("cuda")
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
|
||||
B = T.match_buffer(B_ptr, (N,), dtype, layout=TileLayout(S[N]))
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
_tx = T.thread_id([1])
|
||||
a_buf = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
|
||||
b_buf = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
|
||||
Tx.copy(a_buf, A[0:N])
|
||||
Tx.copy(b_buf, B[0:N])
|
||||
Tx.sub(a_buf, a_buf, b_buf, rounding_mode="rn")
|
||||
Tx.copy(A[0:N], a_buf)
|
||||
|
||||
with target:
|
||||
A_np = np.array([3.14, 2.71], dtype=dtype)
|
||||
B_np = np.array([1.41, 0.57], dtype=dtype)
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
src = mod.mod.imports[0].inspect_source()
|
||||
assert re.search(r"sub\.rn\.ftz\.f32x2", src) or re.search(
|
||||
r"tvm_builtin_ptx_sub_packed_", src
|
||||
), f"Expected packed sub with rn rounding in PTX:\n{src}"
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_np, dev)
|
||||
B_dev = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A_dev, B_dev)
|
||||
expected = A_np - B_np
|
||||
tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1e-6)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_fma_warpgroup_wg_local_layout():
|
||||
rows, cols = 128, 8
|
||||
dtype = "float32"
|
||||
scale_val = 1.5
|
||||
bias_val = -0.25
|
||||
target = tvm.target.Target("cuda")
|
||||
|
||||
@T.prim_func
|
||||
def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
|
||||
B = T.match_buffer(B_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
wg_id = T.warpgroup_id([1])
|
||||
tid = T.thread_id_in_wg([rows])
|
||||
|
||||
reg = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
|
||||
reg_row = reg.local(cols)
|
||||
for i in T.serial(cols):
|
||||
reg_row[i] = A[tid, i]
|
||||
Tx.wg.fma(reg, reg, T.float32(scale_val), T.float32(bias_val))
|
||||
reg_row_1 = reg.local(cols)
|
||||
for i in T.serial(cols):
|
||||
B[tid, i] = reg_row_1[i]
|
||||
|
||||
with target:
|
||||
np.random.seed(0)
|
||||
A_np = np.random.rand(rows, cols).astype(dtype)
|
||||
B_np = np.zeros((rows, cols), dtype=dtype)
|
||||
mod = tvm.IRModule({"main": test_func})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_np, dev)
|
||||
B_dev = tvm.runtime.tensor(B_np, dev)
|
||||
mod(A_dev, B_dev)
|
||||
expected = A_np * scale_val + bias_val
|
||||
tvm.testing.assert_allclose(expected, B_dev.numpy(), atol=1e-5)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Dispatch codegen check (no GPU runtime — explicit target arch).
|
||||
# Complements ``test_fma_warpgroup_wg_local_emits_packed_f32x2`` (which uses
|
||||
# the host-detected ``Target("cuda")`` and skips when arch < sm_100).
|
||||
# -----------------------------------------------------------------------------
|
||||
def test_fma_f32_sm100_packed_f32x2_dispatch():
|
||||
"""fma f32 + all-local → reg.py + fma_f32x2 packed (no T.vectorized)."""
|
||||
shape = (64, 32)
|
||||
lay = TileLayout(S[shape])
|
||||
|
||||
@T.prim_func
|
||||
def k(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, D_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, shape, "float32", layout=lay)
|
||||
B = T.match_buffer(B_ptr, shape, "float32", layout=lay)
|
||||
C = T.match_buffer(C_ptr, shape, "float32", layout=lay)
|
||||
D = T.match_buffer(D_ptr, shape, "float32", layout=lay)
|
||||
T.device_entry()
|
||||
_bx = T.cta_id([1])
|
||||
tx = T.thread_id([64])
|
||||
ra = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
|
||||
rb = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
|
||||
rc = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
|
||||
rd = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
|
||||
Tx.copy(ra, A[tx])
|
||||
Tx.copy(rb, B[tx])
|
||||
Tx.copy(rc, C[tx])
|
||||
Tx.fma(rd, ra, rb, rc)
|
||||
Tx.copy(D[tx], rd)
|
||||
|
||||
target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": k})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
src = mod.mod.imports[0].inspect_source()
|
||||
assert re.search(r"fma\.[a-z]+\.ftz\.f32x2", src) or re.search(
|
||||
r"tvm_builtin_ptx_fma_packed_", src
|
||||
), f"expected packed fma_f32x2; got:\n{src[:2000]}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,690 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
"""Tests for the CUDA synchronous ``gemm`` (mma.sync) tensor-core dispatch.
|
||||
|
||||
The dispatch lowers ``tirx.tile.gemm`` over pure-register fragments to warp-level
|
||||
``mma.sync.aligned.m16n8k16/k8`` for bf16/f16 inputs with f32 accumulation.
|
||||
|
||||
The fragment layouts below are the standard m16n8 register maps (PTX ISA
|
||||
§9.7.13; see ``tests/python/tirx-base/test_tir_ptx_mma.py``):
|
||||
|
||||
lane = 4*g + t (g = lane >> 2 in [0, 8), t = lane & 3 in [0, 4))
|
||||
D/C[M, N]: M = g + 8*rM, N = 2*t + rN, c_id = 2*rM + rN
|
||||
A[M, K]: M = g + 8*rM, K = 2*t + p + 8*kHi, ma = p + 2*rM + 4*kHi
|
||||
B[K, N]: N = g, K = 2*t + p + 8*kHi, mb = p + 2*kHi
|
||||
|
||||
Most assertions run the CPU-only ``LowerTIRx`` transform; the numerical check
|
||||
is guarded by ``requires_cuda`` since it needs a real device.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx.layout import S, TileLayout, laneid
|
||||
from tvm.tirx.operator.tile_primitive import list_registered_schedules
|
||||
|
||||
# Single-tile m16n8k8 fragment layouts -- the smallest unit everything else is
|
||||
# built from. A is 16x8, B is 8x8 as [K, N], D/C is 16x8 (the accumulator does
|
||||
# not depend on K). Every other layout (the k16 single tile, all tilings, and
|
||||
# the transposed orientations) is derived from these three via ``tile_to`` /
|
||||
# ``group`` + ``permute_by_groups``.
|
||||
D_FRAG = TileLayout(S[(2, 8, 4, 2) : (2, 4 @ laneid, 1 @ laneid, 1)])
|
||||
A_FRAG_K8 = TileLayout(S[(2, 8, 4, 2) : (2, 4 @ laneid, 1 @ laneid, 1)])
|
||||
B_FRAG_K8 = TileLayout(S[(4, 2, 8) : (1 @ laneid, 1, 4 @ laneid)])
|
||||
# m16n8k16 single tile = two k8 tiles stacked along K.
|
||||
A_FRAG = A_FRAG_K8.tile_to([16, 16], [16, 8])
|
||||
B_FRAG = B_FRAG_K8.tile_to([16, 8], [8, 8])
|
||||
|
||||
|
||||
def _transpose_frag(layout, shape):
|
||||
"""Swap the two logical axes of a 2D fragment layout.
|
||||
|
||||
The transposed input orientations (A as [K, M], B as [N, K]) hold the exact
|
||||
same per-lane/per-register element distribution as the K-major fragments --
|
||||
only the buffer's logical axes are swapped. So instead of writing them out
|
||||
by hand, derive them: ``group`` the shard into the logical dims, then
|
||||
``permute_by_groups`` to exchange the two groups.
|
||||
"""
|
||||
grouped, seps = layout.group(shape)
|
||||
return grouped.permute_by_groups(seps, [1, 0])
|
||||
|
||||
|
||||
# Transposed input orientations of the same single tile: A as [K, M], B as
|
||||
# [N, K]. The dispatch swaps axes per the transpose flags; the .row.col mma is
|
||||
# unchanged.
|
||||
A_KM_FRAG = _transpose_frag(A_FRAG, [16, 16])
|
||||
B_NK_FRAG = _transpose_frag(B_FRAG, [16, 8])
|
||||
|
||||
|
||||
def _frag(Mt, Nt, Kt, kinst):
|
||||
"""Fragment layouts for an Mt x Nt x Kt tiling of m16n8k{8,16}.
|
||||
|
||||
Logical shapes: A = (16*Mt, kinst*Kt), B = (kinst*Kt, 8*Nt), D/C = (16*Mt, 8*Nt).
|
||||
Each operand's tiled layout is the single-tile base ``tile_to`` the full
|
||||
logical shape -- ``tile_to`` repeats the base's per-lane/per-register element
|
||||
map over the tile grid, so a tiling is just a grid of the single-tile
|
||||
fragments (the base is the single source of truth, k8 and k16 alike).
|
||||
"""
|
||||
A_base = A_FRAG if kinst == 16 else A_FRAG_K8
|
||||
B_base = B_FRAG if kinst == 16 else B_FRAG_K8
|
||||
D = D_FRAG.tile_to([16 * Mt, 8 * Nt], [16, 8])
|
||||
A = A_base.tile_to([16 * Mt, kinst * Kt], [16, kinst])
|
||||
B = B_base.tile_to([kinst * Kt, 8 * Nt], [kinst, 8])
|
||||
return D, A, B
|
||||
|
||||
|
||||
def _build_tiled(Mt, Nt, Kt, kinst, *, beta=0.0, dtype="float16", store=False):
|
||||
"""A single-warp kernel issuing one ``T.gemm`` over an Mt x Nt x Kt tiling.
|
||||
|
||||
With ``store=True`` the result is written back to a global buffer (a full
|
||||
kernel for codegen); otherwise only the ``T.gemm`` is emitted (for
|
||||
``LowerTIRx`` dispatch checks).
|
||||
"""
|
||||
Dl, Al, Bl = _frag(Mt, Nt, Kt, kinst)
|
||||
M, N, K = 16 * Mt, 8 * Nt, kinst * Kt
|
||||
|
||||
if not store:
|
||||
|
||||
@T.prim_func
|
||||
def gemm():
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
_lane = T.lane_id([32])
|
||||
A = T.alloc_buffer((M, K), dtype, scope="local", layout=Al)
|
||||
B = T.alloc_buffer((K, N), dtype, scope="local", layout=Bl)
|
||||
C = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
D = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=beta)
|
||||
|
||||
return gemm
|
||||
|
||||
@T.prim_func
|
||||
def gemm(D_ptr: T.handle):
|
||||
D_g = T.match_buffer(D_ptr, (M, N), "float32")
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
lane = T.lane_id([32])
|
||||
A = T.alloc_buffer((M, K), dtype, scope="local", layout=Al)
|
||||
B = T.alloc_buffer((K, N), dtype, scope="local", layout=Bl)
|
||||
C = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
D = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=beta)
|
||||
# Decode D's per-thread registers (c = ((mt*Nt + nt)*2 + rM)*2 + rN)
|
||||
# back to logical (M, N) and store, exercising the whole tiling.
|
||||
D_reg = D.local(Mt * Nt * 4)
|
||||
for c in T.unroll(Mt * Nt * 4):
|
||||
rN = c % 2
|
||||
rM = (c // 2) % 2
|
||||
nt = (c // 4) % Nt
|
||||
mt = c // (4 * Nt)
|
||||
D_g[mt * 16 + lane // 4 + rM * 8, nt * 8 + (lane % 4) * 2 + rN] = D_reg[c]
|
||||
|
||||
return gemm
|
||||
|
||||
|
||||
def _build_gemm(alpha=1.0, beta=0.0, dtype="bfloat16"):
|
||||
"""A single-warp kernel issuing one ``T.gemm`` over register fragments."""
|
||||
|
||||
@T.prim_func
|
||||
def gemm_min():
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_tid = T.thread_id([32])
|
||||
D = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
C = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
A = T.alloc_buffer((16, 16), dtype, scope="local", layout=A_FRAG)
|
||||
B = T.alloc_buffer((16, 8), dtype, scope="local", layout=B_FRAG)
|
||||
Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=alpha, beta=beta)
|
||||
|
||||
return gemm_min
|
||||
|
||||
|
||||
def _build_transpose(transpose_A, transpose_B, *, store=False):
|
||||
"""Single m16n8k16 tile with the requested A/B input orientations."""
|
||||
Al = A_KM_FRAG if transpose_A else A_FRAG
|
||||
Bl = B_NK_FRAG if transpose_B else B_FRAG
|
||||
A_shape = (16, 16) # [K, M] or [M, K] -- both 16x16 for one tile
|
||||
B_shape = (8, 16) if transpose_B else (16, 8)
|
||||
|
||||
if not store:
|
||||
|
||||
@T.prim_func
|
||||
def gemm():
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
_lane = T.lane_id([32])
|
||||
A = T.alloc_buffer(A_shape, "float16", scope="local", layout=Al)
|
||||
B = T.alloc_buffer(B_shape, "float16", scope="local", layout=Bl)
|
||||
C = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
D = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
Tx.warp.gemm(
|
||||
D,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
transpose_A=transpose_A,
|
||||
transpose_B=transpose_B,
|
||||
alpha=1.0,
|
||||
beta=0.0,
|
||||
)
|
||||
|
||||
return gemm
|
||||
|
||||
@T.prim_func
|
||||
def gemm(D_ptr: T.handle):
|
||||
D_g = T.match_buffer(D_ptr, (16, 8), "float32")
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
lane = T.lane_id([32])
|
||||
A = T.alloc_buffer(A_shape, "float16", scope="local", layout=Al)
|
||||
B = T.alloc_buffer(B_shape, "float16", scope="local", layout=Bl)
|
||||
C = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
D = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
Tx.warp.gemm(
|
||||
D,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
transpose_A=transpose_A,
|
||||
transpose_B=transpose_B,
|
||||
alpha=1.0,
|
||||
beta=0.0,
|
||||
)
|
||||
D_reg = D.local(4)
|
||||
for c in T.unroll(4):
|
||||
D_g[lane // 4 + (c // 2) * 8, (lane % 4) * 2 + c % 2] = D_reg[c]
|
||||
|
||||
return gemm
|
||||
|
||||
|
||||
def _build_dtypes(a_dtype, b_dtype, c_dtype, d_dtype):
|
||||
"""Single tile with explicit per-operand dtypes (for decline checks)."""
|
||||
|
||||
@T.prim_func
|
||||
def gemm_min():
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_tid = T.thread_id([32])
|
||||
D = T.alloc_buffer((16, 8), d_dtype, scope="local", layout=D_FRAG)
|
||||
C = T.alloc_buffer((16, 8), c_dtype, scope="local", layout=D_FRAG)
|
||||
A = T.alloc_buffer((16, 16), a_dtype, scope="local", layout=A_FRAG)
|
||||
B = T.alloc_buffer((16, 8), b_dtype, scope="local", layout=B_FRAG)
|
||||
Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0)
|
||||
|
||||
return gemm_min
|
||||
|
||||
|
||||
def _build_tiled_numeric(Mt, Nt, Kt, kinst, beta, dtype):
|
||||
"""End-to-end ``T.gemm`` over an Mt x Nt x Kt tiling, with the A/B inputs
|
||||
loaded and the D output stored register-by-register.
|
||||
|
||||
Fragments are indexed through their per-register multi-dim ``.local()`` views
|
||||
(the shard's non-lane dims, in shard order): A = [Mt, rM(2), Kt, kHi, kp],
|
||||
B = [Kt, kHi, kp, Nt], D/C = [Mt, rM(2), Nt, rN(2)]. The lane owns g = lane>>2
|
||||
and t = lane&3; within a tile M = mt*16 + rM*8 + g, N = nt*8 + t*2 + rN,
|
||||
K = kt*kinst + kHi*8 + t*2 + kp.
|
||||
"""
|
||||
Dl, Al, Bl = _frag(Mt, Nt, Kt, kinst)
|
||||
M, N, K = 16 * Mt, 8 * Nt, kinst * Kt
|
||||
KP = 2
|
||||
kHi_n = kinst // (4 * KP)
|
||||
|
||||
@T.prim_func
|
||||
def gemm(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, D_ptr: T.handle):
|
||||
A_g = T.match_buffer(A_ptr, (M, K), dtype)
|
||||
B_g = T.match_buffer(B_ptr, (K, N), dtype)
|
||||
C_g = T.match_buffer(C_ptr, (M, N), "float32")
|
||||
D_g = T.match_buffer(D_ptr, (M, N), "float32")
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
lane = T.lane_id([32])
|
||||
A_f = T.alloc_buffer((M, K), dtype, scope="local", layout=Al)
|
||||
B_f = T.alloc_buffer((K, N), dtype, scope="local", layout=Bl)
|
||||
C_f = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
D_f = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
A_reg = A_f.local(Mt, 2, Kt, kHi_n, KP)
|
||||
for mt, rM, kt, kHi, kp in T.grid(Mt, 2, Kt, kHi_n, KP):
|
||||
A_reg[mt, rM, kt, kHi, kp] = A_g[
|
||||
mt * 16 + lane // 4 + 8 * rM,
|
||||
kt * kinst + kHi * 8 + 2 * (lane % 4) + kp,
|
||||
]
|
||||
B_reg = B_f.local(Kt, kHi_n, KP, Nt)
|
||||
for kt, kHi, kp, nt in T.grid(Kt, kHi_n, KP, Nt):
|
||||
B_reg[kt, kHi, kp, nt] = B_g[
|
||||
kt * kinst + kHi * 8 + 2 * (lane % 4) + kp,
|
||||
nt * 8 + lane // 4,
|
||||
]
|
||||
if beta == 1.0:
|
||||
C_reg = C_f.local(Mt, 2, Nt, 2)
|
||||
for mt, rM, nt, rN in T.grid(Mt, 2, Nt, 2):
|
||||
C_reg[mt, rM, nt, rN] = C_g[
|
||||
mt * 16 + lane // 4 + 8 * rM, nt * 8 + 2 * (lane % 4) + rN
|
||||
]
|
||||
Tx.warp.gemm(D_f, A_f, B_f, C_f, transpose_A=False, transpose_B=False, alpha=1.0, beta=beta)
|
||||
D_reg = D_f.local(Mt, 2, Nt, 2)
|
||||
for mt, rM, nt, rN in T.grid(Mt, 2, Nt, 2):
|
||||
D_g[mt * 16 + lane // 4 + 8 * rM, nt * 8 + 2 * (lane % 4) + rN] = D_reg[mt, rM, nt, rN]
|
||||
|
||||
return gemm, M, N, K
|
||||
|
||||
|
||||
def _build_transpose_numeric(transpose_A, transpose_B, dtype="float16"):
|
||||
"""End-to-end single-tile ``T.gemm`` for one A/B input orientation.
|
||||
|
||||
The transposed A fragment (``A_KM_FRAG``) carries its registers in the
|
||||
[kHi, kp, rM] shard order (vs [rM, kHi, kp] for the K-major ``A_FRAG``); B's
|
||||
register order ([kHi, kp]) is the same for both orientations. The buffer
|
||||
index axes swap with the orientation, but each register still holds the same
|
||||
logical (M, K) / (K, N) element.
|
||||
"""
|
||||
Al = A_KM_FRAG if transpose_A else A_FRAG
|
||||
Bl = B_NK_FRAG if transpose_B else B_FRAG
|
||||
A_shape = (16, 16)
|
||||
B_shape = (8, 16) if transpose_B else (16, 8)
|
||||
|
||||
@T.prim_func
|
||||
def gemm(A_ptr: T.handle, B_ptr: T.handle, D_ptr: T.handle):
|
||||
A_g = T.match_buffer(A_ptr, A_shape, dtype)
|
||||
B_g = T.match_buffer(B_ptr, B_shape, dtype)
|
||||
D_g = T.match_buffer(D_ptr, (16, 8), "float32")
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
lane = T.lane_id([32])
|
||||
A_f = T.alloc_buffer(A_shape, dtype, scope="local", layout=Al)
|
||||
B_f = T.alloc_buffer(B_shape, dtype, scope="local", layout=Bl)
|
||||
D_f = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
A_reg = A_f.local(2, 2, 2)
|
||||
if transpose_A:
|
||||
# A_KM_FRAG register order is [kHi, kp, rM]; buffer is [K, M].
|
||||
for kHi, kp, rM in T.grid(2, 2, 2):
|
||||
A_reg[kHi, kp, rM] = A_g[2 * (lane % 4) + kp + 8 * kHi, lane // 4 + 8 * rM]
|
||||
else:
|
||||
# A_FRAG register order is [rM, kHi, kp]; buffer is [M, K].
|
||||
for rM, kHi, kp in T.grid(2, 2, 2):
|
||||
A_reg[rM, kHi, kp] = A_g[lane // 4 + 8 * rM, 2 * (lane % 4) + kp + 8 * kHi]
|
||||
B_reg = B_f.local(2, 2)
|
||||
if transpose_B:
|
||||
# B_NK_FRAG buffer is [N, K].
|
||||
for kHi, kp in T.grid(2, 2):
|
||||
B_reg[kHi, kp] = B_g[lane // 4, 2 * (lane % 4) + kp + 8 * kHi]
|
||||
else:
|
||||
for kHi, kp in T.grid(2, 2):
|
||||
B_reg[kHi, kp] = B_g[2 * (lane % 4) + kp + 8 * kHi, lane // 4]
|
||||
Tx.warp.gemm(
|
||||
D_f,
|
||||
A_f,
|
||||
B_f,
|
||||
D_f,
|
||||
transpose_A=transpose_A,
|
||||
transpose_B=transpose_B,
|
||||
alpha=1.0,
|
||||
beta=0.0,
|
||||
)
|
||||
D_reg = D_f.local(2, 2)
|
||||
for rM, rN in T.grid(2, 2):
|
||||
D_g[lane // 4 + 8 * rM, 2 * (lane % 4) + rN] = D_reg[rM, rN]
|
||||
|
||||
return gemm
|
||||
|
||||
|
||||
def _lower(func):
|
||||
with tvm.target.Target("cuda"):
|
||||
return tvm.tirx.transform.LowerTIRx()(tvm.IRModule({"main": func}))
|
||||
|
||||
|
||||
def test_cuda_gemm_mma_variant_is_registered():
|
||||
# Importing tvm.tirx registers all per-target schedule variants. The new
|
||||
# synchronous CUDA mma path must show up for ("gemm", "cuda"). The registry
|
||||
# keys ops by their full name (``op.name`` == "tirx.tile.gemm").
|
||||
schedules = list_registered_schedules()
|
||||
cuda_gemm = schedules.get("tirx.tile.gemm", {}).get("cuda", [])
|
||||
assert "mma.m16n8k*" in cuda_gemm, (
|
||||
f"mma.m16n8k* not registered; tirx.tile.gemm schedules = {schedules.get('tirx.tile.gemm')}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16", "float16"])
|
||||
def test_cuda_gemm_mma_lowers_to_mma_sync(dtype):
|
||||
"""beta=0: the dispatch clears D, then issues a single accumulating mma with
|
||||
the registers laid out in the fixed PTX fragment order."""
|
||||
script = _lower(_build_gemm(alpha=1.0, beta=0.0, dtype=dtype))["main"].script()
|
||||
|
||||
assert "T.ptx.mma(" in script
|
||||
assert "m16n8k16" in script
|
||||
# beta == 0 clears the accumulator before the K loop.
|
||||
assert "T.float32(0" in script
|
||||
# D accumulator: c_id = 2*rM + rN -> regs 0..3.
|
||||
for r in range(4):
|
||||
assert f"d_local[{r}]" in script
|
||||
# A multiplicand: b32 = rM + 2*kHi (kHi outer) -> ma in {0, 2, 4, 6}.
|
||||
for r in (0, 2, 4, 6):
|
||||
assert f"a_local[{r}]" in script
|
||||
# B multiplicand: b32 = kHi -> mb in {0, 2}.
|
||||
for r in (0, 2):
|
||||
assert f"b_local[{r}]" in script
|
||||
|
||||
|
||||
def test_cuda_gemm_mma_accumulates_c_when_beta_one():
|
||||
"""beta=1: the accumulator is initialized by copying C instead of zeroing."""
|
||||
script = _lower(_build_gemm(alpha=1.0, beta=1.0))["main"].script()
|
||||
|
||||
assert "T.ptx.mma(" in script
|
||||
assert "m16n8k16" in script
|
||||
# The init reads C into D; nothing is zeroed.
|
||||
assert "c_local[" in script
|
||||
assert "T.float32(0" not in script
|
||||
|
||||
|
||||
def test_cuda_gemm_mma_rejects_nonunit_alpha():
|
||||
"""alpha != 1 is unsupported (ptx mma has no scale); dispatch must fail."""
|
||||
with pytest.raises(RuntimeError, match="dispatch failed"):
|
||||
_lower(_build_gemm(alpha=2.0, beta=0.0))
|
||||
|
||||
|
||||
def test_cuda_gemm_mma_rejects_fractional_beta():
|
||||
"""beta must be 0 or 1 (mma only accumulates 1*C); other values must fail."""
|
||||
with pytest.raises(RuntimeError, match="dispatch failed"):
|
||||
_lower(_build_gemm(alpha=1.0, beta=0.5))
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("dtype", ["float16", "bfloat16"])
|
||||
def test_cuda_gemm_mma_numerical(dtype):
|
||||
"""End-to-end D = A @ B on a single m16n8k16 tile (one warp).
|
||||
|
||||
A is [M, K] = [16, 16], B is [K, N] = [16, 8], D is [M, N] = [16, 8].
|
||||
|
||||
The lane-distributed register fragments cannot be filled with a whole-tile
|
||||
``T.copy`` (the per-thread axis can't be matched coordinate-wise), so each
|
||||
of a lane's registers is loaded/stored by decoding the m16n8k16 register map
|
||||
with ``g = lane >> 2`` and ``t = lane & 3``. The per-register *slot* order
|
||||
matches the dispatch's fragment register layout:
|
||||
|
||||
A reg slot = 4*rM + 2*kHi + kp -> M = g + 8*rM, K = 2*t + kp + 8*kHi
|
||||
B reg slot = 2*kHi + kp -> K = 2*t + kp + 8*kHi, N = g
|
||||
D reg slot = 2*rM + rN -> M = g + 8*rM, N = 2*t + rN
|
||||
"""
|
||||
if dtype == "bfloat16":
|
||||
ml_dtypes = pytest.importorskip("ml_dtypes")
|
||||
np_dtype = ml_dtypes.bfloat16
|
||||
else:
|
||||
np_dtype = np.float16
|
||||
|
||||
@T.prim_func
|
||||
def gemm(A_ptr: T.handle, B_ptr: T.handle, D_ptr: T.handle):
|
||||
A_g = T.match_buffer(A_ptr, (16, 16), dtype)
|
||||
B_g = T.match_buffer(B_ptr, (16, 8), dtype)
|
||||
D_g = T.match_buffer(D_ptr, (16, 8), "float32")
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
lane = T.lane_id([32])
|
||||
A_f = T.alloc_buffer((16, 16), dtype, scope="local", layout=A_FRAG)
|
||||
B_f = T.alloc_buffer((16, 8), dtype, scope="local", layout=B_FRAG)
|
||||
D_f = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
A_reg = A_f.local(8)
|
||||
for s in T.unroll(8):
|
||||
kp = s % 2
|
||||
kHi = (s // 2) % 2
|
||||
rM = s // 4
|
||||
A_reg[s] = A_g[lane // 4 + 8 * rM, 2 * (lane % 4) + kp + 8 * kHi]
|
||||
B_reg = B_f.local(4)
|
||||
for s in T.unroll(4):
|
||||
kp = s % 2
|
||||
kHi = s // 2
|
||||
B_reg[s] = B_g[2 * (lane % 4) + kp + 8 * kHi, lane // 4]
|
||||
Tx.warp.gemm(D_f, A_f, B_f, D_f, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0)
|
||||
D_reg = D_f.local(4)
|
||||
for s in T.unroll(4):
|
||||
rN = s % 2
|
||||
rM = s // 2
|
||||
D_g[lane // 4 + 8 * rM, 2 * (lane % 4) + rN] = D_reg[s]
|
||||
|
||||
with tvm.target.Target("cuda"):
|
||||
mod = tvm.compile(tvm.IRModule({"main": gemm}), target="cuda", tir_pipeline="tirx")
|
||||
|
||||
np.random.seed(0)
|
||||
A_np = np.random.uniform(-1, 1, (16, 16)).astype(np.float32)
|
||||
B_np = np.random.uniform(-1, 1, (16, 8)).astype(np.float32)
|
||||
golden = A_np @ B_np
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_np.astype(np_dtype), dev)
|
||||
B_dev = tvm.runtime.tensor(B_np.astype(np_dtype), dev)
|
||||
D_dev = tvm.runtime.tensor(np.zeros((16, 8), np.float32), dev)
|
||||
mod(A_dev, B_dev, D_dev)
|
||||
tvm.testing.assert_allclose(golden, D_dev.numpy(), atol=1e-2, rtol=1e-2)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# (Mt, Nt, Kt, kinst) tilings: single tile, each dim multi-tiled, fully tiled,
|
||||
# M = 64, and the m16n8k8 (kHi == 1) variants including a non-16-divisible K.
|
||||
_TILED_SHAPES = [
|
||||
(1, 1, 1, 16), # single m16n8k16 tile
|
||||
(2, 1, 1, 16), # two M-tiles
|
||||
(1, 2, 1, 16), # two N-tiles
|
||||
(1, 1, 2, 16), # two K-tiles (accumulated in place)
|
||||
(2, 2, 2, 16), # every dim tiled
|
||||
(4, 1, 1, 16), # M = 64
|
||||
(1, 1, 1, 8), # m16n8k8 single tile (kHi == 1)
|
||||
(1, 1, 3, 8), # K = 24 -> three k8 tiles
|
||||
(2, 2, 3, 8), # k8, every dim tiled
|
||||
]
|
||||
# (dtype, beta) input modes crossed against every shape: f16/bf16 inputs, with
|
||||
# beta = 0 (D = A @ B) and beta = 1 (D = A @ B + C, accumulating C in place).
|
||||
_TILED_MODES = [
|
||||
("float16", 0.0),
|
||||
("bfloat16", 0.0),
|
||||
("float16", 1.0),
|
||||
("bfloat16", 1.0),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("Mt, Nt, Kt, kinst", _TILED_SHAPES)
|
||||
@pytest.mark.parametrize("dtype, beta", _TILED_MODES)
|
||||
def test_cuda_gemm_mma_numerical_tiled(dtype, beta, Mt, Nt, Kt, kinst):
|
||||
"""End-to-end D = A @ B (+ C when beta==1) over an Mt x Nt x Kt tiling.
|
||||
|
||||
The two stacked ``parametrize`` decorators form the cartesian product of
|
||||
every tiling shape with every (dtype, beta) input mode, so each combination
|
||||
is an independent pytest item (pytest-xdist runs them in parallel)."""
|
||||
if dtype == "bfloat16":
|
||||
ml_dtypes = pytest.importorskip("ml_dtypes")
|
||||
np_dtype = ml_dtypes.bfloat16
|
||||
else:
|
||||
np_dtype = np.float16
|
||||
|
||||
func, M, N, K = _build_tiled_numeric(Mt, Nt, Kt, kinst, beta, dtype)
|
||||
with tvm.target.Target("cuda"):
|
||||
mod = tvm.compile(tvm.IRModule({"main": func}), target="cuda", tir_pipeline="tirx")
|
||||
|
||||
np.random.seed(0)
|
||||
A_np = np.random.uniform(-1, 1, (M, K)).astype(np.float32)
|
||||
B_np = np.random.uniform(-1, 1, (K, N)).astype(np.float32)
|
||||
C_np = np.random.uniform(-1, 1, (M, N)).astype(np.float32)
|
||||
golden = A_np @ B_np + (C_np if beta == 1.0 else 0.0)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_np.astype(np_dtype), dev)
|
||||
B_dev = tvm.runtime.tensor(B_np.astype(np_dtype), dev)
|
||||
C_dev = tvm.runtime.tensor(C_np, dev)
|
||||
D_dev = tvm.runtime.tensor(np.zeros((M, N), np.float32), dev)
|
||||
mod(A_dev, B_dev, C_dev, D_dev)
|
||||
tvm.testing.assert_allclose(golden, D_dev.numpy(), atol=2e-2, rtol=2e-2)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("dtype", ["float16", "bfloat16"])
|
||||
@pytest.mark.parametrize(
|
||||
"transpose_A, transpose_B",
|
||||
[(False, False), (True, False), (False, True), (True, True)],
|
||||
)
|
||||
def test_cuda_gemm_mma_numerical_transpose(transpose_A, transpose_B, dtype):
|
||||
"""End-to-end D = A @ B for every A/B input orientation, crossed with dtype.
|
||||
|
||||
The orientation and dtype decorators form a cartesian product, so each
|
||||
(transpose_A, transpose_B, dtype) is an independent pytest item."""
|
||||
if dtype == "bfloat16":
|
||||
ml_dtypes = pytest.importorskip("ml_dtypes")
|
||||
np_dtype = ml_dtypes.bfloat16
|
||||
else:
|
||||
np_dtype = np.float16
|
||||
|
||||
func = _build_transpose_numeric(transpose_A, transpose_B, dtype)
|
||||
with tvm.target.Target("cuda"):
|
||||
mod = tvm.compile(tvm.IRModule({"main": func}), target="cuda", tir_pipeline="tirx")
|
||||
|
||||
np.random.seed(0)
|
||||
A_log = np.random.uniform(-1, 1, (16, 16)).astype(np.float32) # logical A[M, K]
|
||||
B_log = np.random.uniform(-1, 1, (16, 8)).astype(np.float32) # logical B[K, N]
|
||||
A_buf = (A_log.T if transpose_A else A_log).astype(np_dtype)
|
||||
B_buf = (B_log.T if transpose_B else B_log).astype(np_dtype)
|
||||
golden = A_log @ B_log
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_buf, dev)
|
||||
B_dev = tvm.runtime.tensor(B_buf, dev)
|
||||
D_dev = tvm.runtime.tensor(np.zeros((16, 8), np.float32), dev)
|
||||
mod(A_dev, B_dev, D_dev)
|
||||
tvm.testing.assert_allclose(golden, D_dev.numpy(), atol=2e-2, rtol=2e-2)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"Mt, Nt, Kt, kinst",
|
||||
[
|
||||
(2, 1, 1, 16), # M = 32 (two M-tiles)
|
||||
(1, 2, 1, 16), # N = 16 (two N-tiles)
|
||||
(1, 1, 2, 16), # K = 32 (two K-tiles, accumulated in place)
|
||||
(2, 2, 2, 16), # 8 tiles
|
||||
(4, 1, 1, 16), # M = 64
|
||||
(1, 1, 1, 8), # m16n8k8 single tile (kHi == 1)
|
||||
(1, 1, 3, 8), # K = 24 -> three k8 tiles (16 does not divide 24)
|
||||
(2, 2, 3, 8), # k8, every dim tiled
|
||||
],
|
||||
)
|
||||
def test_cuda_gemm_mma_lowers_tiled(Mt, Nt, Kt, kinst):
|
||||
"""Every tiling we expect to dispatch must lower, selecting the right mma.
|
||||
|
||||
The k8 cases are the regression guard for the kHi == 1 fragment grouping
|
||||
(an extent-1 high-K register group must not be rejected as a thread axis).
|
||||
"""
|
||||
script = _lower(_build_tiled(Mt, Nt, Kt, kinst))["main"].script()
|
||||
assert "T.ptx.mma(" in script
|
||||
assert f"m16n8k{kinst}" in script
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize(
|
||||
"Mt, Nt, Kt, kinst",
|
||||
[
|
||||
(1, 1, 1, 16),
|
||||
(2, 2, 2, 16),
|
||||
(4, 1, 1, 16),
|
||||
(1, 1, 1, 8),
|
||||
(1, 1, 3, 8),
|
||||
(2, 2, 3, 8),
|
||||
],
|
||||
)
|
||||
def test_cuda_gemm_mma_codegen_issue_count(Mt, Nt, Kt, kinst):
|
||||
"""Full pipeline (UnrollLoop + CUDA codegen) emits one mma per (Mt, Nt, Kt)
|
||||
tile; K-tiles accumulate in place, so D is cleared once per output tile."""
|
||||
target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
|
||||
with target:
|
||||
mod = tvm.compile(
|
||||
tvm.IRModule({"main": _build_tiled(Mt, Nt, Kt, kinst, store=True)}),
|
||||
target=target,
|
||||
tir_pipeline="tirx",
|
||||
)
|
||||
src = mod.mod.imports[0].inspect_source()
|
||||
assert f"mma.sync.aligned.m16n8k{kinst}" in src
|
||||
# mma is emitted as one __device__ helper, invoked once per tile.
|
||||
helper = f"ptx_mma_m16n8k{kinst}_row_col"
|
||||
assert src.count(helper) - 1 == Mt * Nt * Kt
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"transpose_A, transpose_B",
|
||||
[(False, False), (True, False), (False, True), (True, True)],
|
||||
)
|
||||
def test_cuda_gemm_mma_lowers_transpose(transpose_A, transpose_B):
|
||||
"""All four A/B orientations dispatch to the same m16n8k16. transpose only
|
||||
describes the input's logical orientation; the .row.col mma is unchanged."""
|
||||
script = _lower(_build_transpose(transpose_A, transpose_B))["main"].script()
|
||||
assert "T.ptx.mma(" in script
|
||||
assert "m16n8k16" in script
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize(
|
||||
"transpose_A, transpose_B",
|
||||
[(False, False), (True, False), (False, True), (True, True)],
|
||||
)
|
||||
def test_cuda_gemm_mma_codegen_transpose(transpose_A, transpose_B):
|
||||
"""Every orientation codegens to a valid m16n8k16 kernel."""
|
||||
target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
|
||||
with target:
|
||||
mod = tvm.compile(
|
||||
tvm.IRModule({"main": _build_transpose(transpose_A, transpose_B, store=True)}),
|
||||
target=target,
|
||||
tir_pipeline="tirx",
|
||||
)
|
||||
assert "mma.sync.aligned.m16n8k16" in mod.mod.imports[0].inspect_source()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"a, b, c, d",
|
||||
[
|
||||
("float16", "float16", "float16", "float16"), # f16 accumulate
|
||||
("bfloat16", "float16", "float32", "float32"), # mixed A/B inputs
|
||||
("float32", "float32", "float32", "float32"), # f32 (tf32) inputs
|
||||
("int8", "int8", "int32", "int32"), # integer
|
||||
],
|
||||
)
|
||||
def test_cuda_gemm_mma_rejects_unsupported_dtype(a, b, c, d):
|
||||
"""The table holds only (bf16|f16, same, f32, f32); any other dtype
|
||||
signature must decline rather than emit a wrong mma."""
|
||||
with pytest.raises(RuntimeError, match="dispatch failed"):
|
||||
_lower(_build_dtypes(a, b, c, d))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,466 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=missing-function-docstring
|
||||
|
||||
"""Tests for ``T.permute_layout``.
|
||||
|
||||
Coverage:
|
||||
|
||||
- The algorithm helpers (`_bank_free`, `_check_bijection`, `_choose_xor_k`)
|
||||
directly, with a NumPy oracle.
|
||||
- End-to-end compiled-kernel byte-for-byte equivalence on CUDA for the SF
|
||||
fp8-blockwise-gemm transpose shapes (BLK_SFA = 128, 256) plus a few
|
||||
generic linear↔stride-permuted layouts and additional dtypes (u8, fp16,
|
||||
i32, u64).
|
||||
- Reject cases: non-warp scope, dtype mismatch, shape mismatch, swizzle/
|
||||
compose layouts, layouts whose strides don't form a bijection on the
|
||||
slice. Each must surface as a ``RuntimeError`` from the dispatcher and
|
||||
NOT silently emit a wrong kernel.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
|
||||
# Helpers exposed by the dispatcher module for direct algorithm tests.
|
||||
from tvm.tirx.cuda.operator.tile_primitive.permute_layout.warp_xor_swizzle import (
|
||||
_bank_free,
|
||||
_check_bijection,
|
||||
_choose_xor_k,
|
||||
)
|
||||
from tvm.tirx.layout import S, SwizzleLayout, TileLayout
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Algorithm-only tests (no CUDA needed).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _np_layout_offset(extent, strides, multi_idx):
|
||||
return int(sum(s * i for s, i in zip(strides, multi_idx)))
|
||||
|
||||
|
||||
def _expected_permute(src_np, src_strides, dst_strides, extent):
|
||||
"""Compute the expected output: dst at byte offset ``L_dst(i)`` holds the
|
||||
value at ``src`` byte offset ``L_src(i)``, for every logical index i.
|
||||
"""
|
||||
V = math.prod(extent)
|
||||
dst_np = np.zeros_like(src_np)
|
||||
for flat in range(V):
|
||||
idx = []
|
||||
rem = flat
|
||||
for e in reversed(extent):
|
||||
idx.append(rem % e)
|
||||
rem //= e
|
||||
idx = list(reversed(idx))
|
||||
src_off = _np_layout_offset(extent, src_strides, idx)
|
||||
dst_off = _np_layout_offset(extent, dst_strides, idx)
|
||||
dst_np.reshape(-1)[dst_off] = src_np.reshape(-1)[src_off]
|
||||
return dst_np
|
||||
|
||||
|
||||
def test_bank_free_sf_128_u32():
|
||||
"""SF BLK_SFA=128: write phase has 4-way conflict at k=0, free at k=2."""
|
||||
extent = [4, 32]
|
||||
src = [32, 1]
|
||||
dst = [1, 4]
|
||||
bytes_per = 4
|
||||
P = 4
|
||||
assert _bank_free(extent, src, bytes_per, P, 0)
|
||||
assert not _bank_free(extent, dst, bytes_per, P, 0)
|
||||
assert _bank_free(extent, dst, bytes_per, P, 2)
|
||||
assert _choose_xor_k(extent, src, dst, bytes_per, P) == 2
|
||||
|
||||
|
||||
def test_bank_free_sf_256_u32():
|
||||
"""SF BLK_SFA=256: same shift=3 (k=2) handles the high block too."""
|
||||
extent = [2, 4, 32]
|
||||
src = [128, 32, 1]
|
||||
dst = [128, 1, 4]
|
||||
bytes_per = 4
|
||||
P = 8
|
||||
assert _bank_free(extent, src, bytes_per, P, 0)
|
||||
assert not _bank_free(extent, dst, bytes_per, P, 0)
|
||||
assert _bank_free(extent, dst, bytes_per, P, 2)
|
||||
assert _choose_xor_k(extent, src, dst, bytes_per, P) == 2
|
||||
|
||||
|
||||
def test_identity_no_xor():
|
||||
"""L_src == L_dst => k=0 (no XOR needed and the op is essentially a copy)."""
|
||||
assert _choose_xor_k([4, 32], [32, 1], [32, 1], 4, 4) == 0
|
||||
# A 2D buffer with row-major to row-major is a true no-op.
|
||||
assert _bank_free([4, 32], [32, 1], 4, 4, 0)
|
||||
|
||||
|
||||
def test_bijection_check_rejects_aliased():
|
||||
"""If two logical indices map to the same physical byte, reject."""
|
||||
# Stride 0 on a non-singleton extent => alias.
|
||||
assert not _check_bijection([4, 32], [0, 1])
|
||||
# Negative or non-contiguous-but-bijective is still fine.
|
||||
assert _check_bijection([4, 32], [1, 4])
|
||||
|
||||
|
||||
def test_dtype_widths_choose_xor_k():
|
||||
"""Each dtype's outcome:
|
||||
|
||||
The unvectorized algorithm is provably correct only when every per-lane
|
||||
access maps to a single 4-byte bank. For 4-byte dtypes that always holds
|
||||
(one element per bank), so we expect a valid k. For sub-4-byte dtypes
|
||||
with stride-1 reads, multiple lanes share a bank no matter how we permute
|
||||
register slots — the dispatcher correctly rejects those (k is None).
|
||||
"""
|
||||
extent = [4, 32]
|
||||
src = [32, 1] # linear
|
||||
dst = [1, 4] # transposed
|
||||
# u32: this is the SF case; the algorithm must find k=2.
|
||||
assert _choose_xor_k(extent, src, dst, 4, 4) == 2
|
||||
# u16/fp16, u8: stride-1 in bytes < 4 packs >1 lane into the same bank;
|
||||
# register-slot XOR cannot fix that, so the dispatcher rejects.
|
||||
assert _choose_xor_k(extent, src, dst, 2, 4) is None
|
||||
assert _choose_xor_k(extent, src, dst, 1, 4) is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end compiled-kernel tests on CUDA.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _has_cuda():
|
||||
try:
|
||||
return tvm.cuda(0).exist
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
needs_cuda = pytest.mark.skipif(not _has_cuda(), reason="needs CUDA")
|
||||
|
||||
|
||||
def _compile_and_run(prim_func, np_inputs):
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": prim_func})
|
||||
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
tensors = [tvm.runtime.tensor(a, dev) for a in np_inputs]
|
||||
mod(*tensors)
|
||||
return [tensor.numpy() for tensor in tensors]
|
||||
|
||||
outputs = tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
return outputs, mod.mod.imports[0].inspect_source()
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@needs_cuda
|
||||
@pytest.mark.parametrize(
|
||||
"name, pipe, blk, dtype",
|
||||
[
|
||||
("sf_128_u32", 2, 128, "uint32"),
|
||||
("sf_256_u32", 2, 256, "uint32"),
|
||||
("sf_128_i32", 2, 128, "int32"),
|
||||
("sf_128_fp32", 2, 128, "float32"),
|
||||
],
|
||||
)
|
||||
def test_sf_blockwise_transpose(name, pipe, blk, dtype):
|
||||
"""SF blockwise-GEMM scale-factor transpose, the canonical use case."""
|
||||
high = blk // 128 if blk >= 128 else 1
|
||||
# Use 4D logical shape (PIPE, high, 4, 32) to keep the high-block factored.
|
||||
shape = (pipe, high, 4, 32)
|
||||
|
||||
# Element strides for src (linear) and dst (transposed within each
|
||||
# 128-block). Stage stride = blk; each 128-block contributes 128 to the
|
||||
# high stride.
|
||||
src_strides = (blk, 128, 32, 1)
|
||||
dst_strides = (blk, 128, 1, 4)
|
||||
pre = TileLayout(S[shape:src_strides])
|
||||
post = TileLayout(S[shape:dst_strides])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def f(A: T.handle, B: T.handle):
|
||||
A_buf = T.match_buffer(A, shape, dtype, layout=pre)
|
||||
B_buf = T.match_buffer(B, shape, dtype, layout=post)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.thread_id([32])
|
||||
for s in T.serial(0, pipe):
|
||||
Tx.warp.permute_layout(
|
||||
B_buf[s, 0:high, 0:4, 0:32], A_buf[s, 0:high, 0:4, 0:32]
|
||||
)
|
||||
# fmt: on
|
||||
|
||||
np.random.seed(0)
|
||||
A_np = tvm.testing.generate_random_array(dtype, shape)
|
||||
B_np = np.zeros_like(A_np)
|
||||
|
||||
[_, B_out], src = _compile_and_run(f, [A_np, B_np])
|
||||
|
||||
# The dispatcher must have picked the XOR-swizzled variant; check that
|
||||
# the generated CUDA contains the per-lane XOR pattern. This is the
|
||||
# "no perf regression" smoke test: any future variant that omits the
|
||||
# XOR would re-introduce 4-way bank conflicts.
|
||||
assert ">> 3" in src, f"expected XOR-swizzle (lane>>3) in CUDA for {name}"
|
||||
assert "warp_sync" in src or "syncwarp" in src
|
||||
|
||||
# Byte-for-byte equality via numpy reference.
|
||||
for s in range(pipe):
|
||||
A_flat = A_np[s].reshape(-1)
|
||||
B_flat = B_out[s].reshape(-1)
|
||||
ref = _expected_permute(
|
||||
A_flat,
|
||||
list(src_strides[1:]),
|
||||
list(dst_strides[1:]),
|
||||
list(shape[1:]),
|
||||
)
|
||||
np.testing.assert_array_equal(B_flat, ref, err_msg=f"{name} stage {s}")
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@needs_cuda
|
||||
def test_identity_passes_through_as_copy():
|
||||
"""L_src == L_dst should still compile and produce a correct (identity) copy."""
|
||||
shape = (4, 32)
|
||||
layout = TileLayout(S[shape : (32, 1)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def f(A: T.handle, B: T.handle):
|
||||
A_buf = T.match_buffer(A, shape, "uint32", layout=layout)
|
||||
B_buf = T.match_buffer(B, shape, "uint32", layout=layout)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.thread_id([32])
|
||||
Tx.warp.permute_layout(B_buf, A_buf)
|
||||
# fmt: on
|
||||
|
||||
np.random.seed(0)
|
||||
A_np = tvm.testing.generate_random_array("uint32", shape)
|
||||
B_np = np.zeros_like(A_np)
|
||||
|
||||
[_, B_out], _ = _compile_and_run(f, [A_np, B_np])
|
||||
np.testing.assert_array_equal(B_out, A_np)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@needs_cuda
|
||||
@pytest.mark.parametrize("dtype", ["uint32", "int32", "float32"])
|
||||
@pytest.mark.parametrize(
|
||||
"shape, src_strides, dst_strides",
|
||||
[
|
||||
# (8, 32) → (8, 32) transposed: src linear, dst column-major.
|
||||
((8, 32), (32, 1), (1, 8)),
|
||||
# (16, 32): per_thread = 16 — tests P=16 path.
|
||||
((16, 32), (32, 1), (1, 16)),
|
||||
],
|
||||
)
|
||||
def test_generic_transpose(shape, src_strides, dst_strides, dtype):
|
||||
"""Generic linear↔transposed pairs at various P values."""
|
||||
pre = TileLayout(S[shape:src_strides])
|
||||
post = TileLayout(S[shape:dst_strides])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def f(A: T.handle, B: T.handle):
|
||||
A_buf = T.match_buffer(A, shape, dtype, layout=pre)
|
||||
B_buf = T.match_buffer(B, shape, dtype, layout=post)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.thread_id([32])
|
||||
Tx.warp.permute_layout(B_buf, A_buf)
|
||||
# fmt: on
|
||||
|
||||
np.random.seed(0)
|
||||
A_np = tvm.testing.generate_random_array(dtype, shape)
|
||||
B_np = np.zeros_like(A_np)
|
||||
[_, B_out], _ = _compile_and_run(f, [A_np, B_np])
|
||||
|
||||
ref = _expected_permute(A_np.reshape(-1), list(src_strides), list(dst_strides), list(shape))
|
||||
np.testing.assert_array_equal(B_out.reshape(-1), ref)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reject cases: the dispatcher must surface a clear error, never silently
|
||||
# emit a wrong kernel.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_and_assert_rejected(shape, src_layout, dst_layout, dtype, msg_substr):
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def f(A: T.handle, B: T.handle):
|
||||
A_buf = T.match_buffer(A, shape, dtype, layout=src_layout)
|
||||
B_buf = T.match_buffer(B, shape, dtype, layout=dst_layout)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.thread_id([32])
|
||||
Tx.warp.permute_layout(B_buf, A_buf)
|
||||
# fmt: on
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target, pytest.raises(RuntimeError) as exc_info:
|
||||
mod = tvm.IRModule({"main": f})
|
||||
tvm.compile(mod, target=target, tir_pipeline="tirx")
|
||||
assert msg_substr in str(exc_info.value), (
|
||||
f"expected reject reason to mention {msg_substr!r}, got: {exc_info.value}"
|
||||
)
|
||||
|
||||
|
||||
def test_reject_dtype_mismatch():
|
||||
shape = (4, 32)
|
||||
layout = TileLayout(S[shape : (32, 1)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def f(A: T.handle, B: T.handle):
|
||||
A_buf = T.match_buffer(A, shape, "uint32", layout=layout)
|
||||
B_buf = T.match_buffer(B, shape, "uint16", layout=layout)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.thread_id([32])
|
||||
Tx.warp.permute_layout(B_buf, A_buf)
|
||||
# fmt: on
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target, pytest.raises(RuntimeError) as exc_info:
|
||||
tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx")
|
||||
assert "dtype mismatch" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_reject_shape_mismatch():
|
||||
src_layout = TileLayout(S[(4, 32) : (32, 1)])
|
||||
dst_layout = TileLayout(S[(8, 16) : (16, 1)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def f(A: T.handle, B: T.handle):
|
||||
A_buf = T.match_buffer(A, (4, 32), "uint32", layout=src_layout)
|
||||
B_buf = T.match_buffer(B, (8, 16), "uint32", layout=dst_layout)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.thread_id([32])
|
||||
Tx.warp.permute_layout(B_buf, A_buf)
|
||||
# fmt: on
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target, pytest.raises(RuntimeError) as exc_info:
|
||||
tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx")
|
||||
assert "shape mismatch" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_reject_swizzle_layout():
|
||||
"""ComposeLayout(SwizzleLayout, TileLayout) is not supported by the warp variant."""
|
||||
from tvm.tirx.layout import ComposeLayout
|
||||
|
||||
inner = TileLayout(S[(4, 32) : (32, 1)])
|
||||
sw = SwizzleLayout(per_element=2, swizzle_len=2, atom_len=4)
|
||||
swizzled = ComposeLayout(sw, inner)
|
||||
plain = TileLayout(S[(4, 32) : (1, 4)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def f(A: T.handle, B: T.handle):
|
||||
A_buf = T.match_buffer(A, (4, 32), "uint32", layout=swizzled)
|
||||
B_buf = T.match_buffer(B, (4, 32), "uint32", layout=plain)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.thread_id([32])
|
||||
Tx.warp.permute_layout(B_buf, A_buf)
|
||||
# fmt: on
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target, pytest.raises(RuntimeError) as exc_info:
|
||||
tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx")
|
||||
assert "TileLayout" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_reject_non_warp_scope():
|
||||
layout_pre = TileLayout(S[(4, 32) : (32, 1)])
|
||||
layout_post = TileLayout(S[(4, 32) : (1, 4)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def f(A: T.handle, B: T.handle):
|
||||
A_buf = T.match_buffer(A, (4, 32), "uint32", layout=layout_pre)
|
||||
B_buf = T.match_buffer(B, (4, 32), "uint32", layout=layout_post)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
T.thread_id([32])
|
||||
Tx.cta.permute_layout(B_buf, A_buf) # cta scope, not warp
|
||||
# fmt: on
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target, pytest.raises(RuntimeError) as exc_info:
|
||||
tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx")
|
||||
assert "warp" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", ["uint32", "float32"])
|
||||
def test_shared_to_shared_uses_direct_ldst(dtype):
|
||||
"""Compile-only: a shared->shared 32b transpose must take the direct
|
||||
base-ptr + byte-offset ``ld.shared`` / ``st.shared`` path.
|
||||
|
||||
For a 4/8-byte dtype with both operands in shared memory, indexing through
|
||||
``buf[...]`` lowers the swizzled layout to a per-element IMAD flatten. The
|
||||
direct path computes one base ptr (``ptr_to(stride_offset)``) and adds a
|
||||
compile-time ``off * dtype_bytes`` per register slot, then issues
|
||||
``T.ptx.ld/st(..., space="shared")``. The bits move through a uint
|
||||
container, so ``float32`` (whose ``ld.b32`` cannot return a float) lowers
|
||||
the same way as the ``uint32`` SF case.
|
||||
"""
|
||||
shape = (4, 32)
|
||||
pre = TileLayout(S[shape : (32, 1)]) # linear
|
||||
post = TileLayout(S[shape : (1, 4)]) # transposed within the 128-block
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def f(A: T.handle, B: T.handle):
|
||||
A_buf = T.match_buffer(A, shape, dtype, layout=pre)
|
||||
B_buf = T.match_buffer(B, shape, dtype, layout=post)
|
||||
T.device_entry()
|
||||
T.cta_id([1])
|
||||
tid = T.thread_id([32])
|
||||
sA = T.alloc_buffer(shape, dtype, scope="shared", layout=pre)
|
||||
sB = T.alloc_buffer(shape, dtype, scope="shared", layout=post)
|
||||
Tx.cta.copy(sA[:, :], A_buf[:, :])
|
||||
T.cuda.cta_sync()
|
||||
Tx.warp.permute_layout(sB[:, :], sA[:, :])
|
||||
T.cuda.cta_sync()
|
||||
Tx.cta.copy(B_buf[:, :], sB[:, :])
|
||||
# fmt: on
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
with target:
|
||||
mod = tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx")
|
||||
src = mod.mod.imports[0].inspect_source()
|
||||
assert "ld.shared" in src, f"expected direct ld.shared in permute; src=\n{src}"
|
||||
assert "st.shared" in src, f"expected direct st.shared in permute; src=\n{src}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,158 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _import_and_register():
|
||||
# Ensure all schedule registrations (legacy + dispatcher variants) are loaded
|
||||
import tvm.tirx.operator.tile_primitive as _ # noqa: F401
|
||||
|
||||
|
||||
class _DummyKind:
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
def __str__(self) -> str: # used in messages
|
||||
return self.name
|
||||
|
||||
|
||||
class _DummyTarget:
|
||||
def __init__(self, kind_name: str):
|
||||
self.kind = _DummyKind(kind_name)
|
||||
|
||||
|
||||
class _DummyExecScope:
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
|
||||
class _DummySctx:
|
||||
def __init__(self, target_kind: str, exec_scope: str):
|
||||
self.target = _DummyTarget(target_kind)
|
||||
self.exec_scope = _DummyExecScope(exec_scope)
|
||||
self.scope_kind = exec_scope
|
||||
|
||||
|
||||
def test_dispatch_prints_predicate_reasons():
|
||||
"""Validate TRACE mode prints per-variant predicate failure reasons."""
|
||||
_import_and_register()
|
||||
from tvm.ir import Op
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import run_dispatch
|
||||
|
||||
class _OpCall:
|
||||
def __init__(self, op):
|
||||
self.op = op
|
||||
self.args = [] # not used by the tested predicates
|
||||
|
||||
# Use TRN copy; predicate requires exec_scope == "thread".
|
||||
op_call = _OpCall(Op.get("tirx.tile.copy"))
|
||||
sctx = _DummySctx(target_kind="trn", exec_scope="warp") # intentionally wrong
|
||||
|
||||
with pytest.raises(RuntimeError) as e:
|
||||
run_dispatch(op_call, sctx)
|
||||
|
||||
out = str(e.value)
|
||||
print(out)
|
||||
# Header + per-variant reason must be printed in table format
|
||||
assert "TIRx schedule dispatch failed: op=tirx.tile.copy target=trn" in out
|
||||
assert "Variant" in out # table header present
|
||||
assert "default" in out # variant name present
|
||||
assert "rejected: exec_scope" in out
|
||||
# opcall object IR should be printed in the table
|
||||
assert "opcall:" in out
|
||||
|
||||
|
||||
def test_dispatch_forced_variant_missing_table_and_message():
|
||||
_import_and_register()
|
||||
from tvm.ir import Op
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import run_dispatch
|
||||
|
||||
class _OpCall:
|
||||
def __init__(self, op):
|
||||
self.op = op
|
||||
self.dispatch = "__nonexistent__"
|
||||
self.args = []
|
||||
|
||||
op_call = _OpCall(Op.get("tirx.tile.copy"))
|
||||
sctx = _DummySctx(target_kind="trn", exec_scope="thread")
|
||||
|
||||
with pytest.raises(RuntimeError) as e:
|
||||
run_dispatch(op_call, sctx)
|
||||
|
||||
msg = str(e.value)
|
||||
print(msg)
|
||||
assert "TIRx schedule dispatch failed: op=tirx.tile.copy target=trn" in msg
|
||||
assert "no variant named '__nonexistent__' is registered" in msg
|
||||
|
||||
|
||||
def test_dispatch_raises_with_aggregated_reasons():
|
||||
"""Validate STRICT mode raises aggregated error message with reasons."""
|
||||
_import_and_register()
|
||||
from tvm.ir import Op
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import run_dispatch
|
||||
|
||||
class _OpCall:
|
||||
def __init__(self, op):
|
||||
self.op = op
|
||||
self.args = []
|
||||
|
||||
# Use TRN compose_op; variant implementation raises NotImplementedError
|
||||
op_call = _OpCall(Op.get("tirx.tile.compose_op"))
|
||||
sctx = _DummySctx(target_kind="trn", exec_scope="thread")
|
||||
|
||||
with pytest.raises(RuntimeError) as e:
|
||||
run_dispatch(op_call, sctx)
|
||||
|
||||
msg = str(e.value)
|
||||
print(msg)
|
||||
assert "TIRx schedule dispatch failed: op=tirx.tile.compose_op target=trn" in msg
|
||||
assert "default" in msg
|
||||
assert "exception — NotImplementedError" in msg
|
||||
# opcall content and backtrace should be included inside the table
|
||||
assert "opcall:" in msg
|
||||
assert "Traceback (most recent call last):" in msg
|
||||
|
||||
|
||||
def test_dispatch_prints_real_opcall_ir():
|
||||
"""Create a real TilePrimitiveCall via BufferRegions and ensure its IR is in the table."""
|
||||
_import_and_register()
|
||||
from tvm.ir import Op
|
||||
from tvm.tirx.buffer import decl_buffer
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import run_dispatch
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
# Build a real TIRx TilePrimitiveCall: tirx.tile.copy(A[0:64], B[0:64])
|
||||
A = decl_buffer((64,), "float32", scope="global")
|
||||
B = decl_buffer((64,), "float32", scope="shared")
|
||||
real_opcall = TilePrimitiveCall(
|
||||
A[0:64], B[0:64], op=Op.get("tirx.tile.copy"), workspace={}, config={}
|
||||
)
|
||||
|
||||
# Force predicate rejection to trigger formatted error with opcall IR
|
||||
sctx = _DummySctx(target_kind="trn", exec_scope="warp")
|
||||
with pytest.raises(RuntimeError) as e:
|
||||
run_dispatch(real_opcall, sctx)
|
||||
|
||||
out = str(e.value)
|
||||
print(out)
|
||||
# Verify header and that the opcall IR is included in the table
|
||||
assert "TIRx schedule dispatch failed: op=tirx.tile.copy target=trn" in out
|
||||
assert "Variant" in out
|
||||
assert "opcall:" in out
|
||||
# IR should mention the operator name
|
||||
assert "tirx.tile.copy" in out
|
||||
@@ -0,0 +1,354 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir import assert_structural_equal as _assert_structural_equal
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.tirx.layout import F, P, S, TileLayout
|
||||
from tvm.tirx.stmt_functor import ir_transform
|
||||
|
||||
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
|
||||
|
||||
|
||||
def _strip_exec_scope_stmt(stmt):
|
||||
def _postorder(node):
|
||||
if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
|
||||
return node.body
|
||||
return node
|
||||
|
||||
return ir_transform(
|
||||
stmt,
|
||||
preorder=lambda _node: None,
|
||||
postorder=_postorder,
|
||||
only_enable=["tirx.AttrStmt"],
|
||||
)
|
||||
|
||||
|
||||
def assert_structural_equal(lhs, rhs, *args, **kwargs):
|
||||
if isinstance(lhs, tvm.tirx.PrimFunc):
|
||||
lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
|
||||
if isinstance(rhs, tvm.tirx.PrimFunc):
|
||||
rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
|
||||
_assert_structural_equal(lhs, rhs, *args, **kwargs)
|
||||
|
||||
|
||||
Tx_func_map = {"add": Tx.add, "sub": Tx.sub, "mul": Tx.mul, "min": Tx.minimum, "max": Tx.maximum}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "min", "max"])
|
||||
@pytest.mark.parametrize(
|
||||
"operands_type",
|
||||
[
|
||||
"region_region",
|
||||
"const_region",
|
||||
"region_const",
|
||||
"region_broadcast_lhs",
|
||||
"region_broadcast_rhs",
|
||||
],
|
||||
)
|
||||
def test_simple_binary(op_type, operands_type):
|
||||
const = T.float32(3.0)
|
||||
src1_shape = [128, 512] if operands_type != "region_broadcast_lhs" else [128, 1]
|
||||
src1_layout = TileLayout(S[src1_shape : (1 @ P, 1 @ F)])
|
||||
src2_shape = [128, 512] if operands_type != "region_broadcast_rhs" else [128, 1]
|
||||
src2_layout = TileLayout(S[src2_shape : (1 @ P, 1 @ F)])
|
||||
dst_shape = [128, 512]
|
||||
dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
Tx_func = Tx_func_map[op_type]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary() ->None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
if operands_type == "region_region" or operands_type.startswith("region_broadcast"):
|
||||
Tx_func(C_sbuf, A_sbuf, B_sbuf)
|
||||
elif operands_type == "const_region":
|
||||
Tx_func(C_sbuf, const, A_sbuf)
|
||||
elif operands_type == "region_const":
|
||||
Tx_func(C_sbuf, A_sbuf, const)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "binary"})
|
||||
A_sbuf = T.alloc_buffer(src1_shape, scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer(src2_shape, scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer(dst_shape, scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
if operands_type == "region_region":
|
||||
T.nki.tensortensor(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], B_sbuf[p_loop, f_loop], op_type) # noqa: E501
|
||||
elif operands_type == "region_const":
|
||||
T.nki.tensorscalar(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], T.float32(3.0), op_type, T.bool(False)) # noqa: E501
|
||||
elif operands_type == "const_region":
|
||||
T.nki.tensorscalar(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], T.float32(3.0), op_type, T.bool(True)) # noqa: E501
|
||||
elif operands_type == "region_broadcast_rhs":
|
||||
T.nki.tensorscalar(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], B_sbuf[p_loop, 0], op_type, T.bool(False)) # noqa: E501
|
||||
elif operands_type == "region_broadcast_lhs":
|
||||
T.nki.tensorscalar(C_sbuf[p_loop, f_loop], B_sbuf[p_loop, f_loop], A_sbuf[p_loop, 0], op_type, T.bool(True)) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": binary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "min", "max"])
|
||||
@pytest.mark.parametrize(
|
||||
"operands_type",
|
||||
[
|
||||
"region_region",
|
||||
"const_region",
|
||||
"region_const",
|
||||
"region_broadcast_lhs",
|
||||
"region_broadcast_rhs",
|
||||
],
|
||||
)
|
||||
def test_binary_complex(op_type, operands_type):
|
||||
src1_shape = [1024, 512] if operands_type != "region_broadcast_lhs" else [1024, 4]
|
||||
src1_layout_data_iter = (128, 4096) if operands_type != "region_broadcast_lhs" else (128, 32)
|
||||
src1_layout = TileLayout(S[src1_layout_data_iter : (1 @ P, 1 @ F)])
|
||||
src2_shape = [512, 512] if operands_type != "region_broadcast_rhs" else [128, 512]
|
||||
src2_layout_data_iter = (128, 2048) if operands_type != "region_broadcast_rhs" else (128, 512)
|
||||
src2_layout = TileLayout(S[src2_layout_data_iter : (1 @ P, 1 @ F)])
|
||||
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
|
||||
const = T.float32(3.0)
|
||||
Tx_func = Tx_func_map[op_type]
|
||||
|
||||
src1_view_shape = [128, 8, 512]
|
||||
src2_view_shape = [128, 4, 512] if operands_type != "region_broadcast_rhs" else [128, 1, 512]
|
||||
dst_view_shape = [128, 4, 512]
|
||||
if operands_type == "region_broadcast_lhs":
|
||||
src1_view_shape = [128, 8, 4, 1]
|
||||
src2_view_shape = [128, 4, 4, 128]
|
||||
dst_view_shape = [128, 4, 4, 128]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
A_sbuf_view = A_sbuf.view(*src1_view_shape)
|
||||
B_sbuf_view = B_sbuf.view(*src2_view_shape)
|
||||
C_sbuf_view = C_sbuf.view(*dst_view_shape)
|
||||
for i in range(4):
|
||||
if operands_type == "region_region":
|
||||
Tx_func(C_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :], B_sbuf_view[:, i, :])
|
||||
elif operands_type == "region_const":
|
||||
Tx_func(C_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :], const)
|
||||
elif operands_type == "const_region":
|
||||
Tx_func(C_sbuf_view[:, i, :], const, A_sbuf_view[:, i * 2, :])
|
||||
elif operands_type == "region_broadcast_rhs":
|
||||
Tx_func(C_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :], B_sbuf_view[:, 0, :])
|
||||
elif operands_type == "region_broadcast_lhs":
|
||||
Tx_func(C_sbuf_view[:, i, :, :], A_sbuf_view[:, i*2,:, :], B_sbuf_view[:, i, :, :])
|
||||
|
||||
f_extent = 128 if operands_type == "region_broadcast_lhs" else 512
|
||||
b_extent = 4 if operands_type == "region_broadcast_lhs" else 1
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "binary"})
|
||||
A_sbuf = T.alloc_buffer(src1_layout_data_iter, scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer(src2_layout_data_iter, scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
A_sbuf_view = T.decl_buffer(src1_layout_data_iter, data=A_sbuf.data, scope="trn.sbuf", layout=None) # noqa: E501
|
||||
B_sbuf_view = T.decl_buffer(src2_layout_data_iter, data=B_sbuf.data, scope="trn.sbuf", layout=None) # noqa: E501
|
||||
C_sbuf_view = T.decl_buffer((128, 2048), data=C_sbuf.data, scope="trn.sbuf", layout=None)
|
||||
for i, b_loop in T.grid(4, b_extent):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, f_extent, annotations={"nki_dim":"F"}):
|
||||
if operands_type == "region_region":
|
||||
T.nki.tensortensor(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], B_sbuf_view[p_loop, i * 512 + f_loop], op_type) # noqa: E501
|
||||
elif operands_type == "const_region":
|
||||
T.nki.tensorscalar(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], T.float32(3.0), op_type, T.bool(True)) # noqa: E501
|
||||
elif operands_type == "region_const":
|
||||
T.nki.tensorscalar(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], T.float32(3.0), op_type, T.bool(False)) # noqa: E501
|
||||
elif operands_type == "region_broadcast_lhs":
|
||||
T.nki.tensorscalar(C_sbuf_view[p_loop, i * 512 + b_loop * 128 + f_loop], B_sbuf_view[p_loop, i * 512 + b_loop * 128 + f_loop], A_sbuf_view[p_loop, i * 8 + b_loop], op_type, T.bool(True)) # noqa: E501
|
||||
elif operands_type == "region_broadcast_rhs":
|
||||
T.nki.tensortensor(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], B_sbuf_view[p_loop, f_loop], op_type) # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": binary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_binary_broadcast1():
|
||||
src1_shape = [32, 128, 512]
|
||||
src1_layout = TileLayout(S[(32, 128, 4, 128) : (1 @ F, 32 @ F, 32 * 128 @ F, 1 @ P)])
|
||||
src2_shape = [128, 512]
|
||||
src2_layout = TileLayout(S[(512, 128) : (1 @ F, 1 @ P)])
|
||||
dst_shape = src1_shape
|
||||
dst_layout = src1_layout
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.add(C_sbuf, A_sbuf, B_sbuf)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "binary"})
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 512):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorscalar(C_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], B_sbuf[p_loop, b_loop], "add", T.bool(False)) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": binary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_binary_broadcast2():
|
||||
src1_shape = [32, 128, 512]
|
||||
src1_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
|
||||
src2_shape = [128, 512]
|
||||
src2_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
|
||||
dst_shape = src1_shape
|
||||
dst_layout = src1_layout
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.add(C_sbuf, A_sbuf, B_sbuf)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "binary"})
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 128):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensortensor(C_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 128 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 128 + f_loop], B_sbuf[p_loop, b_loop % 4 * 128 + f_loop], "add") # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": binary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_binary_broadcast3():
|
||||
src1_shape = [128, 512]
|
||||
src1_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
|
||||
src2_shape = [32, 128, 512]
|
||||
src2_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
|
||||
dst_shape = src1_shape
|
||||
dst_layout = src1_layout
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.add(C_sbuf, A_sbuf, B_sbuf[0])
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "binary"})
|
||||
A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensortensor(C_sbuf[p_loop, b_loop * 128 + f_loop], A_sbuf[p_loop, b_loop * 128 + f_loop], B_sbuf[p_loop, b_loop * 4096 + f_loop], "add") # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": binary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_binary_with_guard():
|
||||
src1_shape = [32, 128, 512]
|
||||
src1_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
|
||||
src2_shape = [128, 512]
|
||||
src2_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
|
||||
dst_shape = src1_shape
|
||||
dst_layout = src1_layout
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for j in range(4):
|
||||
Tx.add(C_sbuf[:, :, 0:j*128], A_sbuf[:, :, 0:j*128], B_sbuf[:, 0:j*128])
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "binary"})
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
for j, b_loop in T.grid(4, 96):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
|
||||
if b_loop % 3 - j < 0:
|
||||
T.nki.tensortensor(C_sbuf[p_loop, b_loop % 3 * 4096 + b_loop // 3 * 128 + f_loop], A_sbuf[p_loop, b_loop % 3 * 4096 + b_loop // 3 * 128 + f_loop], B_sbuf[p_loop, b_loop % 3 * 128 + f_loop], "add") # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": binary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,772 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir import assert_structural_equal as _assert_structural_equal
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.tirx.layout import F, P, S, TileLayout
|
||||
from tvm.tirx.stmt_functor import ir_transform
|
||||
|
||||
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
|
||||
|
||||
|
||||
def _strip_exec_scope_stmt(stmt):
|
||||
def _postorder(node):
|
||||
if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
|
||||
return node.body
|
||||
return node
|
||||
|
||||
return ir_transform(
|
||||
stmt,
|
||||
preorder=lambda _node: None,
|
||||
postorder=_postorder,
|
||||
only_enable=["tirx.AttrStmt"],
|
||||
)
|
||||
|
||||
|
||||
def assert_structural_equal(lhs, rhs, *args, **kwargs):
|
||||
if isinstance(lhs, tvm.tirx.PrimFunc):
|
||||
lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
|
||||
if isinstance(rhs, tvm.tirx.PrimFunc):
|
||||
rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
|
||||
_assert_structural_equal(lhs, rhs, *args, **kwargs)
|
||||
|
||||
|
||||
def test_simple_activation_reduce():
|
||||
A_shape = (128, 512)
|
||||
A_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
B_shape = (128, 512)
|
||||
B_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
C_shape = (128, 1)
|
||||
C_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def activation_reduce():
|
||||
T.device_entry()
|
||||
A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
|
||||
B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
|
||||
C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
|
||||
Tx.unary_reduce(B, C, A, "sqrt", "sum", reduce_axes=1)
|
||||
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "activation_reduce"})
|
||||
const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
|
||||
A = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
B = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
C = T.alloc_buffer((128, 1), scope="trn.sbuf")
|
||||
for b_loop in range(1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.activation_reduce(C[p_loop, 0], B[p_loop, f_loop], A[p_loop, f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": activation_reduce})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_activation_reduce_in_loop():
|
||||
A_shape = (32, 512, 128)
|
||||
A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
|
||||
B_shape = (16, 512, 128)
|
||||
B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
|
||||
C_shape = (16, 128)
|
||||
C_layout = TileLayout(S[(2, 4, 2, 128) : (2 @ F, 4 @ F, 1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def activation_reduce():
|
||||
T.device_entry()
|
||||
A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
|
||||
B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
|
||||
C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
|
||||
for i in range(2):
|
||||
Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=1)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "activation_reduce"})
|
||||
const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
|
||||
A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
|
||||
C = T.alloc_buffer((128, 16), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(2, 16):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 + b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop % 8 // 2 * 2048 + b_loop // 8 * 1024 + b_loop % 2 * 512 + f_loop], A[p_loop, i * 8192 + b_loop * 512 + f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop]) # noqa: E501
|
||||
# fmt: off
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": activation_reduce})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_activation_reduce_in_loop2():
|
||||
A_shape = (32, 512, 128)
|
||||
A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
|
||||
B_shape = (16, 512, 128)
|
||||
B_layout = TileLayout(S[(16 * 512, 128) : (1 @ F, 1 @ P)])
|
||||
C_shape = (16, 128)
|
||||
C_layout = TileLayout(S[(2, 4, 2, 128) : (2 @ F, 4 @ F, 1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def activation_reduce():
|
||||
T.device_entry()
|
||||
A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
|
||||
B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
|
||||
C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
|
||||
for i in range(2):
|
||||
Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=1)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "activation_reduce"})
|
||||
const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
|
||||
A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
|
||||
C = T.alloc_buffer((128, 16), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(2, 16):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 + b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop * 512 + f_loop], A[p_loop, i * 8192 + b_loop * 512 + f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop]) # noqa: E501
|
||||
# fmt: off
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": activation_reduce})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_activation_reduce_two_stage():
|
||||
A_shape = (32, 512, 128)
|
||||
A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
|
||||
B_shape = (16, 512, 128)
|
||||
B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
|
||||
C_shape = (1, 128)
|
||||
C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def activation_reduce():
|
||||
T.device_entry()
|
||||
A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
|
||||
B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
|
||||
C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
|
||||
for i in range(2):
|
||||
Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1))
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "activation_reduce"})
|
||||
partial_reduce = T.alloc_buffer((128, 8), scope="trn.sbuf")
|
||||
const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
|
||||
A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
|
||||
C = T.alloc_buffer((128, 1), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(2, 1):
|
||||
for reduction_b_loop in range(8):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
|
||||
T.nki.activation_reduce(partial_reduce[p_loop, reduction_b_loop], B[p_loop, reduction_b_loop % 4 * 2048 + reduction_b_loop // 4 * 1024 + f_loop], A[p_loop, i * 8192 + reduction_b_loop * 1024 + f_loop], "sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0)) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(8, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensorreduce(C[p_loop, 0], partial_reduce[p_loop, f_loop], "add", T.bool(False), -1) # noqa: E501
|
||||
# fmt: off
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": activation_reduce})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_activation_reduce_with_bias_scale():
|
||||
A_shape = (32, 512, 128)
|
||||
A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
|
||||
B_shape = (16, 512, 128)
|
||||
B_layout = TileLayout(S[(16 * 512, 128) : (1 @ F, 1 @ P)])
|
||||
C_shape = (16, 128)
|
||||
C_layout = TileLayout(S[(2, 4, 2, 128) : (2 @ F, 4 @ F, 1 @ F, 1 @ P)])
|
||||
bias_shape = 128
|
||||
bias_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def activation_reduce():
|
||||
T.device_entry()
|
||||
A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
|
||||
B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
|
||||
C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
|
||||
bias = T.alloc_buffer(bias_shape, dtype="float32", scope="trn.sbuf", layout=bias_layout)
|
||||
for i in range(2):
|
||||
Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=1, bias=bias, scale=2.0) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "activation_reduce"})
|
||||
A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
|
||||
C = T.alloc_buffer((128, 16), scope="trn.sbuf")
|
||||
bias = T.alloc_buffer((128, 1), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(2, 16):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 + b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop * 512 + f_loop], A[p_loop, i * 8192 + b_loop * 512 + f_loop], "sqrt", "add", bias[p_loop, 0], T.float32(2.0)) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": activation_reduce})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_simple_tensor_scalar_reduce():
|
||||
A_shape = (128, 512)
|
||||
A_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
B_shape = (128, 512)
|
||||
B_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
C_shape = (128, 1)
|
||||
C_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def tensor_scalar_reduce():
|
||||
T.device_entry()
|
||||
A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
|
||||
B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
|
||||
C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
|
||||
Tx.binary_reduce(B, C, A, 1.0, "add", "sum", reduce_axes=1)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "tensor_scalar_reduce"})
|
||||
A = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
B = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
C = T.alloc_buffer((128, 1), scope="trn.sbuf")
|
||||
for b_loop in range(1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorscalar_reduce(C[p_loop, 0], B[p_loop, f_loop], A[p_loop, f_loop], T.float32(1.0), "add", "add", T.bool(False)) # noqa: E501
|
||||
# fmt: off
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": tensor_scalar_reduce})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_tensor_tensor_reduce_fail():
|
||||
A_shape = (128, 512)
|
||||
A_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
B_shape = (128, 512)
|
||||
B_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
D_shape = (128, 512)
|
||||
D_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
C_shape = (128, 1)
|
||||
C_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def tensor_scalar_reduce():
|
||||
T.device_entry()
|
||||
A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
|
||||
B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
|
||||
C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
|
||||
D = T.alloc_buffer(D_shape, dtype="float32", scope="trn.sbuf", layout=D_layout)
|
||||
Tx.binary_reduce(B, C, A, D, "add", "sum", reduce_axes=1)
|
||||
|
||||
# fmt: off
|
||||
with pytest.raises(Exception):
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": tensor_scalar_reduce})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
|
||||
|
||||
def test_tensor_scalar_reduce_complex():
|
||||
src1_shape = [32, 128, 512]
|
||||
src1_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
|
||||
src2_shape = [128, 512]
|
||||
src2_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
|
||||
dst_shape = src1_shape
|
||||
dst_layout = src1_layout
|
||||
reduce_dst_shape = [128, 512]
|
||||
reduce_dst_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def tensor_scalar_reduce() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
D_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout) # noqa: E501
|
||||
Tx.binary_reduce(C_sbuf, D_sbuf, B_sbuf, A_sbuf, "add", "sum", reduce_axes=0)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "tensor_scalar_reduce"})
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
D_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
for b_loop in range(512):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorscalar_reduce(D_sbuf[p_loop, b_loop % 4 * 128 + b_loop // 4], C_sbuf[p_loop, b_loop % 4 * 4096 + f_loop * 128 + b_loop // 4], A_sbuf[p_loop, b_loop % 4 * 4096 + f_loop * 128 + b_loop // 4], B_sbuf[p_loop, b_loop % 4 * 128 + b_loop // 4], "add", "add", T.bool(True)) # noqa: E501
|
||||
# fmt: off
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": tensor_scalar_reduce})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_tensor_scalar_reduce_two_stage():
|
||||
src1_shape = [512, 1024, 4]
|
||||
src1_layout = TileLayout(S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)])
|
||||
dst1_shape = src1_shape
|
||||
dst1_layout = src1_layout
|
||||
reduce_dst_shape = [512]
|
||||
reduce_dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def tensor_scalar_reduce() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(dst1_shape, "float32", scope="trn.sbuf", layout=dst1_layout)
|
||||
C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout) # noqa: E501
|
||||
Tx.binary_reduce(B_sbuf, C_sbuf, A_sbuf, 1.0, "add", "sum", reduce_axes=(1, 2))
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "tensor_scalar_reduce"})
|
||||
partial_reduce = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
for b_loop in range(4):
|
||||
for reduction_b_loop in range(4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensorscalar_reduce(partial_reduce[p_loop, reduction_b_loop], B_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], A_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], T.float32(1.0), "add", "add", T.bool(False)) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(4, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensorreduce(C_sbuf[p_loop, b_loop], partial_reduce[p_loop, f_loop], "add", T.bool(False), -1) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": tensor_scalar_reduce})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_vector_chain():
|
||||
src1_shape = [32, 128, 512]
|
||||
src1_layout = TileLayout(S[(32, 128, 4, 128) : (1 @ F, 32 @ F, 32 * 128 @ F, 1 @ P)])
|
||||
src2_shape = [128, 512]
|
||||
src2_layout = TileLayout(S[(512, 128) : (1 @ F, 1 @ P)])
|
||||
src3_shape = [512]
|
||||
src3_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
|
||||
dst_shape = src1_shape
|
||||
dst_layout = src1_layout
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
|
||||
_C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
D_sbuf = T.alloc_buffer(src3_shape, "float32", scope="trn.sbuf", layout=src3_layout)
|
||||
E_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.binary_chain(E_sbuf, A_sbuf, B_sbuf, D_sbuf, "add", "add", reverse1=True)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "binary"})
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
_C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
D_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
E_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 512):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
|
||||
T.nki.scalar_tensor_scalar(E_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], B_sbuf[p_loop, b_loop], D_sbuf[p_loop, b_loop % 4], "add", "add", T.bool(False), T.bool(True)) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": binary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_vector_chain_2():
|
||||
src1_shape = [32, 128, 512]
|
||||
src1_layout = TileLayout(S[(32, 128, 4, 128) : (1 @ F, 32 @ F, 32 * 128 @ F, 1 @ P)])
|
||||
src2_shape = [128, 512]
|
||||
src2_layout = TileLayout(S[(512, 128) : (1 @ F, 1 @ P)])
|
||||
src3_shape = src1_shape
|
||||
src3_layout = src1_layout
|
||||
dst_shape = src1_shape
|
||||
dst_layout = src1_layout
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
|
||||
_C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
D_sbuf = T.alloc_buffer(src3_shape, "float32", scope="trn.sbuf", layout=src3_layout)
|
||||
E_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.binary_chain(E_sbuf, A_sbuf, B_sbuf, D_sbuf, "add", "add", reverse1=True)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "binary"})
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
_C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
D_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
E_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 512):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
|
||||
T.nki.scalar_tensor_tensor(E_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], B_sbuf[p_loop, b_loop], D_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], "add", "add", T.bool(False), T.bool(True)) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": binary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_reduce_negate():
|
||||
src_shape = [128, 512, 4]
|
||||
src_layout = TileLayout(S[(128, 512, 4) : (1 @ P, 4 @ F, 1 @ F)])
|
||||
dst_shape = [128, 4]
|
||||
dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def reduction():
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
Tx.reduce_negate(B_sbuf[:, i], A_sbuf[:, :, i], reduce_op="sum", reduce_axes=-2)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "reduction"})
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(4, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorreduce(B_sbuf[p_loop, i], A_sbuf[p_loop, f_loop * 4 + i], "add", True, -1) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": reduction})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_binary_reduce_guard():
|
||||
src_shape = [512, 512]
|
||||
src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
reduce_dst_shape = [512]
|
||||
reduce_dst_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary_reduce() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout) # noqa: E501
|
||||
for j in range(4):
|
||||
for i in range(4):
|
||||
Tx.binary_reduce(B_sbuf[0:128*(j+1), 0:128*(i+1)], C_sbuf[0:128*(j+1)], A_sbuf[0:128*(j+1), 0:128*(i+1)], 0.0, "add", "sum", [-1]) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "binary_reduce"})
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
for j, i, b_loop in T.grid(4, 4, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
if b_loop - j < 1 and f_loop < i * 128 + 128:
|
||||
T.nki.tensorscalar_reduce(C_sbuf[p_loop, b_loop], B_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0), "add", "add", T.bool(False)) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": binary_reduce})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_unary_reduce_guard():
|
||||
src_shape = [512, 512]
|
||||
src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
reduce_dst_shape = [512]
|
||||
reduce_dst_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def unary_reduce() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout) # noqa: E501
|
||||
for j in range(4):
|
||||
for i in range(4):
|
||||
Tx.unary_reduce(B_sbuf[0:128*(j+1), 0:128*(i+1)], C_sbuf[0:128*(j+1)], A_sbuf[0:128*(j+1), 0:128*(i+1)], "sqrt", "sum", reduce_axes=[-1]) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "unary_reduce"})
|
||||
const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
for j, i, b_loop in T.grid(4, 4, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
|
||||
if b_loop - j < 1 and f_loop < i * 128 + 128:
|
||||
T.nki.activation_reduce(C_sbuf[p_loop, b_loop], B_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], "sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0)) # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": unary_reduce})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_binary_chain_guard():
|
||||
src_shape = [512, 512]
|
||||
src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
src2_shape = [512, 1]
|
||||
src2_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def binary_chain() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for j in range(4):
|
||||
for i in range(4):
|
||||
Tx.binary_chain(C_sbuf[0:128*(j+1), 0:128*(i+1)], A_sbuf[0:128*(j+1), 0:128*(i+1)], B_sbuf[0:128*(j+1), 0], 1.0, "add", "sub", reverse1=True) # noqa: E501
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "binary_chain"})
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for j, i, b_loop in T.grid(4, 4, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
if b_loop - j < 1 and f_loop < i * 128 + 128:
|
||||
T.nki.scalar_tensor_scalar(C_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], B_sbuf[p_loop, b_loop], T.float32(1.0), "add", "sub", T.bool(False), T.bool(True)) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": binary_chain})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_activation_reduce_two_stage_workspace():
|
||||
A_shape = (32, 512, 128)
|
||||
A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
|
||||
B_shape = (16, 512, 128)
|
||||
B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
|
||||
C_shape = (1, 128)
|
||||
C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def activation_reduce():
|
||||
T.device_entry()
|
||||
intermediate_buffer = T.alloc_buffer((128, 16), scope="trn.sbuf")
|
||||
A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
|
||||
B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
|
||||
C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
|
||||
for i in range(2):
|
||||
Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1), workspace={"partial_reduce": intermediate_buffer}) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "activation_reduce"})
|
||||
const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
|
||||
intermediate_buffer = T.alloc_buffer((128, 16), scope="trn.sbuf")
|
||||
A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
|
||||
C = T.alloc_buffer((128, 1), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(2, 1):
|
||||
for reduction_b_loop in range(8):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
|
||||
T.nki.activation_reduce(intermediate_buffer[p_loop, reduction_b_loop], B[p_loop, reduction_b_loop % 4 * 2048 + reduction_b_loop // 4 * 1024 + f_loop], A[p_loop, i * 8192 + reduction_b_loop * 1024 + f_loop], "sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0)) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(8, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensorreduce(C[p_loop, 0], intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1) # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": activation_reduce})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_tensor_scalar_reduce_two_stage_workspace():
|
||||
src1_shape = [512, 1024, 4]
|
||||
src1_layout = TileLayout(S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)])
|
||||
dst1_shape = src1_shape
|
||||
dst1_layout = src1_layout
|
||||
reduce_dst_shape = [512]
|
||||
reduce_dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def tensor_scalar_reduce() -> None:
|
||||
T.device_entry()
|
||||
intermediate_buffer = T.alloc_buffer((128, 8), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(dst1_shape, "float32", scope="trn.sbuf", layout=dst1_layout)
|
||||
C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout) # noqa: E501
|
||||
Tx.binary_reduce(B_sbuf, C_sbuf, A_sbuf, 1.0, "add", "sum", reduce_axes=(1, 2), workspace={"partial_reduce": intermediate_buffer}) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "tensor_scalar_reduce"})
|
||||
intermediate_buffer = T.alloc_buffer((128, 8), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
for b_loop in range(4):
|
||||
for reduction_b_loop in range(4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensorscalar_reduce(intermediate_buffer[p_loop, reduction_b_loop], B_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], A_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], T.float32(1.0), "add", "add", T.bool(False)) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(4, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensorreduce(C_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": tensor_scalar_reduce})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_unary_reduce_complex():
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def unary_reduce():
|
||||
T.device_entry()
|
||||
p = T.alloc_buffer((128, 8192), "float16", scope="trn.sbuf", layout="PF")
|
||||
rowsum_p = T.alloc_buffer((2, 128, 1), scope="trn.sbuf", layout="FPF")
|
||||
qk = T.alloc_buffer((2, 128, 8192), scope="trn.sbuf", layout="FPF")
|
||||
running_max = T.alloc_buffer((16384, 1), dtype="float32", scope="trn.sbuf", layout="PF")
|
||||
for i in range(4):
|
||||
Tx.unary_reduce(p[0:128, 0:8192], rowsum_p[i % 2, 0:128, 0], qk[i % 2, 0:128, 0:8192], "exp", "sum", bias=running_max[i * 128:i * 128 + 128, 0]) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "unary_reduce"})
|
||||
p = T.alloc_buffer((128, 8192), "float16", scope="trn.sbuf")
|
||||
rowsum_p = T.alloc_buffer((128, 2), scope="trn.sbuf")
|
||||
qk = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
running_max = T.alloc_buffer((128, 128), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(4, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(8192, annotations={"nki_dim": "F"}):
|
||||
T.nki.activation_reduce(rowsum_p[p_loop, i % 2], p[p_loop, f_loop], qk[p_loop, i % 2 * 8192 + f_loop], "exp", "add", running_max[p_loop, i], T.float32(1.0)) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": unary_reduce})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,841 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir import assert_structural_equal as _assert_structural_equal
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.tirx.layout import F, P, S, TileLayout
|
||||
from tvm.tirx.stmt_functor import ir_transform
|
||||
|
||||
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
|
||||
|
||||
|
||||
def _strip_exec_scope_stmt(stmt):
|
||||
def _postorder(node):
|
||||
if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
|
||||
return node.body
|
||||
return node
|
||||
|
||||
return ir_transform(
|
||||
stmt,
|
||||
preorder=lambda _node: None,
|
||||
postorder=_postorder,
|
||||
only_enable=["tirx.AttrStmt"],
|
||||
)
|
||||
|
||||
|
||||
def assert_structural_equal(lhs, rhs, *args, **kwargs):
|
||||
if isinstance(lhs, tvm.tirx.PrimFunc):
|
||||
lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
|
||||
if isinstance(rhs, tvm.tirx.PrimFunc):
|
||||
rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
|
||||
_assert_structural_equal(lhs, rhs, *args, **kwargs)
|
||||
|
||||
|
||||
def test_simple_copy():
|
||||
src_shape = [128, 512]
|
||||
src_layout = T.TileLayout(T.S[(128, 512) : (512, 1)])
|
||||
dst_shape = [128, 512]
|
||||
dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.copy(A_sbuf, A)
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (128, 512), layout=None)
|
||||
A_1 = T.decl_buffer((65536,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
|
||||
T.nki.load(A_sbuf[p_loop, f_loop], A_1[p_loop * 512 + f_loop])
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_simple_copy_2():
|
||||
src_shape = [128, 512]
|
||||
src_layout = TileLayout(S[(128, 4, 128) : (512, 128, 1)])
|
||||
|
||||
dst_shape = [128, 512]
|
||||
dst_layout = TileLayout(S[(128, 4, 128) : (4 @ F, 1 @ F, 1 @ P)])
|
||||
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.copy(A_sbuf, A)
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (128, 512), layout=None)
|
||||
A_1 = T.decl_buffer((65536,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 512):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(0, 1, annotations={"nki_dim": "F"}):
|
||||
T.nki.load(A_sbuf[p_loop, b_loop], A_1[b_loop * 128 + p_loop])
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_in_a_loop():
|
||||
src_shape = [512, 512]
|
||||
src_layout = T.TileLayout(T.S[(4, 128, 512) : (512 * 128, 512, 1)])
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
Tx.copy(A_sbuf[i * 128 : i * 128 + 128, :], A[i * 128 : i * 128 + 128, :])
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (512, 512), layout=None)
|
||||
A_1 = T.decl_buffer((262144,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(4, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
|
||||
T.nki.load(
|
||||
A_sbuf[p_loop, i * 512 + f_loop], A_1[i * 65536 + p_loop * 512 + f_loop]
|
||||
)
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_in_a_loop_2():
|
||||
src_shape = [512, 512]
|
||||
src_layout = T.TileLayout(T.S[(128, 2048) : (2048, 1)])
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
|
||||
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
A_sbuf_view = A_sbuf.view(128, 4, 512)
|
||||
A_view = A.view(128, 4, 512)
|
||||
for i in range(4):
|
||||
Tx.copy(A_sbuf_view[:, i, :], A_view[:, i, :])
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (512, 512), layout=None)
|
||||
_A_flat = T.decl_buffer((262144,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
A_sbuf_view = T.decl_buffer((128, 2048), data=A_sbuf.data, scope="trn.sbuf", layout=None)
|
||||
A_view = T.decl_buffer((262144,), data=A.data, layout=None)
|
||||
for i, b_loop in T.grid(4, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
|
||||
T.nki.load(
|
||||
A_sbuf_view[p_loop, i * 512 + f_loop],
|
||||
A_view[p_loop * 2048 + i * 512 + f_loop],
|
||||
)
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod.show()
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_transpose():
|
||||
src_shape = [512, 512]
|
||||
src_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.copy(B_sbuf, A_sbuf)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
|
||||
acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for b_loop in range(16):
|
||||
for extend_b_loop in range(1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
|
||||
T.nki.matmul(acc_psum[b_loop % 8, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True)) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + b_loop], acc_psum[b_loop % 8, p_loop, f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_transpose_2():
|
||||
src_shape = [65536]
|
||||
src_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [4, 65536]
|
||||
dst_layout = TileLayout(S[(4, 128, 128, 4) : (4 @ F, 16 @ F, 1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
Tx.copy(B_sbuf[i, :], A_sbuf)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
|
||||
acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
|
||||
A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for i in range(4):
|
||||
for b_loop in range(4):
|
||||
for extend_b_loop in range(1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
|
||||
T.nki.matmul(acc_psum[b_loop, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, lhs_f_loop * 4 + b_loop], identity[p_loop, rhs_f_loop], T.bool(True)) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + i * 4 + b_loop], acc_psum[b_loop, p_loop, f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_different_f():
|
||||
src_shape = [512, 64]
|
||||
src_layout = TileLayout(S[(4, 128, 4, 4, 4) : (64 @ F, 1 @ P, 16 @ F, 4 @ F, 1 @ F)])
|
||||
dst_shape = [512, 64]
|
||||
dst_layout = TileLayout(S[(4, 128, 4, 4, 4) : (64 @ F, 1 @ P, 4 @ F, 16 @ F, 1 @ F)])
|
||||
|
||||
@T.prim_func
|
||||
def copy() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.copy(B_sbuf, A_sbuf)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
A_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 64):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(0, 4, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensor_copy(
|
||||
B_sbuf[
|
||||
p_loop,
|
||||
b_loop // 16 * 64 + b_loop % 4 * 16 + b_loop % 16 // 4 * 4 + f_loop,
|
||||
],
|
||||
A_sbuf[p_loop, b_loop * 4 + f_loop],
|
||||
)
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_different_shape():
|
||||
src_shape = [512, 64]
|
||||
src_layout = TileLayout(S[(4, 128, 4, 4, 4) : (64 @ F, 1 @ P, 16 @ F, 4 @ F, 1 @ F)])
|
||||
dst_shape = [4, 128, 4]
|
||||
dst_layout = TileLayout(S[(4, 128, 4) : (4 @ F, 1 @ P, 1 @ F)])
|
||||
|
||||
@T.prim_func
|
||||
def copy() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
B_sbuf_view = B_sbuf.view(512, 4)
|
||||
Tx.copy(B_sbuf_view, A_sbuf[:, 0:4])
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
A_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 16), scope="trn.sbuf")
|
||||
_B_sbuf_view = T.decl_buffer((128, 16), data=B_sbuf.data, scope="trn.sbuf", layout=None)
|
||||
for b_loop in T.serial(0, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(0, 4, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensor_copy(
|
||||
B_sbuf[p_loop, b_loop * 4 + f_loop],
|
||||
A_sbuf[p_loop, b_loop * 64 + f_loop],
|
||||
)
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_irregular_shape():
|
||||
src_shape = [128, 10000]
|
||||
src_layout = TileLayout(S[(128, 10000) : (10000, 1)])
|
||||
dst_shape = [128, 512]
|
||||
dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
Tx.copy(A[:, i * 512 : i * 512 + 512], A_sbuf)
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (128, 10000), layout=None)
|
||||
A_1 = T.decl_buffer((1280000,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(4, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
|
||||
T.nki.store(A_1[p_loop * 10000 + i * 512 + f_loop], A_sbuf[p_loop, f_loop])
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_different_shape_dim():
|
||||
src_shape = [32, 128, 512]
|
||||
src_layout = TileLayout(S[(32, 128, 512) : (128 * 512, 128, 1)])
|
||||
dst_shape = [128, 512]
|
||||
dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(32):
|
||||
Tx.copy(A_sbuf, A[i, :, :])
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (32, 128, 512), layout=None)
|
||||
A_1 = T.decl_buffer((2097152,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(32, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.load(A_sbuf[p_loop, f_loop], A_1[i * 65536 + p_loop * 128 + f_loop])
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_with_offset():
|
||||
src_shape = [256, 512]
|
||||
src_layout = TileLayout(S[(256, 512) : (512, 1)])
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(2):
|
||||
Tx.copy(A_sbuf[i * 256 : i * 256 + 256, :], A)
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (256, 512), layout=None)
|
||||
A_1 = T.decl_buffer((131072,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(2, 2):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
|
||||
T.nki.load(
|
||||
A_sbuf[p_loop, i * 1024 + b_loop * 512 + f_loop],
|
||||
A_1[b_loop * 65536 + p_loop * 512 + f_loop],
|
||||
)
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_large_dma_copy():
|
||||
src_shape = [512, 4096]
|
||||
src_layout = T.TileLayout(T.S[(4, 128, 4096) : (4096 * 128, 4096, 1)])
|
||||
dst_shape = [512, 4096]
|
||||
dst_layout = TileLayout(S[(4, 128, 4096) : (4096 @ F, 1 @ P, 1 @ F)])
|
||||
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
Tx.copy(A_sbuf[i * 128 : i * 128 + 128, :], A[i * 128 : i * 128 + 128, :])
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (512, 4096), layout=None)
|
||||
A_1 = T.decl_buffer((2097152,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(4, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(0, 4096, annotations={"nki_dim": "F"}):
|
||||
T.nki.load(
|
||||
A_sbuf[p_loop, i * 4096 + f_loop],
|
||||
A_1[i * 524288 + p_loop * 4096 + f_loop],
|
||||
)
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_with_inst_size_limit():
|
||||
src_shape = [512, 4096]
|
||||
src_layout = dst_layout = TileLayout(S[(4, 128, 4096) : (4096 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
T.device_entry()
|
||||
B_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
Tx.copy(A_sbuf[i * 128 : i * 128 + 128, :], B_sbuf[i * 128 : i * 128 + 128, :])
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(4, 8):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensor_copy(
|
||||
A_sbuf[p_loop, i * 4096 + b_loop * 512 + f_loop],
|
||||
B_sbuf[p_loop, i * 4096 + b_loop * 512 + f_loop],
|
||||
)
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_with_complex_index():
|
||||
A_shape = [4096, 4096]
|
||||
A_layout = T.TileLayout(T.S[(4096, 4096) : (1, 4096)])
|
||||
A_sbuf_shape = (2, 2048, 1024)
|
||||
A_sbuf_layout = TileLayout(S[(2, 2048, 8, 128) : (16384 @ F, 1 @ F, 2048 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle, ) -> None:
|
||||
A = T.match_buffer(A_ptr, A_shape, "float32", layout=A_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(A_sbuf_shape, "float32", scope="trn.sbuf", layout=A_sbuf_layout)
|
||||
Tx.copy(A_sbuf[1, 0:2048, 0:1024], A[2048: 4096, 3072:4096])
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (4096, 4096), layout=None)
|
||||
A_1 = T.decl_buffer((16777216,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 32768), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 8):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
|
||||
T.nki.load(A_sbuf[p_loop, b_loop * 2048 + f_loop + 16384], A_1[b_loop * 524288 + p_loop * 4096 + f_loop + 12584960]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_with_complex_index_2():
|
||||
A_sbuf_shape = [4096, 4096]
|
||||
A_sbuf_layout = T.TileLayout(T.S[(4096, 32, 128) : (1 @ F, 4096 @ F, 1 @ P)])
|
||||
A_shape = (2, 2048, 1024)
|
||||
A_layout = T.TileLayout(T.S[(2, 2048, 1024) : (2048 * 1024, 1, 2048)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle, ) -> None:
|
||||
A = T.match_buffer(A_ptr, A_shape, "float32", layout=A_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(A_sbuf_shape, "float32", scope="trn.sbuf", layout=A_sbuf_layout)
|
||||
Tx.copy(A_sbuf[2048: 4096, 3072:4096], A[1, 0:2048, 0:1024])
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (2, 2048, 1024), layout=None)
|
||||
A_1 = T.decl_buffer((4194304,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 131072), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 8):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
|
||||
T.nki.load(A_sbuf[p_loop, b_loop * 4096 + f_loop + 100352], A_1[b_loop * 262144 + p_loop * 2048 + f_loop + 2097152]) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_transpose_with_workspace():
|
||||
src_shape = [512, 512]
|
||||
src_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
identity = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf")
|
||||
acc_psum = T.alloc_buffer((1, 128, 512), "float32", scope="trn.psum", allocated_addr=(0, 0))
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
|
||||
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
|
||||
Tx.copy(B_sbuf, A_sbuf, workspace={"identity": identity, "acc_psum": acc_psum})
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
|
||||
acc_psum = T.alloc_buffer((1, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
|
||||
for b_loop in range(16):
|
||||
for extend_b_loop in range(1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
|
||||
T.nki.matmul(acc_psum[0, lhs_f_loop, extend_b_loop * 128 + rhs_f_loop], A_sbuf[p_loop, b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True)) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + b_loop], acc_psum[0, p_loop, f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_with_guard():
|
||||
src_shape = [512, 512]
|
||||
src_layout = T.TileLayout(T.S[(4, 128, 512) : (512 * 128, 512, 1)])
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for j in range(4):
|
||||
for i in range(4):
|
||||
Tx.copy(A_sbuf[i * 128 : i * 128 + 128, 0:128*j], A[i * 128 : i * 128 + 128, 0:128*j]) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (512, 512), layout=None)
|
||||
A_1 = T.decl_buffer((262144,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for j, i, b_loop in T.grid(4, 4, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 384, annotations={"nki_dim":"F"}):
|
||||
if f_loop < j * 128:
|
||||
T.nki.load(A_sbuf[p_loop, i * 512 + f_loop], A_1[i * 65536 + p_loop * 512 + f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_with_guard_2():
|
||||
src_shape = [512, 512]
|
||||
src_layout = T.TileLayout(T.S[(4, 128, 512) : (512 * 128, 512, 1)])
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for j in range(4):
|
||||
for i in range(4):
|
||||
Tx.copy(A_sbuf[0:128*j, 0:128*i], A[0:128*j, 0:128*i])
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
|
||||
A = T.match_buffer(A_ptr, (512, 512), layout=None)
|
||||
A_1 = T.decl_buffer((262144,), data=A.data, layout=None)
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for j, i, b_loop in T.grid(4, 4, 3):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 384, annotations={"nki_dim":"F"}):
|
||||
if b_loop - j < 0 and f_loop < i * 128:
|
||||
T.nki.load(A_sbuf[p_loop, b_loop * 512 + f_loop], A_1[b_loop * 65536 + p_loop * 512 + f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_transpose_with_guard():
|
||||
src_shape = [512, 512]
|
||||
src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
for j in range(4):
|
||||
Tx.copy(B_sbuf[i * 128 : i * 128 + 128, 0:128*j], A_sbuf[i * 128 : i * 128 + 128, 0:128*j]) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
|
||||
acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for i, j, b_loop in T.grid(4, 4, 3):
|
||||
for extend_b_loop in range(1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
|
||||
if b_loop - j < 0:
|
||||
T.nki.matmul(acc_psum[b_loop, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, i * 512 + b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True)) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
if b_loop - j < 0:
|
||||
T.nki.tensor_copy(B_sbuf[p_loop, i * 512 + f_loop * 4 + b_loop], acc_psum[b_loop, p_loop, f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_with_specified_max_inst_size():
|
||||
src_shape = [128, 512]
|
||||
src_layout = "PF"
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.copy(A_sbuf, B_sbuf, max_inst_size=128)
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf", layout=None)
|
||||
B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf", layout=None)
|
||||
for b_loop in T.serial(0, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensor_copy(A_sbuf[p_loop, b_loop * 128 + f_loop], B_sbuf[p_loop, b_loop * 128 + f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_copy_transpose_with_extended_f():
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((128, 2048), "float32", scope="trn.sbuf", layout="PF")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), "float32", scope="trn.sbuf", layout="FP")
|
||||
Tx.copy(B_sbuf, A_sbuf)
|
||||
|
||||
@T.prim_func
|
||||
def expected(A_ptr: T.handle):
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
|
||||
acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for b_loop in range(4):
|
||||
for extend_b_loop in range(4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
|
||||
T.nki.matmul(acc_psum[b_loop, lhs_f_loop, extend_b_loop * 128 + rhs_f_loop], A_sbuf[p_loop, b_loop * 512 + extend_b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True)) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
|
||||
T.nki.tensor_copy(B_sbuf[p_loop, b_loop * 512 + f_loop], acc_psum[b_loop, p_loop, f_loop]) # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,583 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir import assert_structural_equal as _assert_structural_equal
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.tirx.layout import F, P, S, TileLayout
|
||||
from tvm.tirx.stmt_functor import ir_transform
|
||||
|
||||
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
|
||||
|
||||
|
||||
def _strip_exec_scope_stmt(stmt):
|
||||
def _postorder(node):
|
||||
if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
|
||||
return node.body
|
||||
return node
|
||||
|
||||
return ir_transform(
|
||||
stmt,
|
||||
preorder=lambda _node: None,
|
||||
postorder=_postorder,
|
||||
only_enable=["tirx.AttrStmt"],
|
||||
)
|
||||
|
||||
|
||||
def assert_structural_equal(lhs, rhs, *args, **kwargs):
|
||||
if isinstance(lhs, tvm.tirx.PrimFunc):
|
||||
lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
|
||||
if isinstance(rhs, tvm.tirx.PrimFunc):
|
||||
rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
|
||||
_assert_structural_equal(lhs, rhs, *args, **kwargs)
|
||||
|
||||
|
||||
def test_simple_gemm():
|
||||
A_layout = TileLayout(S[(128, 128) : (1 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)]).to_psum()
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((128, 128), "float32", scope="trn.psum", layout=C_layout)
|
||||
Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
A_sbuf = T.alloc_buffer((128, 128), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 128), scope="trn.sbuf")
|
||||
C_psum = T.alloc_buffer((1, 128, 128), scope="trn.psum")
|
||||
for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(1, 1, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
|
||||
T.nki.matmul(C_psum[0, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, lhs_f_loop], B_sbuf[p_loop, rhs_f_loop], True) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_larger_gemm():
|
||||
A_layout = TileLayout(S[(2, 128, 4, 128) : (512 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(2, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((256, 512), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((256, 256), "float32", scope="trn.psum", layout=C_layout)
|
||||
Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
A_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
C_psum = T.alloc_buffer((1, 128, 512), scope="trn.psum")
|
||||
for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 1, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
|
||||
T.nki.matmul(C_psum[0, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, lhs_b_loop * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, reduction_b_loop * 256 + rhs_f_loop], True) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm_in_a_loop():
|
||||
A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
|
||||
for i in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
|
||||
B_sbuf[512 * k : 512 * k + 512, :],
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
|
||||
for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 1, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
|
||||
T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + rhs_f_loop], True) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm_with_stride():
|
||||
A_layout = TileLayout(S[(4, 128, 128, 8) : (1024 @ F, 1 @ F, 1 @ P, 128 @ F)])
|
||||
B_layout = TileLayout(S[(128, 8, 2, 128) : (1 @ P, 512 @ F, 256 @ F, 2 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((512, 512, 2), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((512, 2, 256), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
|
||||
for i in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
A_sbuf[256 * i : 256 * i + 256, :, k],
|
||||
B_sbuf[:, k, :],
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 4095), scope="trn.sbuf")
|
||||
C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
|
||||
for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 1, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
|
||||
T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + reduction_b_loop * 256 + k * 128 + lhs_f_loop], B_sbuf[p_loop, reduction_b_loop * 1024 + k * 512 + rhs_f_loop * 2], True) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm_swap_lhs_rhs():
|
||||
A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)]).to_psum()
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
|
||||
for i in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
|
||||
B_sbuf[512 * k : 512 * k + 512, :],
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
|
||||
for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 2, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
|
||||
T.nki.matmul(C_psum[i, lhs_f_loop, rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop], True) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm_with_sbuf_output():
|
||||
A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
|
||||
for i in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_sbuf[256 * i : 256 * i + 256, :],
|
||||
A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
|
||||
B_sbuf[512 * k : 512 * k + 512, :],
|
||||
C_sbuf[256 * i : 256 * i + 256, :],
|
||||
)
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
buffer = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
for i, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2):
|
||||
for reduction_b_loop in range(4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
|
||||
T.nki.matmul(buffer[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop], True) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensor_copy(C_sbuf[lhs_f_loop, i * 512 + rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], buffer[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm_different_shape():
|
||||
A_layout = TileLayout(S[(2, 4, 128, 8, 128) : (4096 @ F, 1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)]).to_psum()
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((2, 512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
|
||||
for i in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
A_sbuf[1, 256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
|
||||
B_sbuf[512 * k : 512 * k + 512, :],
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
A_sbuf = T.alloc_buffer((128, 8192), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
|
||||
for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 2, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
|
||||
T.nki.matmul(C_psum[i, lhs_f_loop, rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop + 4096], True) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm_too_large_f_size():
|
||||
A_layout = TileLayout(S[(256, 128) : (1 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(128, 1024) : (1 @ P, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(2, 128, 1024) : (1024 @ F, 1 @ P, 1 @ F)]).to_psum()
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((256, 128), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((128, 1024), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((256, 1024), "float32", scope="trn.psum", layout=C_layout)
|
||||
Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
A_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
C_psum = T.alloc_buffer((4, 128, 512), scope="trn.psum")
|
||||
for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 512, annotations={"nki_dim":"rhs_F"}):
|
||||
T.nki.matmul(C_psum[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, lhs_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, rhs_b_loop * 512 + rhs_f_loop], True) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm_sbuf_output_with_workspace():
|
||||
A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
|
||||
C_psum = T.alloc_buffer((1, 128, 512), "float32", scope="trn.psum", allocated_addr=(0, 0))
|
||||
for i in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_sbuf[256 * i : 256 * i + 256, :],
|
||||
A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
|
||||
B_sbuf[512 * k : 512 * k + 512, :],
|
||||
C_sbuf[256 * i : 256 * i + 256, :],
|
||||
workspace={"acc_psum": C_psum}
|
||||
)
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
C_psum = T.alloc_buffer((1, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
|
||||
for i, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2):
|
||||
for reduction_b_loop in range(4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
|
||||
T.nki.matmul(C_psum[0, lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop], True) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensor_copy(C_sbuf[lhs_f_loop, i * 512 + rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], C_psum[0, lhs_f_loop, rhs_f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm_pf_mismatch_fail():
|
||||
A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(2, 128, 8, 128) : (128 @ F, 1 @ F, 256 @ F, 1 @ P)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((256, 1024), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
|
||||
for i in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
|
||||
B_sbuf[:, 512 * k : 512 * k + 512],
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
)
|
||||
# fmt: on
|
||||
with pytest.raises(Exception):
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
|
||||
|
||||
def test_gemm_transpose_AB():
|
||||
A_layout = TileLayout(S[(8, 128, 4, 128) : (128 @ F, 1 @ P, 1024 @ F, 1 @ F)])
|
||||
B_layout = TileLayout(S[(2, 128, 8, 128) : (128 @ F, 1 @ F, 256 @ F, 1 @ P)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((1024, 512), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((256, 1024), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
|
||||
for i in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
A_sbuf[512 * k : 512 * k + 512, 256 * i : 256 * i + 256],
|
||||
B_sbuf[:, 512 * k : 512 * k + 512],
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
transpose_A=True,
|
||||
transpose_B=True,
|
||||
)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
|
||||
for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 1, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
|
||||
T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + rhs_f_loop], True) # noqa: E501
|
||||
|
||||
#fmt: off
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm_guard():
|
||||
A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
|
||||
for i in range(2):
|
||||
for j in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_sbuf[0: 256 * i, 0: 128 * (j + 1)],
|
||||
A_sbuf[0: 256 * i, 0: 512 * (k + 1)],
|
||||
B_sbuf[0: 512 * (k + 1), 0: 128 * (j + 1)],
|
||||
C_sbuf[0: 256 * i, 0: 128 * (j + 1)],
|
||||
)
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
for i, j, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2, 2):
|
||||
for reduction_b_loop in range(8):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
|
||||
if reduction_b_loop - k * 4 < 4 and lhs_b_loop - j < 1 and 0 < i and reduction_b_loop - k * 4 < 4: # noqa: E501
|
||||
T.nki.matmul(acc_psum[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, rhs_b_loop * 1024 + reduction_b_loop * 128 + rhs_f_loop], True) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
|
||||
if 0 < i and lhs_b_loop - j < 1:
|
||||
T.nki.tensor_copy(C_sbuf[lhs_f_loop, rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], acc_psum[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop]) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm_guard2():
|
||||
A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
|
||||
for j in range(4):
|
||||
for i in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + (j+1) * 128],
|
||||
B_sbuf[512 * k : 512 * k + (j+1) * 128, :],
|
||||
C_psum[256 * i : 256 * i + 256, :],
|
||||
)
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
|
||||
for j, i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(4, 2, 2, 2, 1, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
|
||||
for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
|
||||
if reduction_b_loop - j < 1 and reduction_b_loop - j < 1:
|
||||
T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + rhs_f_loop], True) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,402 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir import assert_structural_equal
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.tirx.layout import F, P, S, TileLayout
|
||||
from tvm.tirx.trn.transform import TrnPrivateBufferAlloc
|
||||
|
||||
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
|
||||
|
||||
|
||||
def test_copy_transpose():
|
||||
src_shape = [512, 512]
|
||||
src_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.copy(B_sbuf, A_sbuf)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "copy"})
|
||||
T.device_entry()
|
||||
identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
|
||||
acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
|
||||
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
|
||||
A_sbuf = T.alloc_buffer((512, 512), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(128, 2048) : (1 @ P, 1@F)]))
|
||||
B_sbuf = T.alloc_buffer((512, 512), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(2048, 128) : (1@F, 1@P)]))
|
||||
Tx.copy(B_sbuf[0:512, 0:512], A_sbuf[0:512, 0:512], workspace={"acc_psum": acc_psum, "identity": identity}) # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_normal_copy():
|
||||
src_shape = [128, 512]
|
||||
src_layout = TileLayout(S[(128, 512) : (512, 1)])
|
||||
dst_shape = [128, 512]
|
||||
dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def copy(A_ptr: T.handle) -> None:
|
||||
A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.copy(A_sbuf, A)
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": copy})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], copy)
|
||||
|
||||
|
||||
def test_unary_with_bias_scale():
|
||||
src_shape = [512, 1024]
|
||||
src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
bias = T.float32(1.0)
|
||||
scale = T.float32(2.0)
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def unary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.exp(C_sbuf, A_sbuf, bias=bias, scale=scale)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "unary"})
|
||||
T.device_entry()
|
||||
const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(1.0))
|
||||
A_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(128, 4096) : (1@P, 1@F)]))
|
||||
C_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(128, 4096) : (1@P, 1@F)]))
|
||||
Tx.exp(C_sbuf[0:512, 0:1024], A_sbuf[0:512, 0:1024], T.float32(1.0), T.float32(2.0), workspace={"const_bias": const_bias}) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": unary})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_reduction_two_stage():
|
||||
src_shape = [128, 32, 4, 32]
|
||||
src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [128, 4]
|
||||
dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def reduction():
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.sum(B_sbuf, A_sbuf, axes=(1, 3))
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "reduction"})
|
||||
T.device_entry()
|
||||
partial_reduce = T.alloc_buffer((128, 32), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer((128, 32, 4, 32), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(128, 32 * 32 * 4) : (1@P, 1@F)]))
|
||||
B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(128, 4) : (1@P, 1@F)]))
|
||||
Tx.sum(B_sbuf[0:128, 0:4], A_sbuf[0:128, 0:32, 0:4, 0:32], [1, 3], False, workspace={"partial_reduce": partial_reduce}) # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": reduction})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gemm():
|
||||
A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 1 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
|
||||
for i in range(2):
|
||||
for k in range(2):
|
||||
Tx.gemm(
|
||||
C_sbuf[256 * i : 256 * i + 256, :],
|
||||
A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
|
||||
B_sbuf[512 * k : 512 * k + 512, :],
|
||||
C_sbuf[256 * i : 256 * i + 256, :],
|
||||
)
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "gemm"})
|
||||
T.device_entry()
|
||||
acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
|
||||
A_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(4, 128, 8, 128) : (1024@F, 1@F, 1@F, 1@P)])) # noqa: E501
|
||||
B_sbuf = T.alloc_buffer((1024, 256), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(8, 128, 2, 128) : (256@F, 1@P, 128@F, 1@F)])) # noqa: E501
|
||||
C_sbuf = T.alloc_buffer((512, 256), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(4, 128, 2, 128) : (256@F, 1@F, 128@F, 1@P)])) # noqa: E501
|
||||
for i, k in T.grid(2, 2):
|
||||
Tx.gemm(C_sbuf[256 * i:256 * i + 256, 0:256], A_sbuf[256 * i:256 * i + 256, 512 * k:512 * k + 512], B_sbuf[512 * k:512 * k + 512, 0:256], C_sbuf[256 * i:256 * i + 256, 0:256], False, False, T.float32(1.0), T.float32(0.0), workspace={"acc_psum": acc_psum}) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_binary_reduce_two_stage():
|
||||
src1_shape = [512, 1024, 4]
|
||||
src1_layout = TileLayout(S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)])
|
||||
dst1_shape = src1_shape
|
||||
dst1_layout = src1_layout
|
||||
reduce_dst_shape = [512]
|
||||
reduce_dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def tensor_scalar_reduce() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
|
||||
B_sbuf = T.alloc_buffer(dst1_shape, "float32", scope="trn.sbuf", layout=dst1_layout)
|
||||
C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout) # noqa: E501
|
||||
Tx.binary_reduce(B_sbuf, C_sbuf, A_sbuf, 1.0, "add", "sum", reduce_axes=(1, 2))
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "tensor_scalar_reduce"})
|
||||
T.device_entry()
|
||||
partial_reduce = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer((512, 1024, 4), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)]))
|
||||
B_sbuf = T.alloc_buffer((512, 1024, 4), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)]))
|
||||
C_sbuf = T.alloc_buffer((512,), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(128, 4) : (1 @ P, 1 @ F)]))
|
||||
Tx.binary_reduce(B_sbuf[0:512, 0:1024, 0:4], C_sbuf[0:512], A_sbuf[0:512, 0:1024, 0:4], T.float32(1.0), "add", "sum", [1, 2], workspace={"partial_reduce": partial_reduce}) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": tensor_scalar_reduce})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_activation_reduce_two_stage():
|
||||
A_shape = (32, 512, 128)
|
||||
A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
|
||||
B_shape = (16, 512, 128)
|
||||
B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
|
||||
C_shape = (1, 128)
|
||||
C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def activation_reduce():
|
||||
T.device_entry()
|
||||
A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
|
||||
B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
|
||||
C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
|
||||
for i in range(2):
|
||||
Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1))
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "activation_reduce"})
|
||||
T.device_entry()
|
||||
partial_reduce = T.alloc_buffer((128, 8), scope="trn.sbuf")
|
||||
const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
|
||||
A = T.alloc_buffer((32, 512, 128), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(16 * 1024, 128) : (1@F, 1@P)]))
|
||||
B = T.alloc_buffer((16, 512, 128), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(2, 4, 1024, 128) : (1024@F, 2048@F, 1@F, 1@P)]))
|
||||
C = T.alloc_buffer((1, 128), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(1, 128) : (1@F, 1@P)]))
|
||||
for i in range(2):
|
||||
Tx.unary_reduce(B[0:16, 0:512, 0:128], C[0, 0:128], A[i * 16:i * 16 + 16, 0:512, 0:128], "sqrt", "sum", None, None, [0, 1], workspace={"const_bias": const_bias, "partial_reduce": partial_reduce}) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": activation_reduce})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_partial_workspace_specify():
|
||||
A_shape = (32, 512, 128)
|
||||
A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
|
||||
B_shape = (16, 512, 128)
|
||||
B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
|
||||
C_shape = (1, 128)
|
||||
C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def activation_reduce():
|
||||
T.device_entry()
|
||||
partial_reduce = T.alloc_buffer((128, 16), scope="trn.sbuf")
|
||||
A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
|
||||
B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
|
||||
C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
|
||||
for i in range(2):
|
||||
Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1), workspace={"partial_reduce": partial_reduce}) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "activation_reduce"})
|
||||
T.device_entry()
|
||||
const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
|
||||
partial_reduce = T.alloc_buffer((128, 16), scope="trn.sbuf")
|
||||
A = T.alloc_buffer((32, 512, 128), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(16 * 1024, 128) : (1@F, 1@P)]))
|
||||
B = T.alloc_buffer((16, 512, 128), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(2, 4, 1024, 128) : (1024@F, 2048@F, 1@F, 1@P)]))
|
||||
C = T.alloc_buffer((1, 128), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(1, 128) : (1@F, 1@P)]))
|
||||
for i in range(2):
|
||||
Tx.unary_reduce(B[0:16, 0:512, 0:128], C[0, 0:128], A[i * 16:i * 16 + 16, 0:512, 0:128], "sqrt", "sum", None, None, [0, 1], workspace={"const_bias": const_bias, "partial_reduce": partial_reduce}) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": activation_reduce})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_workspace_reuse():
|
||||
src_shape = [512, 1024]
|
||||
src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
scale = T.float32(2.0)
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def unary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.exp(C_sbuf, A_sbuf, bias=0.0, scale=scale, max_inst_size=1024)
|
||||
Tx.exp(C_sbuf, C_sbuf)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "unary"})
|
||||
T.device_entry()
|
||||
const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
|
||||
A_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(128, 4096) : (1 @ P, 1 @ F)]))
|
||||
C_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
|
||||
layout=T.TileLayout(T.S[(128, 4096) : (1 @ P, 1 @ F)]))
|
||||
Tx.exp(C_sbuf[0:512, 0:1024], A_sbuf[0:512, 0:1024], T.float32(0.0), T.float32(2.0), workspace={"const_bias": const_bias}, max_inst_size=1024) # noqa: E501
|
||||
Tx.exp(C_sbuf[0:512, 0:1024], C_sbuf[0:512, 0:1024], None, None, workspace={"const_bias": const_bias}) # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": unary})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_no_rewrite_with_existing_workspace():
|
||||
src_shape = [128, 32, 4, 32]
|
||||
src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [128, 4]
|
||||
dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def reduction():
|
||||
T.device_entry()
|
||||
intermediate_buffer = T.alloc_buffer((128, 64), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.sum(B_sbuf, A_sbuf, axes=(1, 3), workspace={"partial_reduce": intermediate_buffer})
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": reduction})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], reduction)
|
||||
|
||||
|
||||
def test_no_rewrite_with_psum_output():
|
||||
A_layout = TileLayout(S[(128, 128) : (1 @ F, 1 @ P)])
|
||||
B_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)])
|
||||
|
||||
C_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def gemm() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=A_layout)
|
||||
B_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=B_layout)
|
||||
C_psum = T.alloc_buffer((128, 128), "float32", scope="trn.psum", layout=C_layout)
|
||||
Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": gemm})
|
||||
mod = TrnPrivateBufferAlloc()(mod)
|
||||
assert_structural_equal(mod["main"], gemm)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,283 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir import assert_structural_equal as _assert_structural_equal
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.tirx.layout import F, P, S, TileLayout
|
||||
from tvm.tirx.stmt_functor import ir_transform
|
||||
|
||||
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
|
||||
|
||||
|
||||
def _strip_exec_scope_stmt(stmt):
|
||||
def _postorder(node):
|
||||
if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
|
||||
return node.body
|
||||
return node
|
||||
|
||||
return ir_transform(
|
||||
stmt,
|
||||
preorder=lambda _node: None,
|
||||
postorder=_postorder,
|
||||
only_enable=["tirx.AttrStmt"],
|
||||
)
|
||||
|
||||
|
||||
def assert_structural_equal(lhs, rhs, *args, **kwargs):
|
||||
if isinstance(lhs, tvm.tirx.PrimFunc):
|
||||
lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
|
||||
if isinstance(rhs, tvm.tirx.PrimFunc):
|
||||
rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
|
||||
_assert_structural_equal(lhs, rhs, *args, **kwargs)
|
||||
|
||||
|
||||
opcode_map = {"sum": "add", "max": "max", "min": "min"}
|
||||
|
||||
Tx_func_map = {"sum": Tx.sum, "max": Tx.max, "min": Tx.min}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_type", ["sum", "max", "min"])
|
||||
def test_simple_reduction(op_type):
|
||||
src_shape = [128, 512]
|
||||
src_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [128, 1]
|
||||
dst_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
|
||||
|
||||
opcode = opcode_map[op_type]
|
||||
tx_func = Tx_func_map[op_type]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def reduction() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
tx_func(B_sbuf, A_sbuf, axes=-1)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "reduction"})
|
||||
A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 1), scope="trn.sbuf")
|
||||
for b_loop in range(1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorreduce(B_sbuf[p_loop, 0], A_sbuf[p_loop, f_loop], opcode, False, -1)
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": reduction})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_reduction_with_multiple_axes():
|
||||
src_shape = [128, 512, 4]
|
||||
src_layout = TileLayout(S[(128, 512, 4) : (1 @ P, 1 @ F, 512 @ F)])
|
||||
dst_shape = [128]
|
||||
dst_layout = TileLayout(S[128 : 1 @ P])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def reduction():
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.sum(B_sbuf, A_sbuf, axes=(1, 2), max_inst_size=2048)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "reduction"})
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 1), scope="trn.sbuf")
|
||||
for b_loop in range(1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorreduce(B_sbuf[p_loop, 0], A_sbuf[p_loop, f_loop], "add", False, -1)
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": reduction})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_reduction_in_loop():
|
||||
src_shape = [128, 512, 4]
|
||||
src_layout = TileLayout(S[(128, 512, 4) : (1 @ P, 4 @ F, 1 @ F)])
|
||||
dst_shape = [128, 4]
|
||||
dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def reduction():
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
Tx.sum(B_sbuf[:, i], A_sbuf[:, :, i], axes=-2)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "reduction"})
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(4, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorreduce(B_sbuf[p_loop, i], A_sbuf[p_loop, f_loop * 4 + i], "add", False, -1) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": reduction})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_reduction_two_stage():
|
||||
src_shape = [128, 32, 4, 32]
|
||||
src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [128, 4]
|
||||
dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def reduction():
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.sum(B_sbuf, A_sbuf, axes=(1, 3))
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "reduction"})
|
||||
intermediate_buffer = T.alloc_buffer((128, 32), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
for b_loop in range(4):
|
||||
for reduction_b_loop in range(32):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop], A_sbuf[p_loop, reduction_b_loop * 128 + b_loop * 32 + f_loop], "add", False, -1) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorreduce(B_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", False, -1) # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": reduction})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_reduction_with_guard():
|
||||
src_shape = [512, 2048]
|
||||
src_layout = TileLayout(S[(4, 128, 2048) : (2048 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = [512, 1]
|
||||
dst_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def reduction() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
for j in range(4):
|
||||
Tx.sum(B_sbuf[0: (i+1) * 128, 0], A_sbuf[0: (i+1) * 128, 0: (j+1) * 256], max_inst_size=512) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "reduction"})
|
||||
intermediate_buffer = T.alloc_buffer((128, 2), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer((128, 8192), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
for i, j in T.grid(4, 4):
|
||||
for b_loop in range(4):
|
||||
for reduction_b_loop in range(2):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
|
||||
if (
|
||||
b_loop - i < 1
|
||||
and reduction_b_loop * 512 + f_loop < j * 256 + 256
|
||||
):
|
||||
T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop], A_sbuf[p_loop, b_loop * 2048 + reduction_b_loop * 512 + f_loop], "add", T.bool(False), -1) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(2, annotations={"nki_dim": "F"}):
|
||||
if b_loop - i < 1 and f_loop * 2 - j < 1:
|
||||
T.nki.tensorreduce(B_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": reduction})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_reduction_two_stage_workspace():
|
||||
src_shape = [128, 32, 4, 32]
|
||||
src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [128, 4]
|
||||
dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def reduction():
|
||||
T.device_entry()
|
||||
intermediate_buffer = T.alloc_buffer((128, 64), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.sum(B_sbuf, A_sbuf, axes=(1, 3), workspace={"partial_reduce": intermediate_buffer})
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "reduction"})
|
||||
intermediate_buffer = T.alloc_buffer((128, 64), scope="trn.sbuf")
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
for b_loop in range(4):
|
||||
for reduction_b_loop in range(32):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop], A_sbuf[p_loop, reduction_b_loop * 128 + b_loop * 32 + f_loop], "add", False, -1) # noqa: E501
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
|
||||
T.nki.tensorreduce(B_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", False, -1) # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": reduction})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,186 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir import assert_structural_equal as _assert_structural_equal
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.tirx.layout import F, P, S, TileLayout
|
||||
from tvm.tirx.stmt_functor import ir_transform
|
||||
|
||||
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
|
||||
|
||||
|
||||
def _strip_exec_scope_stmt(stmt):
|
||||
def _postorder(node):
|
||||
if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
|
||||
return node.body
|
||||
return node
|
||||
|
||||
return ir_transform(
|
||||
stmt,
|
||||
preorder=lambda _node: None,
|
||||
postorder=_postorder,
|
||||
only_enable=["tirx.AttrStmt"],
|
||||
)
|
||||
|
||||
|
||||
def assert_structural_equal(lhs, rhs, *args, **kwargs):
|
||||
if isinstance(lhs, tvm.tirx.PrimFunc):
|
||||
lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
|
||||
if isinstance(rhs, tvm.tirx.PrimFunc):
|
||||
rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
|
||||
_assert_structural_equal(lhs, rhs, *args, **kwargs)
|
||||
|
||||
|
||||
def test_select():
|
||||
src_shape = [128, 512]
|
||||
src_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [128, 512]
|
||||
dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def select() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.select(B_sbuf, A_sbuf, 0.0, lambda i, j: i < j)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "select"})
|
||||
A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.affine_select(B_sbuf[p_loop, f_loop], p_loop < f_loop, A_sbuf[p_loop, f_loop], T.float32(0.0)) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": select})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_select_in_loop():
|
||||
src_shape = [32, 128, 512]
|
||||
src_layout = TileLayout(S[(32, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = [128, 512]
|
||||
dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def select() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(2):
|
||||
Tx.select(B_sbuf, A_sbuf[i*16, :, :], 0.0, lambda a, b: (i+1)* a < b)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "select"})
|
||||
A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
for i, b_loop in T.grid(2, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.affine_select(B_sbuf[p_loop, f_loop], (i + 1) * p_loop < f_loop, A_sbuf[p_loop, i * 8192 + f_loop], T.float32(0.0)) # noqa: E501
|
||||
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": select})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_select_expr_affine():
|
||||
src_shape = [512, 512]
|
||||
src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def select() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.select(B_sbuf, A_sbuf, 0.0, lambda i, j: i < j)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "select"})
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.affine_select(B_sbuf[p_loop, b_loop * 512 + f_loop], b_loop * 128 + p_loop < f_loop, A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0)) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": select})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_select_with_guard():
|
||||
src_shape = [512, 512]
|
||||
src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def select() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
for j in range(4):
|
||||
Tx.select(B_sbuf[0: (i+1) * 128, 0: (j+1) * 128], A_sbuf[0: (i+1) * 128, 0: (j+1) * 128], 0.0, lambda a, b: a < b) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "select"})
|
||||
A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
for i, j, b_loop in T.grid(4, 4, 4):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
if b_loop - i < 1 and f_loop < j * 128 + 128:
|
||||
T.nki.affine_select(B_sbuf[p_loop, b_loop * 512 + f_loop], b_loop * 128 + p_loop < f_loop, A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0)) # noqa: E501
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": select})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,288 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir import assert_structural_equal as _assert_structural_equal
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.tirx.layout import F, P, S, TileLayout
|
||||
from tvm.tirx.stmt_functor import ir_transform
|
||||
|
||||
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
|
||||
|
||||
|
||||
def _strip_exec_scope_stmt(stmt):
|
||||
def _postorder(node):
|
||||
if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
|
||||
return node.body
|
||||
return node
|
||||
|
||||
return ir_transform(
|
||||
stmt,
|
||||
preorder=lambda _node: None,
|
||||
postorder=_postorder,
|
||||
only_enable=["tirx.AttrStmt"],
|
||||
)
|
||||
|
||||
|
||||
def assert_structural_equal(lhs, rhs, *args, **kwargs):
|
||||
if isinstance(lhs, tvm.tirx.PrimFunc):
|
||||
lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
|
||||
if isinstance(rhs, tvm.tirx.PrimFunc):
|
||||
rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
|
||||
_assert_structural_equal(lhs, rhs, *args, **kwargs)
|
||||
|
||||
|
||||
Tx_func_map = {"reciprocal": Tx.reciprocal, "sqrt": Tx.sqrt, "memset": Tx.memset, "exp": Tx.exp}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_type", ["reciprocal", "memset"])
|
||||
def test_simple_unary(op_type):
|
||||
src_shape = [128, 512]
|
||||
src_layout = T.TileLayout(T.S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [128, 512]
|
||||
dst_layout = T.TileLayout(T.S[(128, 512) : (1 @ P, 1 @ F)])
|
||||
tx_func = Tx_func_map[op_type]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def unary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
if op_type == "memset":
|
||||
tx_func(B_sbuf, T.float32(0.0))
|
||||
else:
|
||||
tx_func(B_sbuf, A_sbuf)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "unary"})
|
||||
A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
if op_type == "reciprocal":
|
||||
T.nki.reciprocal(
|
||||
B_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop]
|
||||
)
|
||||
elif op_type == "memset":
|
||||
T.nki.memset(B_sbuf[p_loop, f_loop], 0.0)
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": unary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_type", ["reciprocal", "memset"])
|
||||
def test_unary_in_a_loop(op_type):
|
||||
src_shape = [1024, 512]
|
||||
src_layout = T.TileLayout(T.S[(128, 4096) : (1 @ P, 1 @ F)])
|
||||
dst_shape = [512, 512]
|
||||
dst_layout = T.TileLayout(T.S[(128, 2048) : (1 @ P, 1 @ F)])
|
||||
|
||||
Tx_func = Tx_func_map[op_type]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def unary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
A_sbuf_view = A_sbuf.view(128, 8, 512)
|
||||
B_sbuf_view = B_sbuf.view(128, 4, 512)
|
||||
for i in range(4):
|
||||
if op_type == "memset":
|
||||
Tx_func(B_sbuf_view[:, i, :], T.float32(0.0))
|
||||
else:
|
||||
Tx_func(B_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :])
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "unary"})
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
|
||||
A_sbuf_view = T.decl_buffer((128, 4096), data=A_sbuf.data, scope="trn.sbuf", layout=None)
|
||||
B_sbuf_view = T.decl_buffer((128, 2048), data=B_sbuf.data, scope="trn.sbuf", layout=None)
|
||||
for i, b_loop in T.grid(4, 1):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
if op_type == "reciprocal":
|
||||
T.nki.reciprocal(B_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop]) # noqa: E501
|
||||
elif op_type == "memset":
|
||||
T.nki.memset(B_sbuf[p_loop, i * 512 + f_loop], 0.0)
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": unary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_unary_complex1():
|
||||
dst_layout = TileLayout(S[(32, 128, 256) : (256 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = [4096, 256]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def unary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
Tx.memset(A_sbuf, T.float32(0.0))
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "unary"})
|
||||
A_sbuf = T.alloc_buffer((128, 8192), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 16):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.memset(A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0))
|
||||
# fmt: on
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": unary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_type", ["sqrt", "exp"])
|
||||
def test_unary_with_bias_scale(op_type):
|
||||
src_shape = [512, 1024]
|
||||
src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
bias_shape = [512, 1]
|
||||
bias_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
|
||||
scale = T.float32(2.0)
|
||||
tx_func = Tx_func_map[op_type]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def unary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(bias_shape, "float32", scope="trn.sbuf", layout=bias_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
tx_func(C_sbuf, A_sbuf, bias=B_sbuf, scale=scale)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "unary"})
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 8):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
T.nki.activation(C_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], op_type, B_sbuf[p_loop, b_loop//2], T.float32(2.0)) # noqa: E501
|
||||
# fmt: off
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": unary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op_type", ["sqrt", "exp"])
|
||||
def test_unary_with_bias_scale_2(op_type):
|
||||
src_shape = [512, 1024]
|
||||
src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
bias = T.float32(1.0)
|
||||
scale = T.float32(2.0)
|
||||
tx_func = Tx_func_map[op_type]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def unary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
tx_func(C_sbuf, A_sbuf, bias=bias, scale=scale)
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "unary"})
|
||||
const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
|
||||
T.nki.memset(const_bias[p_loop, f_loop], T.float32(1.0))
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
for b_loop in T.serial(0, 8):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
|
||||
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
|
||||
T.nki.activation(C_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], op_type, const_bias[p_loop, f_loop], T.float32(2.0)) # noqa: E501
|
||||
# fmt: off
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": unary})
|
||||
mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_unary_with_guard():
|
||||
src_shape = [512, 1024]
|
||||
src_layout = TileLayout(S[(4, 128, 1024) : (1024 @ F, 1 @ P, 1 @ F)])
|
||||
dst_shape = src_shape
|
||||
dst_layout = src_layout
|
||||
bias_shape = [512, 1]
|
||||
bias_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
|
||||
scale = T.float32(2.0)
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def unary() -> None:
|
||||
T.device_entry()
|
||||
A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
|
||||
B_sbuf = T.alloc_buffer(bias_shape, "float32", scope="trn.sbuf", layout=bias_layout)
|
||||
C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
|
||||
for i in range(4):
|
||||
for j in range(4):
|
||||
Tx.sqrt(C_sbuf[0: (i+1) * 128, 0: (j+1)*256], A_sbuf[0: (i+1) * 128, 0: (j+1)*256], bias=B_sbuf[0: (i+1) * 128, 0], scale=scale) # noqa: E501
|
||||
|
||||
@T.prim_func
|
||||
def expected():
|
||||
T.func_attr({"global_symbol": "unary"})
|
||||
A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
|
||||
C_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
|
||||
for i, j, b_loop in T.grid(4, 4, 8):
|
||||
T.attr(0, "tensorized_nki_instruction", 1)
|
||||
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
|
||||
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
|
||||
if b_loop // 2 - i < 1 and b_loop % 2 * 512 + f_loop < j * 256 + 256:
|
||||
T.nki.activation(C_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], "sqrt", B_sbuf[p_loop, b_loop // 2], T.float32(2.0)) # noqa: E501
|
||||
# fmt: off
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": unary})
|
||||
mod = tvm.tirx.transform.LowerTIRx()(mod)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user