From 0d013a68a251fc2434a61e24aa3566dcc940b068 Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Tue, 10 Nov 2026 18:33:41 -0500 Subject: [PATCH] [FIX] Let the profiler run on fake tensors with atomics and reinterpreted tensors With SANITIZER_ENABLE_FAKE_TENSOR=1, kernel arguments become FakeTensors, which carry the original tensor's address but no storage. Two things broke on them: - The profiler skips loads, stores, and dots, but atomics still ran in Triton's interpreter, which dereferenced the original address on the host: a GPU tensor's address segfaults the process, and a CPU tensor is written. Under fake tensors the profiler now overrides create_atomic_rmw and create_atomic_cas like loads, returning zeros without touching memory. With real tensors, atomics still run. - An argument wrapped by triton.reinterpret is rewrapped around its FakeTensor, and TensorWrapper.__init__ reads base.data, which FakeTensor lacked. FakeTensor.data now returns the FakeTensor itself. Tests: test_fake_tensor_atomics_do_not_touch_memory and test_fake_tensor_accepts_reinterpreted_tensor in tests/end_to_end/test_profiler.py (both fail without the fix). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_0169QRxKtsivkLDCS5f7rrsm --- tests/end_to_end/test_profiler.py | 68 +++++++++++++++++++++++++++ tilelens/clients/profiler/profiler.py | 20 +++++++- tilelens/core/frontend/triton.py | 5 ++ 3 files changed, 92 insertions(+), 1 deletion(-) diff --git a/tests/end_to_end/test_profiler.py b/tests/end_to_end/test_profiler.py index 04726f169..84901f8f7 100644 --- a/tests/end_to_end/test_profiler.py +++ b/tests/end_to_end/test_profiler.py @@ -421,3 +421,71 @@ def test_load_store_skip_enabled(_isolate_profiler_cfg): assert torch.allclose( y, torch.zeros_like(y) ), f"Skipped execution should leave output unchanged, got {y[:10]}" + + +# ======== Fake Tensors ======== +@pytest.fixture +def _isolate_virtual_memory(): + """Save and restore config.virtual_memory around a test.""" + saved = cfg.virtual_memory + yield + cfg.virtual_memory = saved + + +@triton.jit +def fake_tensor_atomic_kernel(x_ptr, counter_ptr, lock_ptr, BLOCK_SIZE: tl.constexpr): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + tl.load(x_ptr + offsets) + tl.atomic_add(counter_ptr, 1.0) + tl.atomic_cas(lock_ptr, 0, 7) + + +def test_fake_tensor_atomics_do_not_touch_memory( + _isolate_profiler_cfg, _isolate_virtual_memory +): + """Fake tensors have no storage, so the profiler must not perform atomics through them. + + On a GPU tensor, an atomic through a fake tensor dereferences a device address on the + host and crashes; on a CPU tensor it writes into the caller's tensor. + """ + BLOCK_SIZE = 16 + x = torch.ones(4 * BLOCK_SIZE, dtype=torch.float32) + counter = torch.zeros(1, dtype=torch.float32) + lock = torch.zeros(1, dtype=torch.int32) + + cfg.virtual_memory = True + cfg.profiler_enable_block_sampling = False + cfg.profiler_disable_buffer_load_check = True + profiler = Profiler(disable_load_mask_percentage_check=True) + tilelens.trace(profiler)(fake_tensor_atomic_kernel)[(4,)]( + x, counter, lock, BLOCK_SIZE + ) + + assert counter.item() == 0.0 + assert lock.item() == 0 + assert profiler.load_bytes.total_bytes_attempted == 4 * BLOCK_SIZE * 4 + + +@triton.jit +def fake_tensor_copy_kernel(x_ptr, y_ptr, BLOCK_SIZE: tl.constexpr): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + tl.store(y_ptr + offsets, tl.load(x_ptr + offsets).to(tl.float32)) + + +def test_fake_tensor_accepts_reinterpreted_tensor( + _isolate_profiler_cfg, _isolate_virtual_memory +): + """triton.reinterpret wraps a tensor in a TensorWrapper, which reads the tensor's .data.""" + BLOCK_SIZE = 16 + x = torch.zeros(2 * BLOCK_SIZE, dtype=torch.uint8) + y = torch.zeros(2 * BLOCK_SIZE, dtype=torch.float32) + + cfg.virtual_memory = True + cfg.profiler_enable_block_sampling = False + cfg.profiler_disable_buffer_load_check = True + profiler = Profiler(disable_load_mask_percentage_check=True) + x_f8 = triton.reinterpret(x, tl.float8e5) + tilelens.trace(profiler)(fake_tensor_copy_kernel)[(2,)](x_f8, y, BLOCK_SIZE) + + assert profiler.load_bytes.total_bytes_attempted == 2 * BLOCK_SIZE + assert profiler.store_bytes.total_bytes_attempted == 2 * BLOCK_SIZE * 4 diff --git a/tilelens/clients/profiler/profiler.py b/tilelens/clients/profiler/profiler.py index e95723549..a894c0492 100644 --- a/tilelens/clients/profiler/profiler.py +++ b/tilelens/clients/profiler/profiler.py @@ -2,7 +2,7 @@ from ...core.client import Client from ...core.callbacks import OpCallbacks, ForLoopCallbacks -from ...core.data import Op, Load, Store, AddPtr, Dot +from ...core.data import Op, Load, Store, AddPtr, Dot, AtomicCas, AtomicRMW from ...core.config import config as cfg from ...core.patch import LoopSite from .data import LoadStoreBytes @@ -289,6 +289,20 @@ def dot_overrider(a, b, d, input_precision, max_num_imprecise_acc): # This replaces np.matmul(a_data, b_data, dtype=d.data.dtype) + d.data return TensorHandle(np.zeros_like(d.data), d.dtype.scalar) + def atomic_rmw_overrider(rmwOp, ptr, val, mask, sem, scope): + # Fake tensors have no storage: skip the access, return zeros + dtype_tt = ptr.get_element_ty() + return TensorHandle( + np.zeros_like(ptr.data, dtype=_get_np_dtype(dtype_tt)), dtype_tt + ) + + def atomic_cas_overrider(ptr, cmp, val, sem, scope): + # Fake tensors have no storage: skip the access, return zeros + dtype_tt = ptr.get_element_ty() + return TensorHandle( + np.zeros_like(ptr.data, dtype=_get_np_dtype(dtype_tt)), dtype_tt + ) + @self.lock_fn def pre_addptr_callback(ptr, offset): dtype_tt = ptr.get_element_ty() @@ -326,6 +340,10 @@ def pre_addptr_callback(ptr, offset): return OpCallbacks(op_overrider=dot_overrider) elif op_type is AddPtr: return OpCallbacks(before_callback=pre_addptr_callback) + elif op_type is AtomicRMW and cfg.virtual_memory: + return OpCallbacks(op_overrider=atomic_rmw_overrider) + elif op_type is AtomicCas and cfg.virtual_memory: + return OpCallbacks(op_overrider=atomic_cas_overrider) return OpCallbacks() diff --git a/tilelens/core/frontend/triton.py b/tilelens/core/frontend/triton.py index 30b4846a9..3ba0e9fdc 100644 --- a/tilelens/core/frontend/triton.py +++ b/tilelens/core/frontend/triton.py @@ -231,6 +231,11 @@ def numel(self) -> int: def element_size(self) -> int: return self._element_size + @property + def data(self) -> "FakeTensor": + # triton.reinterpret wraps tensors in a TensorWrapper, which reads .data + return self + class TritonFrontend(Frontend): def __init__(self):