diff --git a/README.md b/README.md index a0413df4..862444c4 100644 --- a/README.md +++ b/README.md @@ -187,7 +187,7 @@ simulation runtimes live under `tilelens/core/simulation/`. Analyze kernels across visualization, profiling, and sanitization with a single line of code. -- Visualizer: currently supports load, store, and matmul operations for 1/2/3D tensors (more operations and dimensions coming soon). Triton's interpreter runs loads, stores, and atomics on raw host memory, so the tracer raises `IndexError` (wrapped in an `InterpreterError` under Gluon) before a Triton load, store, or atomic, or a Gluon `gl.load`/`gl.store`, that runs off the storage of a tensor argument, instead of letting it corrupt memory. This guard is best-effort: accesses that never touch a tensor argument (for example a whole program past the end), programs that cast an integer to a pointer (`x.to(tl.pointer_type(...))`, as pointer tables do), and Gluon atomics and async copies are not checked, so use the Sanitizer for a complete out-of-bounds report. +- Visualizer: currently supports load, store, and matmul operations for 1/2/3D tensors (more operations and dimensions coming soon). Triton's interpreter runs loads, stores, and atomics on raw host memory, so the tracer raises `IndexError` (wrapped in an `InterpreterError` under Gluon) before a Triton load, store, or atomic, or a Gluon `gl.load`, `gl.store`, `gl.atomic_*`, AMD `buffer_atomic_*`, or Ampere/AMD `async_copy` operation, that runs off the storage of a tensor argument, instead of letting it corrupt memory. This guard is best-effort: accesses that never touch a tensor argument (for example a whole program past the end), programs that cast an integer to a pointer (`x.to(tl.pointer_type(...))`, as pointer tables do), and AMD `buffer_load`/`buffer_store` are not checked, so use the Sanitizer for a complete out-of-bounds report. - Profiler: flags non-unrolled loops, inefficient mask usage, and missing buffer_load optimizations while tracking load/store byte counts with low-overhead sampling. - Sanitizer: symbolically checks tensor memory accesses for out-of-bounds errors and emits reports with tensor metadata, call stack, and expression trees; optional fake-memory storage avoids real reads. diff --git a/tests/end_to_end/test_gluon.py b/tests/end_to_end/test_gluon.py index f7ba26fe..c1bcc881 100644 --- a/tests/end_to_end/test_gluon.py +++ b/tests/end_to_end/test_gluon.py @@ -7,7 +7,11 @@ from triton import knobs from triton.experimental import gluon from triton.experimental.gluon import language as gl +from triton.experimental.gluon.language.amd import cdna3 from triton.experimental.gluon.language.amd.cdna4 import async_copy as amd_cdna4_cp +from triton.experimental.gluon.language.amd.gfx1250 import ( + async_copy as amd_gfx1250_cp, +) 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 ( @@ -23,6 +27,7 @@ from tilelens.clients.sanitizer.sanitizer import SymbolicSanitizer from tilelens.core.data import Load from tilelens.core.simulation.gluon import GluonInterpretedFunction, gluon_builder +from tilelens.visualizer.draw import get_visualization_data try: from triton.experimental.gluon.language.nvidia.ampere import async_copy as cp @@ -119,6 +124,58 @@ def _unmasked_store_1d_memcpy_kernel( gl.store(out_ptr + offsets, value) +@gluon.jit +def _pointer_table_kernel(table_ptr, arg_ptr, out_ptr, BLOCK: gl.constexpr): + layout: gl.constexpr = gl.BlockedLayout([1, 1], [2, 16], [1, 1], [1, 0]) + rows = gl.arange(0, 2, layout=gl.SliceLayout(1, layout)) + cols = gl.arange(0, BLOCK, layout=gl.SliceLayout(0, layout)) + row_ptrs = gl.load(table_ptr + rows).to(gl.pointer_type(gl.float32)) + values = gl.load(row_ptrs[:, None] + cols[None, :]) + gl.store(out_ptr + rows[:, None] * BLOCK + cols[None, :], values) + + +@gluon.jit +def _int_to_ptr_bitcast_copy_kernel(addr, src_ptr, out_ptr, BLOCK: gl.constexpr): + layout: gl.constexpr = gl.BlockedLayout([1], [32], [1], [0]) + offs = gl.arange(0, BLOCK, layout=layout) + words = addr.to(gl.int64).to(gl.pointer_type(gl.int32)) + floats = words.to(gl.pointer_type(gl.float32), bitcast=True) + gl.store(out_ptr + offs, gl.load(floats + offs)) + + +@gluon.jit +def _atomic_add_kernel(out_ptr, xnumel, BLOCK: gl.constexpr, layout: gl.constexpr): + offsets = gl.program_id(0) * BLOCK + gl.arange(0, BLOCK, layout=layout) + gl.atomic_add(out_ptr + offsets, 1.0, mask=offsets < xnumel) + + +@gluon.jit +def _atomic_cas_kernel(out_ptr, xnumel, BLOCK: gl.constexpr, layout: gl.constexpr): + offsets = gl.program_id(0) * BLOCK + gl.arange(0, BLOCK, layout=layout) + gl.atomic_cas( + out_ptr + offsets, + gl.full([BLOCK], -7, gl.int32, layout), + gl.full([BLOCK], 1, gl.int32, layout), + ) + + +@gluon.jit +def _buffer_atomic_add_kernel( + out_ptr, xnumel, BLOCK: gl.constexpr, layout: gl.constexpr +): + offsets = gl.program_id(0) * BLOCK + gl.arange(0, BLOCK, layout=layout) + values = gl.full([BLOCK], 1.0, gl.float32, layout) + cdna3.buffer_atomic_add(out_ptr, offsets, values, mask=offsets < xnumel) + + +@gluon.jit +def _spin_lock_kernel(lock_ptr, out_ptr): + while gl.atomic_cas(lock_ptr, 0, 1) == 1: + pass + gl.store(out_ptr, gl.load(out_ptr) + 1) + gl.atomic_xchg(lock_ptr, 0) + + @gluon.jit def _masked_1d_memcpy_kernel( in_ptr, @@ -247,6 +304,125 @@ def _amd_async_copy_other_kernel(in_ptr, out_ptr, xnumel, BLOCK: gl.constexpr): gl.store(out_ptr + offsets, values) +@gluon.jit +def _ampere_async_load_kernel( + in_ptr, + out_ptr, + BLOCK: gl.constexpr, + layout: gl.constexpr, + smem_layout: gl.constexpr, +): + offsets = gl.arange(0, BLOCK, layout=layout) + smem = gl.allocate_shared_memory(gl.float32, [BLOCK], layout=smem_layout) + cp.async_load(smem, in_ptr + offsets) + cp.commit_group() + cp.wait_group(0) + gl.store(out_ptr + offsets, smem.load(layout)) + + +@gluon.jit +def _ampere_async_copy_global_to_shared_kernel( + in_ptr, + out_ptr, + BLOCK: gl.constexpr, + layout: gl.constexpr, + smem_layout: gl.constexpr, +): + offsets = gl.arange(0, BLOCK, layout=layout) + smem = gl.allocate_shared_memory(gl.float32, [BLOCK], layout=smem_layout) + cp.async_copy_global_to_shared(smem, in_ptr + offsets) + cp.commit_group() + cp.wait_group(0) + gl.store(out_ptr + offsets, smem.load(layout)) + + +@gluon.jit +def _cdna4_global_load_to_shared_kernel( + in_ptr, + out_ptr, + BLOCK: gl.constexpr, + layout: gl.constexpr, + smem_layout: gl.constexpr, +): + offsets = gl.arange(0, BLOCK, layout=layout) + smem = gl.allocate_shared_memory(gl.float32, [BLOCK], layout=smem_layout) + amd_cdna4_cp.global_load_to_shared(smem, in_ptr + offsets) + amd_cdna4_cp.wait_group(0) + gl.store(out_ptr + offsets, smem.load(layout)) + + +@gluon.jit +def _cdna4_buffer_load_to_shared_kernel( + in_ptr, + out_ptr, + BLOCK: gl.constexpr, + layout: gl.constexpr, + smem_layout: gl.constexpr, +): + offsets = gl.arange(0, BLOCK, layout=layout) + smem = gl.allocate_shared_memory(gl.float32, [BLOCK], layout=smem_layout) + # the CPU simulation of buffer loads needs an explicit mask and other + amd_cdna4_cp.buffer_load_to_shared( + smem, in_ptr, offsets, mask=offsets < BLOCK, other=0.0 + ) + amd_cdna4_cp.wait_group(0) + gl.store(out_ptr + offsets, smem.load(layout)) + + +@gluon.jit +def _gfx1250_global_to_shared_kernel( + in_ptr, + out_ptr, + BLOCK: gl.constexpr, + layout: gl.constexpr, + smem_layout: gl.constexpr, +): + offsets = gl.arange(0, BLOCK, layout=layout) + smem = gl.allocate_shared_memory(gl.float32, [BLOCK], layout=smem_layout) + amd_gfx1250_cp.global_to_shared(smem, in_ptr + offsets) + amd_gfx1250_cp.wait_group(0) + gl.store(out_ptr + offsets, smem.load(layout)) + + +@gluon.jit +def _gfx1250_shared_to_global_kernel( + out_ptr, BLOCK: gl.constexpr, layout: gl.constexpr, smem_layout: gl.constexpr +): + offsets = gl.arange(0, BLOCK, layout=layout) + smem = gl.allocate_shared_memory(gl.float32, [BLOCK], layout=smem_layout) + smem.store(gl.full([BLOCK], 3.0, gl.float32, layout)) + # the CPU simulation of shared_to_global needs an explicit mask + amd_gfx1250_cp.shared_to_global(out_ptr + offsets, smem, mask=offsets < BLOCK) + amd_gfx1250_cp.wait_group(0) + + +@gluon.jit +def _ampere_async_load_col_mask_kernel( + in_ptr, + out_ptr, + n_cols, + M: gl.constexpr, + N: gl.constexpr, + layout: gl.constexpr, + smem_layout: gl.constexpr, +): + # a matmul-style K-tail mask that the builtin broadcasts over the tile + rows = gl.arange(0, M, layout=gl.SliceLayout(1, layout)) + cols = gl.arange(0, N, layout=gl.SliceLayout(0, layout)) + smem = gl.allocate_shared_memory(gl.float32, [M, N], layout=smem_layout) + cp.async_load( + smem, + in_ptr + rows[:, None] * n_cols + cols[None, :], + mask=cols[None, :] < n_cols, + ) + cp.commit_group() + cp.wait_group(0) + mask = cols[None, :] < n_cols + gl.store( + out_ptr + rows[:, None] * n_cols + cols[None, :], smem.load(layout), mask=mask + ) + + @gluon.jit def _async_copy_elementwise_add_kernel( a_ptr, @@ -860,6 +1036,22 @@ def test_gluon_sanitizer_allows_masked_in_bounds_kernel(): assert sanitizer.records == [] +def test_gluon_sanitizer_runs_spin_lock_kernel(): + # Gluon atomics reach clients as Gluon-specific op types, which the + # symbolic clients do not override; as symbolic atomics the spin loop + # could not be evaluated + sanitizer = SymbolicSanitizer(abort_on_error=False) + kernel = tilelens.trace(client=sanitizer, frontend="gluon")(_spin_lock_kernel) + lock = torch.zeros(1, dtype=torch.int32) + out = torch.zeros(1) + + ret = kernel[(3,)](lock, out, num_warps=1) + + assert ret is None + assert sanitizer.records == [] + assert lock.item() == 0 + + def test_gluon_core_ops_run_scalar_range_memcpy_on_cpu(): inp = torch.arange(40, dtype=torch.float32) out = torch.full_like(inp, -1) @@ -902,6 +1094,155 @@ def test_gluon_trace_refuses_out_of_bounds_store(): assert (buf[:64] == -7).all() and (buf[104:] == -7).all() +@pytest.mark.parametrize( + "kernel_fn", + [ + _ampere_async_load_kernel, + _ampere_async_copy_global_to_shared_kernel, + _cdna4_global_load_to_shared_kernel, + _cdna4_buffer_load_to_shared_kernel, + _gfx1250_global_to_shared_kernel, + ], +) +def test_gluon_trace_checks_async_copy_loads(kernel_fn): + layout = gl.BlockedLayout([1], [32], [1], [0]) + smem_layout = gl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[0]) + kernel = tilelens.trace("tracer", frontend="gluon")(kernel_fn) + inp = torch.arange(64, dtype=torch.float32) + out = torch.full_like(inp, -1) + + kernel[(1,)](inp, out, 64, layout, smem_layout, num_warps=1) + + torch.testing.assert_close(out, inp, atol=0, rtol=0) + # the async copy is the kernel's only global read + assert any(isinstance(r, Load) for r in kernel.client_manager.launch.records) + + # the same unmasked copy now reads 64 floats from a 40-float storage + buf = np.full(40 + 128, -7, dtype=np.float32) + short_inp = torch.from_numpy(buf[64:104]) + with pytest.raises(Exception) as exc_info: + kernel[(1,)](short_inp, out, 64, layout, smem_layout, num_warps=1) + + cause = exc_info.value.__cause__ + assert isinstance(cause, IndexError) + assert "out-of-bounds load" in str(cause) + + +def test_gluon_trace_records_async_copy_with_broadcast_mask(): + tilelens.clear() + layout = gl.BlockedLayout([1, 1], [4, 8], [1, 1], [1, 0]) + smem_layout = gl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[1, 0]) + kernel = tilelens.trace("tracer", frontend="gluon")( + _ampere_async_load_col_mask_kernel + ) + inp = torch.arange(24, dtype=torch.float32) + out = torch.full_like(inp, -1) + + kernel[(1,)](inp, out, 6, 4, 8, layout, smem_layout, num_warps=1) + + torch.testing.assert_close(out, inp, atol=0, rtol=0) + loads = [r for r in kernel.client_manager.launch.records if isinstance(r, Load)] + assert loads and all(r.masks.shape == r.offsets.shape for r in loads) + get_visualization_data() + + +def test_gluon_trace_refuses_out_of_bounds_async_copy_store(): + # out's storage is exactly 40 floats; the rest of buf is sentinel memory + buf = np.full(40 + 128, -7, dtype=np.float32) + out = torch.from_numpy(buf[64:104]) + layout = gl.BlockedLayout([1], [32], [1], [0]) + smem_layout = gl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[0]) + kernel = tilelens.trace("tracer", frontend="gluon")( + _gfx1250_shared_to_global_kernel + ) + + with pytest.raises(Exception) as exc_info: + kernel[(1,)](out, 64, layout, smem_layout, num_warps=1) + + cause = exc_info.value.__cause__ + assert isinstance(cause, IndexError) + assert "out-of-bounds store" in str(cause) + assert (buf[:64] == -7).all() and (buf[104:] == -7).all() + + +def test_gluon_trace_allows_pointer_table_gather(): + # one row is a kernel arg and one is not, so a single gather touches known + # and unknown storage; pointers built from integers are not judged + arg_row = torch.arange(4, dtype=torch.float32) + other_row = torch.arange(4, 8, dtype=torch.float32) + table = torch.tensor([arg_row.data_ptr(), other_row.data_ptr()]) + out = torch.empty(8) + kernel = tilelens.trace("tracer", frontend="gluon")(_pointer_table_kernel) + + kernel[(1,)](table, arg_row, out, 4, num_warps=1) + + torch.testing.assert_close(out, torch.arange(8, dtype=torch.float32)) + + +def test_gluon_sanitizer_allows_pointer_cast_after_int_to_ptr(): + # the sanitizer must keep seeing a concrete pointer here: the follow-up + # pointer bitcast runs on the simulation builder without interception + sanitizer = SymbolicSanitizer(abort_on_error=False) + kernel = tilelens.trace(client=sanitizer, frontend="gluon")( + _int_to_ptr_bitcast_copy_kernel + ) + inp = torch.arange(32, dtype=torch.float32) + out = torch.empty_like(inp) + + # inp also goes in as src_ptr so the sanitizer knows its storage + ret = kernel[(1,)](inp.data_ptr(), inp, out, 32, num_warps=1) + + assert ret is None + assert sanitizer.records == [] + + +@pytest.mark.parametrize( + "kernel, dtype, op_name", + [ + (_atomic_add_kernel, np.float32, "atomic"), + (_atomic_cas_kernel, np.int32, "atomic_cas"), + (_buffer_atomic_add_kernel, np.float32, "atomic"), + ], +) +def test_gluon_trace_refuses_out_of_bounds_atomics(kernel, dtype, op_name): + # out's storage is exactly 40 elements; the kernels' masks admit 48 lanes + # (atomic_cas has no mask) and the rest of buf is sentinel memory + buf = np.full(40 + 128, -7, dtype=dtype) + out = torch.from_numpy(buf[64:104]) + layout = gl.BlockedLayout([1], [32], [1], [0]) + traced = tilelens.trace("tracer", frontend="gluon")(kernel) + + with pytest.raises(Exception) as exc_info: + traced[(1,)](out, 48, 64, layout, num_warps=1) + + cause = exc_info.value.__cause__ + assert isinstance(cause, IndexError) + assert f"out-of-bounds {op_name} " in str(cause) + assert (buf[:64] == -7).all() and (buf[104:] == -7).all() + + +def test_gluon_trace_runs_in_bounds_atomics(): + layout = gl.BlockedLayout([1], [32], [1], [0]) + added = torch.zeros(40) + buffer_added = torch.zeros(40) + swapped = torch.full((64,), -7, dtype=torch.int32) + + # the second program's lanes past the end are masked off + tilelens.trace("tracer", frontend="gluon")(_atomic_add_kernel)[(2,)]( + added, 40, 32, layout, num_warps=1 + ) + tilelens.trace("tracer", frontend="gluon")(_buffer_atomic_add_kernel)[(1,)]( + buffer_added, 40, 64, layout, num_warps=1 + ) + tilelens.trace("tracer", frontend="gluon")(_atomic_cas_kernel)[(1,)]( + swapped, 64, 64, layout, num_warps=1 + ) + + torch.testing.assert_close(added, torch.ones(40), atol=0, rtol=0) + torch.testing.assert_close(buffer_added, torch.ones(40), atol=0, rtol=0) + assert (swapped == 1).all() + + def test_gluon_core_ops_run_masked_2d_memcpy_on_cpu(): inp = torch.arange(24, dtype=torch.float32).reshape(4, 6) out = torch.full_like(inp, -1) diff --git a/tests/unit/test_adapters.py b/tests/unit/test_adapters.py index 176e64c9..68f74141 100644 --- a/tests/unit/test_adapters.py +++ b/tests/unit/test_adapters.py @@ -367,6 +367,31 @@ def test_gluon_load_store_adapters_keep_first_positional_mask(): ) +def test_gluon_async_copy_load_adapters_accept_cdna4_keywords(): + from tilelens.core.frontend import gluon as gluon_frontend + + async_copy = GLUON_ADAPTERS[gluon_frontend.GluonAsyncCopyLoad] + buffer_to_shared = GLUON_ADAPTERS[gluon_frontend.GluonBufferLoadToShared] + + assert async_copy("smem", "ptr", "mask").args == ("ptr", "mask", None) + # CDNA4 spells the operands dest, ptr (and offsets for buffer_load_to_shared) + assert async_copy(dest="smem", ptr="ptr", mask="mask").args == ("ptr", "mask", None) + assert buffer_to_shared(dest="smem", ptr=100, offsets=4, mask="mask").args == ( + 104, + "mask", + None, + ) + + +def test_gluon_ampere_async_copies_do_not_get_the_tma_descriptor_adapter(): + """Ampere async_load shares its name with the TMA load but takes (smem, pointer, mask).""" + from triton.experimental.gluon.language.nvidia.ampere import async_copy + from tilelens.core.frontend import gluon as gluon_frontend + + for fn in (async_copy.async_load, async_copy.async_copy_global_to_shared): + assert fn not in gluon_frontend.GLUON_CALLABLE_ADAPTERS + + def test_gluon_load_adapter_does_not_treat_pred_as_mask(): from tilelens.core.frontend import gluon as gluon_frontend @@ -377,6 +402,40 @@ def test_gluon_load_adapter_does_not_treat_pred_as_mask(): ) +def test_gluon_builder_atomic_adapters_skip_the_builder(): + """Gluon atomics are patched on the simulation Builder class, so self comes first.""" + from tilelens.core.frontend import gluon as gluon_frontend + + builder, rmw_op, ptr, cmp, val, mask, sem, scope = (object() for _ in range(8)) + rmw = GLUON_ADAPTERS[gluon_frontend.GluonAtomicRMW] + cas = GLUON_ADAPTERS[gluon_frontend.GluonAtomicCas] + + assert rmw(builder, rmw_op, ptr, val, mask, sem, scope).args == (ptr, mask) + assert cas(builder, ptr, cmp, val, sem, scope).args == (ptr,) + + +def test_gluon_buffer_atomic_adapter_returns_per_lane_pointers(): + """create_buffer_atomic_rmw(rmw_op, ptr, offsets, value, sem, scope, mask) -> (ptrs, mask).""" + from tilelens.core.frontend import gluon as gluon_frontend + from tilelens.core.simulation.gluon import gluon_builder + + adapter = GLUON_ADAPTERS[gluon_frontend.GluonBufferAtomicRMW] + ptr = TensorHandle(np.array([4096], dtype=np.uint64), tl.pointer_type(tl.float32)) + offsets = TensorHandle(np.arange(4, dtype=np.int32), tl.int32) + mask = TensorHandle(np.array([True, True, False, True]), tl.int1) + + ptrs, result_mask = adapter( + gluon_builder, "add", ptr, offsets, "value", "sem", "scope", mask + ).args + assert np.array_equal(ptrs.data, 4096 + 4 * np.arange(4)) + assert result_mask is mask + # unmasked calls pass an empty ir.value() instead of a mask + unmasked = adapter( + gluon_builder, "add", ptr, offsets, "value", "sem", "scope", object() + ) + assert unmasked.args[1] is None + + def test_gluon_frontend_routes_tensor_descriptor_access_to_overrider(monkeypatch): from tilelens.core.frontend import gluon as gluon_frontend diff --git a/tests/unit/test_tracer.py b/tests/unit/test_tracer.py index 4e842b84..6b628f6d 100644 --- a/tests/unit/test_tracer.py +++ b/tests/unit/test_tracer.py @@ -8,7 +8,7 @@ from triton.runtime.interpreter import TensorHandle from tilelens.clients.tracer.tracer import Tracer, _convert_grid_idx -from tilelens.core.data import Transfer +from tilelens.core.data import AtomicRMW, IntToPtr, Load, Transfer # ======== _convert_grid_idx Tests =========== @@ -180,3 +180,31 @@ def ptr(*offsets): tracer._check_in_bounds("load", ptr(0, 4, 8, 12), True) with pytest.raises(IndexError, match="out-of-bounds load"): tracer._check_in_bounds("load", ptr(8, 12, 16), True) + + +def test_tracer_traces_op_subclasses_as_their_core_op(): + # frontends subclass a core op to pick another argument adapter (Gluon's + # builder atomics, int->pointer casts and async copies) + class FrontendLoad(Load): + pass + + class FrontendAtomicRMW(AtomicRMW): + pass + + class FrontendIntToPtr(IntToPtr): + pass + + tracer = Tracer() + tensor = torch.zeros(4) + tracer.arg_callback("x_ptr", tensor, None) + lanes = np.array([tensor.data_ptr() + offset for offset in (8, 12, 16)]) + ptr = TensorHandle(lanes.astype(np.uint64), tl.pointer_type(tl.float32)) + + with pytest.raises(IndexError, match="out-of-bounds load"): + tracer.register_op_callback(FrontendLoad).before_callback(ptr, None, None) + with pytest.raises(IndexError, match="out-of-bounds atomic"): + tracer.register_op_callback(FrontendAtomicRMW).before_callback(ptr, None) + + tracer.grid_idx_callback((0, 0, 0)) + tracer.register_op_callback(FrontendIntToPtr).after_callback(None) + assert tracer._get_thread_local("int_to_ptr") is True diff --git a/tilelens/clients/tracer/tracer.py b/tilelens/clients/tracer/tracer.py index c787bbba..a14d6fbc 100644 --- a/tilelens/clients/tracer/tracer.py +++ b/tilelens/clients/tracer/tracer.py @@ -18,6 +18,7 @@ from ...utils.traceback_utils import extract_user_frames from tilelens.core.masked_load_store import masked_load import numpy as np +from triton.runtime.interpreter import TensorHandle def _convert_grid_idx(grid_idx) -> tuple[int, int, int] | None: @@ -47,6 +48,27 @@ def _storage_ranges(name, arg): yield name, tensor, start, start + storage.nbytes() +def _paired_lanes(ptr, mask): + """Broadcast pointer and mask lanes together. + + Gluon adapters see a builtin's operands before it broadcasts them, so an + async copy can hand over a row or scalar mask for a full block of pointers. + """ + data = getattr(ptr, "data", None) + mask_data = getattr(mask, "data", None) + if ( + not isinstance(data, np.ndarray) + or not isinstance(mask_data, np.ndarray) + or data.shape == mask_data.shape + ): + return ptr, mask + try: + data, mask_data = np.broadcast_arrays(data, mask_data) + except ValueError: + return ptr, mask + return TensorHandle(data, ptr.dtype), TensorHandle(mask_data, mask.dtype) + + class Tracer(Client): NAME = "tracer" @@ -209,6 +231,7 @@ def _convert_keys_to_numpy(keys): @self.lock_fn def pre_load_callback(ptr, mask, keys): if keys is None: + ptr, mask = _paired_lanes(ptr, mask) self._check_in_bounds("load", ptr, mask) if not self.sample: return @@ -229,6 +252,7 @@ def pre_load_callback(ptr, mask, keys): @self.lock_fn def pre_store_callback(ptr, mask, keys): if keys is None: + ptr, mask = _paired_lanes(ptr, mask) self._check_in_bounds("store", ptr, mask) if not self.sample: return @@ -344,7 +368,13 @@ def post_dot_callback(ret, input, other): ReduceSum: OpCallbacks(after_callback=post_reduce_sum_callback), Dot: OpCallbacks(after_callback=post_dot_callback), } - return callbacks.get(op_type, OpCallbacks()) + # Frontends subclass a core op to pick another argument adapter (Gluon's + # builder atomics, int->pointer casts and async copies); the adapter + # hands over the core op's arguments, so trace them as that op. + for op_class in op_type.__mro__: + if op_class in callbacks: + return callbacks[op_class] + return OpCallbacks() def register_for_loop_callback(self): return ForLoopCallbacks() diff --git a/tilelens/core/frontend/gluon.py b/tilelens/core/frontend/gluon.py index 7f7a1460..bb9701c0 100644 --- a/tilelens/core/frontend/gluon.py +++ b/tilelens/core/frontend/gluon.py @@ -21,6 +21,9 @@ from triton.experimental.gluon.language.amd.gfx1250 import ( # type: ignore tdm as gluon_amd_tdm, ) +from triton.experimental.gluon.language.nvidia.ampere import ( # type: ignore + async_copy as gluon_ampere_async_copy, +) from triton.experimental.gluon.language.nvidia.blackwell import ( # type: ignore tma as gluon_blackwell_tma, ) @@ -32,11 +35,14 @@ from ..data import ( AddPtr, Allocate, + AtomicCas, + AtomicRMW, BinaryOp, Broadcast, Dot, ExpandDims, Fma, + IntToPtr, Join, Load, MakeRange, @@ -56,6 +62,7 @@ TensorDescriptorAccess, ) +from ..simulation.gluon import Builder as GluonBuilder from .base import AdapterResult, Frontend, _LangPatchScope, register_frontend from .triton import TritonFrontend @@ -75,6 +82,26 @@ class GluonBufferLoadToShared(Load): pass +# Atomics and int->pointer casts that Gluon builds on the simulation builder. +# Each gets its own subclass for its builder-signature adapter, so clients that +# match op types exactly (the symbolic clients) still run them concretely, as +# before; their symbolic results would reach Gluon ops that are not dispatched. +class GluonAtomicRMW(AtomicRMW): + pass + + +class GluonAtomicCas(AtomicCas): + pass + + +class GluonBufferAtomicRMW(AtomicRMW): + pass + + +class GluonIntToPtr(IntToPtr): + pass + + def _set_warp_specialize_scheduler(scheduler: Any) -> Any: global _WARP_SPECIALIZE_SCHEDULER previous = _WARP_SPECIALIZE_SCHEDULER @@ -133,14 +160,17 @@ def adapter( def _gluon_async_copy_load_adapter( - smem: Any, - pointer: Any, + smem: Any = None, + pointer: Any = None, mask: Any = None, *_args: Any, + ptr: Any = None, **_kwargs: Any, ) -> AdapterResult: # Async copy load traces source global memory, not the destination SMEM. - return _gluon_pointer_load_adapter(pointer, mask=mask) + # CDNA4 global_load_to_shared names its operands (dest, ptr), not + # (smem, pointer), so accept its keyword spelling too. + return _gluon_pointer_load_adapter(ptr if pointer is None else pointer, mask=mask) def _gluon_async_copy_store_adapter( @@ -155,7 +185,7 @@ def _gluon_async_copy_store_adapter( def _gluon_buffer_load_to_shared_adapter( - smem: Any, + dest: Any, ptr: Any, offsets: Any, *args: Any, @@ -190,6 +220,48 @@ def _gluon_pointer_store_adapter( _gluon_store_adapter = _gluon_pointer_store_adapter +# The builder methods below are patched on the Builder class, so the builder +# instance arrives first. +def _gluon_atomic_rmw_adapter( + _builder: Any, + _rmw_op: Any, + ptr: Any, + _val: Any, + mask: Any, + *_args: Any, + **_kwargs: Any, +) -> AdapterResult: + return AdapterResult(ptr, mask) + + +def _gluon_atomic_cas_adapter( + _builder: Any, + ptr: Any, + *_args: Any, + **_kwargs: Any, +) -> AdapterResult: + return AdapterResult(ptr) + + +def _gluon_buffer_atomic_rmw_adapter( + builder: Any, + _rmw_op: Any, + ptr: Any, + offsets: Any, + _value: Any, + _sem: Any, + _scope: Any, + mask: Any = None, + *_args: Any, + **_kwargs: Any, +) -> AdapterResult: + # A buffer atomic addresses a scalar base pointer plus offsets; hand clients + # the per-lane pointers. Unmasked calls pass an empty ir.value() as the mask. + if not isinstance(mask, TensorHandle): + mask = None + return AdapterResult(builder.create_addptr(ptr, offsets), mask) + + def _is_global_tensor_descriptor_like(value: Any) -> bool: type_name = type(value).__name__ return type_name in { @@ -348,6 +420,12 @@ def _existing_ops(namespace: Any, attrs: dict[str, type[Op]]) -> dict[str, type[ "greater_than": BinaryOp, "greater_equal": BinaryOp, }, + GluonBuilder: { + "create_atomic_rmw": GluonAtomicRMW, + "create_atomic_cas": GluonAtomicCas, + "create_buffer_atomic_rmw": GluonBufferAtomicRMW, + "create_int_to_ptr": GluonIntToPtr, + }, } _TMA_NAMESPACES: tuple[tuple[Any, dict[str, type[Op]]], ...] = ( ( @@ -459,6 +537,13 @@ def _existing_ops(namespace: Any, attrs: dict[str, type[Op]]) -> dict[str, type[ "shared_to_global": GluonAsyncCopyStore, }, ), + ( + gluon_ampere_async_copy, + { + "async_load": GluonAsyncCopyLoad, + "async_copy_global_to_shared": GluonAsyncCopyLoad, + }, + ), ) for namespace, attrs in _TMA_NAMESPACES: existing = _existing_ops(namespace, attrs) @@ -472,6 +557,10 @@ def _existing_ops(namespace: Any, attrs: dict[str, type[Op]]) -> dict[str, type[ GluonAsyncCopyLoad: _gluon_async_copy_load_adapter, GluonAsyncCopyStore: _gluon_async_copy_store_adapter, GluonBufferLoadToShared: _gluon_buffer_load_to_shared_adapter, + GluonAtomicRMW: _gluon_atomic_rmw_adapter, + GluonAtomicCas: _gluon_atomic_cas_adapter, + GluonBufferAtomicRMW: _gluon_buffer_atomic_rmw_adapter, + GluonIntToPtr: lambda _builder, *args, **kwargs: AdapterResult(*args, **kwargs), MakeRange: _gluon_make_range_adapter, Splat: lambda shape, value, *_args, **_kwargs: AdapterResult(shape, value), Allocate: _gluon_allocate_adapter, @@ -542,7 +631,10 @@ def _existing_ops(namespace: Any, attrs: dict[str, type[Op]]) -> dict[str, type[ descriptor_kwarg="src", coords_kwarg="offsets", ) - elif attr in _TMA_LOAD_PRED_ARG_INDICES: + elif ( + namespace in (gluon_hopper_tma, gluon_blackwell_tma) + and attr in _TMA_LOAD_PRED_ARG_INDICES + ): GLUON_CALLABLE_ADAPTERS[original] = _gluon_descriptor_load_adapter( _TMA_LOAD_PRED_ARG_INDICES[attr], descriptor_kwarg="tensor_desc",