Bohan Hou
859498dc01
## Summary This PR adds the initial TIRx support needed for low-level programming of Blackwell-class GPU architectures. As part of the ongoing TIRx refactor, it introduces TVMScript support for directly scripting advanced hardware features without relying on scheduling as the primary programming interface. The change keeps existing `s_tir` script support intact while making direct scripting a first-class path for TIRx programs. ## Main Changes - Add TIRx operator dispatch and layout infrastructure. - Add TVMScript support for new low-level TIRx operations. - Add analysis, transform, and lowering support for TIRx IR nodes. - Add CUDA/Blackwell-oriented codegen and intrinsic coverage. - Add Python and C++ integration points for TIRx scripting and runtime support. ## Validation - `pre-commit run --all-files` - `ninja -C build -j32` - `CUDA_VISIBLE_DEVICES=2 pytest tests/python/tirx/ -n 16` - `1723 passed, 47 skipped, 32 warnings` - `CUDA_VISIBLE_DEVICES=2 python -m pytest -v tests/python/all-platform-minimal-test` - `37 passed, 105 skipped` - `TVM_TEST_TARGETS=llvm python -m pytest -v tests/python/tirx-analysis tests/python/tirx-base tests/python/tirx-transform -n 16` - `664 passed, 25 skipped, 9 xfailed, 1 xpassed` ## Local CI Notes Some full CI-equivalent jobs were not locally reproducible because this machine is missing parts of the Apache TVM CI environment, including `llvm-config-15/17`, Vulkan, ROCm, Maven, Sphinx, Doxygen, Emscripten, and ARM/QEMU cross-toolchain components. Metal-specific tests were skipped locally because no Metal runtime is available.
429 lines
16 KiB
Python
429 lines
16 KiB
Python
# 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.
|
|
"""Unit tests for ExecContext (RFC v3 §6). Cases mirror RFC §8.1 -- §8.10."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from tvm.tirx.exec_context import (
|
|
CLUSTER,
|
|
CTA,
|
|
LANE_CTA_THREAD,
|
|
LANE_FLAT,
|
|
LANE_W_INNER,
|
|
LANE_WG_OUTER,
|
|
LANE_WG_THREAD,
|
|
THREAD,
|
|
WARP,
|
|
WARPGROUP,
|
|
AxisRange,
|
|
ExecContext,
|
|
ExecContextError,
|
|
LaneBinding,
|
|
filter_modulo,
|
|
filter_narrow,
|
|
initial_A,
|
|
scope_switch,
|
|
)
|
|
|
|
# -- canonical bindings declared at kernel entry (see RFC §8 naming conv) --
|
|
WARP_FLAT = LaneBinding(axis="warpid", kind=LANE_FLAT, declared_extent=16)
|
|
WG_OUTER = LaneBinding(axis="warpid", kind=LANE_WG_OUTER, declared_extent=4)
|
|
W_INNER = LaneBinding(axis="warpid", kind=LANE_W_INNER, declared_extent=4)
|
|
LANE_BIND = LaneBinding(axis="laneid", kind=LANE_FLAT, declared_extent=32)
|
|
CTA_BIND = LaneBinding(axis="cta_id", kind=LANE_FLAT, declared_extent=1)
|
|
CTA_THREAD_BIND = LaneBinding(axis="thread", kind=LANE_CTA_THREAD, declared_extent=256)
|
|
WG_THREAD_BIND = LaneBinding(axis="thread", kind=LANE_WG_THREAD, declared_extent=128)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §3 scope_switch: split table
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_initial_A_single_cta():
|
|
A = initial_A(warp_ext=16)
|
|
assert A.laneid == AxisRange(32, 0)
|
|
assert A.warpid == AxisRange(16, 0)
|
|
assert A.cta_id == AxisRange(1, 0)
|
|
assert A.size == 512
|
|
|
|
|
|
def test_initial_A_cluster():
|
|
A = initial_A(warp_ext=16, cta_ext=4)
|
|
assert A.cta_id == AxisRange(4, 0)
|
|
assert A.size == 2048
|
|
|
|
|
|
def test_axis_modulo_filter_uses_stride():
|
|
A = initial_A(warp_ext=16, cta_ext=4)
|
|
A = filter_modulo(A, "cta_id", 2, 0)
|
|
assert A.cta_id == AxisRange(2, 0, 2)
|
|
A = filter_narrow(A, CTA_BIND, 1, 4)
|
|
assert A.cta_id == AxisRange(1, 2, 2)
|
|
|
|
|
|
def test_axis_modulo_filter_two_cta_pair_residues():
|
|
A = initial_A(warp_ext=16, cta_ext=2)
|
|
assert filter_modulo(A, "cta_id", 2, 0).cta_id == AxisRange(1, 0, 2)
|
|
assert filter_modulo(A, "cta_id", 2, 1).cta_id == AxisRange(1, 1, 2)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"kappa,expected_inter_axes,expected_intra_axes",
|
|
[
|
|
(THREAD, {"laneid", "warpid", "cta_id"}, set()),
|
|
(WARP, {"warpid", "cta_id"}, {"laneid"}),
|
|
(CTA, {"cta_id"}, {"laneid", "warpid"}),
|
|
(CLUSTER, set(), {"laneid", "warpid", "cta_id"}),
|
|
],
|
|
)
|
|
def test_scope_switch_trivial(kappa, expected_inter_axes, expected_intra_axes):
|
|
A = initial_A(warp_ext=16, cta_ext=4)
|
|
split = scope_switch(A, kappa)
|
|
assert set(split.inter) == expected_inter_axes
|
|
assert set(split.intra) == expected_intra_axes
|
|
|
|
|
|
def test_scope_switch_warpgroup_aligned():
|
|
A = initial_A(warp_ext=16)
|
|
split = scope_switch(A, WARPGROUP)
|
|
assert split.inter["wgid"] == AxisRange(4, 0)
|
|
assert split.inter["cta_id"] == AxisRange(1, 0)
|
|
assert split.intra["laneid"] == AxisRange(32, 0)
|
|
assert split.intra["wid_in_wg"] == AxisRange(4, 0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §4.2 warpgroup factoring: 3 cases
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_factor_case1_aligned():
|
|
A = initial_A(warp_ext=8) # ext=8, off=0 -- aligned
|
|
split = scope_switch(A, WARPGROUP)
|
|
assert split.inter["wgid"] == AxisRange(2, 0)
|
|
assert split.intra["wid_in_wg"] == AxisRange(4, 0)
|
|
|
|
|
|
def test_factor_case2_fits_in_one_wg():
|
|
# warpid ext=2, off=0 -- fits in one wg
|
|
A = initial_A(warp_ext=16)
|
|
A = filter_narrow(A, WARP_FLAT, 0, 2)
|
|
split = scope_switch(A, WARPGROUP)
|
|
assert split.inter["wgid"] == AxisRange(1, 0)
|
|
assert split.intra["wid_in_wg"] == AxisRange(2, 0)
|
|
|
|
|
|
def test_factor_case2_offset():
|
|
# warpid ext=2, off=6 -> wid_off=2, fits (2 <= 4-2)
|
|
A = initial_A(warp_ext=16)
|
|
A = filter_narrow(A, WARP_FLAT, 6, 8)
|
|
split = scope_switch(A, WARPGROUP)
|
|
assert split.inter["wgid"] == AxisRange(1, 1)
|
|
assert split.intra["wid_in_wg"] == AxisRange(2, 2)
|
|
|
|
|
|
def test_factor_case3_fails():
|
|
# RFC §8.6: warpid[2:6] crosses wg boundary unaligned
|
|
A = initial_A(warp_ext=16)
|
|
A = filter_narrow(A, WARP_FLAT, 2, 6)
|
|
assert A.warpid == AxisRange(4, 2)
|
|
with pytest.raises(ExecContextError, match="crosses warpgroup boundary"):
|
|
scope_switch(A, WARPGROUP)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §8.1 -- Pure narrowing CTA -> WG -> W
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ex_8_1_cta_wg_warp():
|
|
ctx = ExecContext.at_kernel_entry(warp_ext=16)
|
|
# with T.cta()
|
|
ctx = ctx.with_scope_switch(CTA)
|
|
assert ctx.inter == {"cta_id": AxisRange(1, 0)}
|
|
assert ctx.intra == {"laneid": AxisRange(32, 0), "warpid": AxisRange(16, 0)}
|
|
# with T.warpgroup()
|
|
ctx = ctx.with_scope_switch(WARPGROUP)
|
|
assert ctx.inter == {"wgid": AxisRange(4, 0), "cta_id": AxisRange(1, 0)}
|
|
assert ctx.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)}
|
|
# with T.warp()
|
|
ctx = ctx.with_scope_switch(WARP)
|
|
assert ctx.inter == {"warpid": AxisRange(16, 0), "cta_id": AxisRange(1, 0)}
|
|
assert ctx.intra == {"laneid": AxisRange(32, 0)}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §8.2 -- Filter + scope_switch
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ex_8_2_filter_then_warpgroup():
|
|
ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
|
|
ctx = ctx.with_filter(WARP_FLAT, 0, 8)
|
|
assert ctx.A.warpid == AxisRange(8, 0)
|
|
# recompute at cta: intra=(lane:32, warp:8)
|
|
assert ctx.intra == {"laneid": AxisRange(32, 0), "warpid": AxisRange(8, 0)}
|
|
# enter warpgroup: factor(8, 0) -> case 1
|
|
ctx = ctx.with_scope_switch(WARPGROUP)
|
|
assert ctx.inter == {"wgid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)}
|
|
assert ctx.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §8.3 -- Sugar form T.warp(warpid[2:4])
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ex_8_3_sugar_warp_range():
|
|
ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
|
|
# desugar: filter warpid[2:4], then warp
|
|
ctx = ctx.with_filter(WARP_FLAT, 2, 4).with_scope_switch(WARP)
|
|
assert ctx.A.warpid == AxisRange(2, 2)
|
|
assert ctx.inter == {"warpid": AxisRange(2, 2), "cta_id": AxisRange(1, 0)}
|
|
assert ctx.intra == {"laneid": AxisRange(32, 0)}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §8.4 -- Widen after filter (warp -> warpgroup)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ex_8_4_widen_warp_to_wg():
|
|
ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
|
|
ctx = ctx.with_filter(WARP_FLAT, 0, 4).with_scope_switch(WARP)
|
|
# widen to warpgroup
|
|
ctx = ctx.with_scope_switch(WARPGROUP)
|
|
assert ctx.inter == {"wgid": AxisRange(1, 0), "cta_id": AxisRange(1, 0)}
|
|
assert ctx.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §8.5 -- Partial warp selection -> warpgroup (partial intra)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ex_8_5_partial_wg():
|
|
ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
|
|
ctx = ctx.with_filter(WARP_FLAT, 0, 2).with_scope_switch(WARPGROUP)
|
|
# case 2: 2 <= 4-0
|
|
assert ctx.inter == {"wgid": AxisRange(1, 0), "cta_id": AxisRange(1, 0)}
|
|
assert ctx.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(2, 0)}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §8.6 -- Cross warpgroup boundary (factor fails)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ex_8_6_factor_fail():
|
|
ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
|
|
# with_filter recomputes (inter, intra) for current scope_kind=cta -- still OK
|
|
ctx2 = ctx.with_filter(WARP_FLAT, 2, 6)
|
|
assert ctx2.A.warpid == AxisRange(4, 2)
|
|
# scope_switch to warpgroup is the one that must fail
|
|
with pytest.raises(ExecContextError, match="crosses warpgroup boundary"):
|
|
ctx2.with_scope_switch(WARPGROUP)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §8.7 -- Deep mixed nesting
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ex_8_7_deep_nested():
|
|
ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
|
|
ctx = ctx.with_filter(WARP_FLAT, 0, 8).with_scope_switch(WARPGROUP)
|
|
assert ctx.inter == {"wgid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)}
|
|
ctx = ctx.with_filter(WARP_FLAT, 0, 2)
|
|
# recompute at warpgroup: factor(2, 0) -> case 2
|
|
assert ctx.inter == {"wgid": AxisRange(1, 0), "cta_id": AxisRange(1, 0)}
|
|
assert ctx.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(2, 0)}
|
|
ctx = ctx.with_scope_switch(WARP)
|
|
assert ctx.inter == {"warpid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)}
|
|
assert ctx.intra == {"laneid": AxisRange(32, 0)}
|
|
ctx = ctx.with_filter(LANE_BIND, 0, 8)
|
|
assert ctx.intra == {"laneid": AxisRange(8, 0)}
|
|
assert ctx.inter == {"warpid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §8.8 -- FA4 pattern: 3 sibling filter branches
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ex_8_8_fa4_pattern():
|
|
root = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
|
|
|
|
# Branch 1: warp 12 (single warp, tcgen05 MMA elected)
|
|
b1 = root.with_filter(WARP_FLAT, 12, 13)
|
|
assert b1.A.warpid == AxisRange(1, 12)
|
|
assert b1.intra == {"laneid": AxisRange(32, 0), "warpid": AxisRange(1, 12)}
|
|
|
|
# Branch 2: softmax warpgroups (warps 0-7)
|
|
b2 = root.with_filter(WARP_FLAT, 0, 8).with_scope_switch(WARPGROUP)
|
|
assert b2.inter == {"wgid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)}
|
|
assert b2.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)}
|
|
|
|
# Branch 3: correction warpgroup (warps 8-11 = wg2)
|
|
b3 = root.with_filter(WARP_FLAT, 8, 12)
|
|
assert b3.A.warpid == AxisRange(4, 8)
|
|
assert b3.intra == {"laneid": AxisRange(32, 0), "warpid": AxisRange(4, 8)}
|
|
# And should factor cleanly when entering warpgroup
|
|
b3wg = b3.with_scope_switch(WARPGROUP)
|
|
assert b3wg.inter == {"wgid": AxisRange(1, 2), "cta_id": AxisRange(1, 0)}
|
|
assert b3wg.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §8.9 -- Cross-CTA with widening to cluster
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ex_8_9_cross_cta_cluster():
|
|
ctx = ExecContext.at_kernel_entry(warp_ext=16, cta_ext=4).with_scope_switch(CTA)
|
|
assert ctx.inter == {"cta_id": AxisRange(4, 0)}
|
|
# filter to warp 0, then warp
|
|
w = ctx.with_filter(WARP_FLAT, 0, 1).with_scope_switch(WARP)
|
|
assert w.inter == {"warpid": AxisRange(1, 0), "cta_id": AxisRange(4, 0)}
|
|
assert w.intra == {"laneid": AxisRange(32, 0)}
|
|
|
|
# back at cta scope, enter warpgroup
|
|
wg = ctx.with_scope_switch(WARPGROUP)
|
|
assert wg.inter == {"wgid": AxisRange(4, 0), "cta_id": AxisRange(4, 0)}
|
|
# widen to cluster
|
|
cl = wg.with_scope_switch(CLUSTER)
|
|
assert cl.inter == {}
|
|
assert cl.intra == {
|
|
"laneid": AxisRange(32, 0),
|
|
"warpid": AxisRange(16, 0),
|
|
"cta_id": AxisRange(4, 0),
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# §8.10 -- identical to 8.3 modulo prose; covered above
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Rule 1 & 5: filter can only shrink A; saved/restored across scope exit
|
|
# (Restoration is the caller's (IR walker's) responsibility -- ExecContext
|
|
# is immutable, each with_filter returns a fresh ctx. Test that the parent
|
|
# is untouched.)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_filter_is_pure():
|
|
ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
|
|
child = ctx.with_filter(WARP_FLAT, 0, 8)
|
|
assert ctx.A.warpid == AxisRange(16, 0) # parent not mutated
|
|
assert child.A.warpid == AxisRange(8, 0)
|
|
|
|
|
|
def test_filter_empty_range_rejected():
|
|
A = initial_A(warp_ext=16)
|
|
with pytest.raises(ExecContextError, match="empty or inverted"):
|
|
filter_narrow(A, WARP_FLAT, 5, 5)
|
|
|
|
|
|
def test_filter_out_of_range_rejected():
|
|
A = initial_A(warp_ext=16)
|
|
A = filter_narrow(A, WARP_FLAT, 0, 4)
|
|
with pytest.raises(ExecContextError, match="empty range"):
|
|
filter_narrow(A, WARP_FLAT, 8, 12) # disjoint from [0, 4)
|
|
|
|
|
|
def test_filter_flat_cta_thread_full_warp_range():
|
|
A = initial_A(warp_ext=8)
|
|
A = filter_narrow(A, CTA_THREAD_BIND, 0, 128)
|
|
assert A.warpid == AxisRange(4, 0)
|
|
assert A.laneid == AxisRange(32, 0)
|
|
|
|
|
|
def test_filter_flat_cta_thread_single_warp_lane_range():
|
|
A = initial_A(warp_ext=8)
|
|
A = filter_narrow(A, CTA_THREAD_BIND, 34, 40)
|
|
assert A.warpid == AxisRange(1, 1)
|
|
assert A.laneid == AxisRange(6, 2)
|
|
|
|
|
|
def test_filter_flat_cta_thread_nonrectangular_rejected():
|
|
A = initial_A(warp_ext=8)
|
|
with pytest.raises(ExecContextError, match="non-rectangular"):
|
|
filter_narrow(A, CTA_THREAD_BIND, 20, 50)
|
|
|
|
|
|
def test_filter_flat_warpgroup_thread_range_inside_one_warpgroup():
|
|
A = initial_A(warp_ext=8)
|
|
A = filter_narrow(A, WG_OUTER, 1, 2)
|
|
A = filter_narrow(A, WG_THREAD_BIND, 32, 64)
|
|
assert A.warpid == AxisRange(1, 5)
|
|
assert A.laneid == AxisRange(32, 0)
|
|
|
|
|
|
def test_filter_flat_warpgroup_thread_full_range_across_warpgroups_is_noop():
|
|
A = initial_A(warp_ext=8)
|
|
A2 = filter_narrow(A, WG_THREAD_BIND, 0, 128)
|
|
assert A2.warpid == AxisRange(8, 0)
|
|
assert A2.laneid == AxisRange(32, 0)
|
|
|
|
|
|
def test_filter_flat_warpgroup_thread_partial_range_across_warpgroups_rejected():
|
|
A = initial_A(warp_ext=8)
|
|
with pytest.raises(ExecContextError, match="multiple warpgroups"):
|
|
filter_narrow(A, WG_THREAD_BIND, 0, 64)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Factor-lane bindings: wg_outer and w_inner
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_filter_wg_outer():
|
|
A = initial_A(warp_ext=16)
|
|
A2 = filter_narrow(A, WG_OUTER, 1, 3) # wg 1..2 -> warps 4..11
|
|
assert A2.warpid == AxisRange(8, 4)
|
|
|
|
|
|
def test_filter_wg_outer_unaligned_rejected():
|
|
A = initial_A(warp_ext=16)
|
|
A = filter_narrow(A, WARP_FLAT, 2, 6) # warp offset 2 (not WG-aligned)
|
|
with pytest.raises(ExecContextError, match="aligned to WG_SIZE"):
|
|
filter_narrow(A, WG_OUTER, 0, 1)
|
|
|
|
|
|
def test_filter_w_inner():
|
|
A = initial_A(warp_ext=16)
|
|
# First narrow into a single warpgroup, then inner filter is valid
|
|
A = filter_narrow(A, WARP_FLAT, 4, 8) # wg1: warps 4..7
|
|
A2 = filter_narrow(A, W_INNER, 1, 3) # pick inner lanes 1..2
|
|
assert A2.warpid == AxisRange(2, 5)
|
|
|
|
|
|
def test_filter_w_inner_spanning_wg_rejected():
|
|
A = initial_A(warp_ext=16) # spans all 4 wgs
|
|
with pytest.raises(ExecContextError, match="spans multiple warpgroups"):
|
|
filter_narrow(A, W_INNER, 0, 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main([__file__, "-v"]))
|