diff --git a/tests/end_to_end/test_gluon.py b/tests/end_to_end/test_gluon.py index 76409c2c..a1f40f1b 100644 --- a/tests/end_to_end/test_gluon.py +++ b/tests/end_to_end/test_gluon.py @@ -6,7 +6,7 @@ from triton import knobs from triton.experimental import gluon from triton.experimental.gluon import language as gl -from triton.experimental.gluon.language.nvidia import hopper +from triton.experimental.gluon.language.nvidia import blackwell, hopper from triton.experimental.gluon.language.nvidia.blackwell import tma as blackwell_tma from triton.experimental.gluon.language.nvidia.hopper import ( mbarrier, @@ -15,10 +15,9 @@ from triton.experimental.gluon.nvidia.hopper import TensorDescriptor import triton_viz -from triton_viz.core.callbacks import ForLoopCallbacks, OpCallbacks -from triton_viz.core.client import Client from triton_viz.clients.sanitizer.sanitizer import SymbolicSanitizer from triton_viz.core.data import Load +from triton_viz.core.simulation.gluon import GluonInterpretedFunction, gluon_builder try: from triton.experimental.gluon.language.amd.cdna4 import ( @@ -47,38 +46,40 @@ _HAS_TMA_IM2COL = TensorDescriptorIm2Col is not None -class _NoOpClient(Client): - NAME = "noop" - - def pre_run_callback(self, fn): - return True - - def post_run_callback(self, fn): - return True - - def arg_callback(self, name, arg, arg_cvt): - pass - - def grid_callback(self, grid): - pass - - def grid_idx_callback(self, grid_idx): - pass - - def register_op_callback(self, op_type, *args, **kwargs): - return OpCallbacks() - - def register_for_loop_callback(self): - return ForLoopCallbacks() - - def finalize(self): - return [] - - def pre_warmup_callback(self, jit_fn, *args, **kwargs): - return True - - def post_warmup_callback(self, jit_fn, ret): - pass +def _run_gluon_on_cpu(fn, grid, *args, **kwargs): + GluonInterpretedFunction(fn.fn).run(*args, grid=grid, **kwargs) + + +# CI can run against Gluon versions whose TensorMemoryLayout._to_ir uses older +# builder argument order; keep this CPU accuracy test on the current builder shape. +class _TensorMemoryLayoutForCpuRoundtrip(blackwell.TensorMemoryLayout): + def __init__(self, block, col_stride, cga_layout=None, two_ctas=False): + object.__setattr__(self, "block", tuple(block)) + object.__setattr__(self, "col_stride", int(col_stride)) + object.__setattr__( + self, + "cga_layout", + [] if cga_layout is None else [list(basis) for basis in cga_layout], + ) + object.__setattr__(self, "two_ctas", bool(two_ctas)) + object.__setattr__(self, "fp4_padded", False) + object.__setattr__(self, "cta_split_num", None) + + def _to_ir(self, builder): + return builder.get_tensor_memory_layout( + self.block, + self.col_stride, + self.cga_layout, + self.two_ctas, + False, + ) + + def mangle(self): + cga_layout_str = "_".join( + "~".join(map(str, basis)) for basis in self.cga_layout + ) + two_ctas_str = "2CT" if self.two_ctas else "" + return f"TL{self.block[0]}x{self.block[1]}C{self.col_stride}{cga_layout_str}{two_ctas_str}TL" @gluon.jit @@ -518,9 +519,9 @@ def _run_im2col_case( pixel_box_upper_corner=pixel_box_upper_corner, ) out_desc = TensorDescriptor.from_tensor(out, [16, 32], layout) - kernel = triton_viz.trace(_NoOpClient(), frontend="gluon")(_tma_im2col_kernel) - - kernel[(1,)](in_desc, out_desc, *coord, *offsets, num_warps=1) + _run_gluon_on_cpu( + _tma_im2col_kernel, (1,), in_desc, out_desc, *coord, *offsets, num_warps=1 + ) return out @@ -614,13 +615,53 @@ def _run_small_wgmma_case(seed: int, lhs_in_reg: bool): b_desc = TensorDescriptor.from_tensor(b, list(b.shape), b_layout) c_desc = TensorDescriptor.from_tensor(c, list(c.shape), c_layout) out_desc = TensorDescriptor.from_tensor(out, list(out.shape), c_layout) - kernel = triton_viz.trace(_NoOpClient(), frontend="gluon")(_small_wgmma_kernel) - - kernel[(1,)](a_desc, b_desc, c_desc, out_desc, lhs_in_reg, num_warps=4) + _run_gluon_on_cpu( + _small_wgmma_kernel, + (1,), + a_desc, + b_desc, + c_desc, + out_desc, + lhs_in_reg, + num_warps=4, + ) torch.testing.assert_close(out, a.float() @ b.float() + c, atol=1e-2, rtol=1e-2) +@gluon.jit +def _tensor_memory_roundtrip_kernel( + in_ptr, + out_ptr, + M: gl.constexpr, + N: gl.constexpr, + num_warps: gl.constexpr, +): + global_layout: gl.constexpr = gl.BlockedLayout( + [1, 1], + [1, 32], + [1, num_warps], + [1, 0], + ) + offs_m = gl.arange(0, M, gl.SliceLayout(1, global_layout)) + offs_n = gl.arange(0, N, gl.SliceLayout(0, global_layout)) + offsets = offs_m[:, None] * N + offs_n[None, :] + value = gl.load(in_ptr + offsets) + tmem_layout: gl.constexpr = _TensorMemoryLayoutForCpuRoundtrip( + block=(64, 64), + col_stride=32 // in_ptr.dtype.element_ty.primitive_bitwidth, + cga_layout=[[1, 0], [0, 1]], + ) + tmem = blackwell.allocate_tensor_memory( + element_ty=in_ptr.dtype.element_ty, + shape=[M, N], + layout=tmem_layout, + ) + tmem.store(value) + out = tmem.load(global_layout) + gl.store(out_ptr + offsets, out) + + def test_gluon_trace_runs_copy_scalar_kernel(): kernel = triton_viz.trace("tracer", frontend="gluon")(_copy_scalar_kernel) @@ -820,9 +861,7 @@ def test_gluon_tma_runs_1d_copy_on_cpu(): layout = gl.NVMMASharedLayout.get_default_for([64], gl.float32) in_desc = TensorDescriptor.from_tensor(inp, [64], layout) out_desc = TensorDescriptor.from_tensor(out, [64], layout) - kernel = triton_viz.trace(_NoOpClient(), frontend="gluon")(_tma_copy_1d_kernel) - - kernel[(1,)](in_desc, out_desc, 64, num_warps=1) + _run_gluon_on_cpu(_tma_copy_1d_kernel, (1,), in_desc, out_desc, 64, num_warps=1) torch.testing.assert_close(out, inp, atol=0, rtol=0) @@ -835,12 +874,18 @@ def test_gluon_tma_runs_staged_elementwise_add_on_cpu(): a_desc = TensorDescriptor.from_tensor(a, [4, 8], layout) b_desc = TensorDescriptor.from_tensor(b, [4, 8], layout) out_desc = TensorDescriptor.from_tensor(out, [4, 8], layout) - kernel = triton_viz.trace(_NoOpClient(), frontend="gluon")( - _tma_elementwise_add_kernel + _run_gluon_on_cpu( + _tma_elementwise_add_kernel, + (1,), + a_desc, + b_desc, + out_desc, + *a.shape, + 4, + 8, + num_warps=4, ) - kernel[(1,)](a_desc, b_desc, out_desc, *a.shape, 4, 8, num_warps=4) - torch.testing.assert_close(out, a + b, atol=0, rtol=0) @@ -858,16 +903,23 @@ def test_gluon_tma_runs_float_atomics_on_cpu(): add_desc = TensorDescriptor.from_tensor(add_dst, [block_m, block_n], layout) min_desc = TensorDescriptor.from_tensor(min_dst, [block_m, block_n], layout) max_desc = TensorDescriptor.from_tensor(max_dst, [block_m, block_n], layout) - kernel = triton_viz.trace(_NoOpClient(), frontend="gluon")(_tma_atomic_float_kernel) - expected_add = add_dst.clone() expected_min = min_dst.clone() expected_max = max_dst.clone() expected_add[1:3, 2:6] += src expected_min[1:3, 2:6] = torch.minimum(expected_min[1:3, 2:6], src) expected_max[1:3, 2:6] = torch.maximum(expected_max[1:3, 2:6], src) - - kernel[(1,)](add_desc, min_desc, max_desc, src, block_m, block_n, num_warps=4) + _run_gluon_on_cpu( + _tma_atomic_float_kernel, + (1,), + add_desc, + min_desc, + max_desc, + src, + block_m, + block_n, + num_warps=4, + ) torch.testing.assert_close(add_dst, expected_add, atol=0, rtol=0) torch.testing.assert_close(min_dst, expected_min, atol=0, rtol=0) @@ -925,6 +977,29 @@ def test_gluon_wgmma_runs_small_mma_with_lhs_registers_on_cpu(): _run_small_wgmma_case(seed=1, lhs_in_reg=True) +def test_gluon_builder_preserves_tensor_memory_fp4_padding(): + layout = gluon_builder.get_tensor_memory_layout( + (64, 64), + 1, + [[1, 0], [0, 1]], + False, + True, + ) + if not hasattr(layout, "fp4_padded"): + pytest.skip("Gluon TensorMemoryLayout has no fp4_padded field") + assert layout.fp4_padded is True + + +def test_gluon_blackwell_tensor_memory_roundtrips_tile_on_cpu(): + inp = torch.arange(64 * 64, dtype=torch.float32).reshape(64, 64) + out = torch.empty_like(inp) + _run_gluon_on_cpu( + _tensor_memory_roundtrip_kernel, (1,), inp, out, *inp.shape, num_warps=4 + ) + + torch.testing.assert_close(out, inp, atol=0, rtol=0) + + def test_gluon_blackwell_tma_runs_gather_on_cpu(): block_x = 8 block_y = 8 @@ -934,11 +1009,9 @@ def test_gluon_blackwell_tma_runs_gather_on_cpu(): out = torch.full((block_x, block_y), -1.0) layout = gl.NVMMASharedLayout.get_default_for([block_x, block_y], gl.float32) desc = TensorDescriptor.from_tensor(inp, [1, block_y], layout) - kernel = triton_viz.trace(_NoOpClient(), frontend="gluon")( - _blackwell_tma_gather_kernel - ) - - kernel[(1,)]( + _run_gluon_on_cpu( + _blackwell_tma_gather_kernel, + (1,), out, *out.stride(), desc, @@ -972,18 +1045,23 @@ def test_gluon_blackwell_tma_runs_bitwise_atomics_on_cpu(): and_desc = TensorDescriptor.from_tensor(and_dst, [block_m, block_n], layout) or_desc = TensorDescriptor.from_tensor(or_dst, [block_m, block_n], layout) xor_desc = TensorDescriptor.from_tensor(xor_dst, [block_m, block_n], layout) - kernel = triton_viz.trace(_NoOpClient(), frontend="gluon")( - _blackwell_tma_bitwise_atomic_kernel - ) - expected_and = and_dst.clone() expected_or = or_dst.clone() expected_xor = xor_dst.clone() expected_and[1:3, 1:5] = torch.bitwise_and(expected_and[1:3, 1:5], src) expected_or[1:3, 1:5] = torch.bitwise_or(expected_or[1:3, 1:5], src) expected_xor[1:3, 1:5] = torch.bitwise_xor(expected_xor[1:3, 1:5], src) - - kernel[(1,)](and_desc, or_desc, xor_desc, src, block_m, block_n, num_warps=4) + _run_gluon_on_cpu( + _blackwell_tma_bitwise_atomic_kernel, + (1,), + and_desc, + or_desc, + xor_desc, + src, + block_m, + block_n, + num_warps=4, + ) torch.testing.assert_close(and_dst, expected_and, atol=0, rtol=0) torch.testing.assert_close(or_dst, expected_or, atol=0, rtol=0) @@ -1002,11 +1080,9 @@ def test_gluon_blackwell_tma_runs_scatter_on_cpu(): ) layout = gl.NVMMASharedLayout.get_default_for([block_x, block_y], gl.float32) desc = TensorDescriptor.from_tensor(out, [1, block_y], layout) - kernel = triton_viz.trace(_NoOpClient(), frontend="gluon")( - _blackwell_tma_scatter_kernel - ) - - kernel[(1,)]( + _run_gluon_on_cpu( + _blackwell_tma_scatter_kernel, + (1,), desc, x_offsets, y_offset, diff --git a/triton_viz/core/simulation/gluon.py b/triton_viz/core/simulation/gluon.py index b2b926bb..4960bdb4 100644 --- a/triton_viz/core/simulation/gluon.py +++ b/triton_viz/core/simulation/gluon.py @@ -1387,16 +1387,37 @@ def get_tensor_memory_layout( self, block: Any, col_stride: Any, - cga_layout: Any, - two_ctas: Any, - fp4_padded: Any, + cga_layout: Any = None, + two_ctas: Any = False, + fp4_padded: Any = False, ): + cga_bases = [] if cga_layout is None else [list(basis) for basis in cga_layout] + layout_fields = getattr( + gluon_blackwell.TensorMemoryLayout, + "__dataclass_fields__", + {}, + ) + # Older Gluon builds used cta_split_num and had no FP4 padding field. + if "fp4_padded" in layout_fields: + return gluon_blackwell.TensorMemoryLayout( + tuple(block), + int(col_stride), + cga_bases, + bool(two_ctas), + bool(fp4_padded), + ) + if "cta_split_num" in layout_fields: + return gluon_blackwell.TensorMemoryLayout( + tuple(block), + int(col_stride), + cga_bases or None, + bool(two_ctas), + ) return gluon_blackwell.TensorMemoryLayout( tuple(block), int(col_stride), - [] if cga_layout is None else [list(basis) for basis in cga_layout], + cga_bases, bool(two_ctas), - bool(fp4_padded), ) def get_tensor_memory_scales_layout(self, cga_layout: Any): @@ -2431,6 +2452,18 @@ def create_warp_return(self): ) +_GLUON_NON_SEMANTIC_BUILTINS: tuple[tuple[Any, str, Callable], ...] = tuple( + (module, name, getattr(module, name)) + for module, name in ( + (gluon_amd_cdna4, "_compute_efficient_padded_shared_layout_impl"), + (gluon_amd_cdna4, "_get_mfma_scale_layout_impl"), + (gluon_amd_gfx1250, "_get_wmma_scale_layout_impl"), + (gluon_blackwell, "_compute_tmem_reg_layout"), + ) + if module is not None and hasattr(module, name) +) + + def _make_tma_atomic_reduce(kind: Any) -> Callable: @gluon_core.builtin def async_atomic_reduce(tensor_desc, coord, src, _semantic=None): @@ -2479,6 +2512,21 @@ def _yield_warp_specialize() -> None: gluon_frontend._maybe_yield_warp_specialize() +def _patch_gluon_non_semantic_builtins(scope: _LangPatchScope) -> None: + for module, name, member in _GLUON_NON_SEMANTIC_BUILTINS: + + def new_member( + *args: Any, + member: Callable = member, + **kwargs: Any, + ): + _yield_warp_specialize() + kwargs.pop("_semantic", None) + return member(*args, **kwargs) + + scope.set_attr(module, name, new_member) + + def _patch_gluon_builtins(pkg: Any, scope: _LangPatchScope) -> None: for name, member in inspect.getmembers(pkg): if not callable(member): @@ -2489,7 +2537,11 @@ def _patch_gluon_builtins(pkg: Any, scope: _LangPatchScope) -> None: ): continue - def new_member(*args: Any, member: Callable = member, **kwargs: Any): + def new_member( + *args: Any, + member: Callable = member, + **kwargs: Any, + ): _yield_warp_specialize() # Gluon builtins may pass a compile-time semantic object; replace it # with this builder-backed semantic while preserving user arguments. @@ -2509,6 +2561,7 @@ def patch_lang(fn: Callable) -> _LangPatchScope: _patch_gluon_builtins(module, scope) for cls in _GLUON_BUILTIN_CLASSES: _patch_gluon_builtins(cls, scope) + _patch_gluon_non_semantic_builtins(scope) def allocate_mbarrier(*_args: Any, **_kwargs: Any): return _tensor_result(np.array([0], dtype=np.int32), tl.int32)