# 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"]))