From 17fdca1ef484202513afce2b427dd525e0c36e0b Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Sun, 8 Nov 2026 13:10:07 -0500 Subject: [PATCH] [FIX] Accept NumPy-array masks in the profiler's load/store statistics A tensor-descriptor load or store (and, on Triton 3.7, a block-pointer one) reaches the masked load/store with its bounds mask as a plain NumPy array rather than a TensorHandle. The profiler read `mask.data.size`, and an array's `.data` is a memoryview, so every such access raised AttributeError: 'memoryview' object has no attribute 'size'. `_mask_data` now returns the array itself, and both the mask statistics and the byte counts use it. Co-Authored-By: Claude Opus 5.5 --- tests/end_to_end/test_profiler.py | 38 +++++++++++++++++++++++++++ tilelens/clients/profiler/profiler.py | 25 +++++++++++++----- 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/tests/end_to_end/test_profiler.py b/tests/end_to_end/test_profiler.py index 04726f169..275890a97 100644 --- a/tests/end_to_end/test_profiler.py +++ b/tests/end_to_end/test_profiler.py @@ -259,6 +259,44 @@ def test_mask_percentage(_isolate_profiler_cfg): ), f"Expected store masked percentage {expected_store_masked_percentage:.2f}%, got {actual_store_masked_percentage:.2f}%" +@triton.jit +def descriptor_copy_kernel(in_ptr, out_ptr, N, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + src = tl.make_tensor_descriptor( + in_ptr, shape=[N], strides=[1], block_shape=[BLOCK_SIZE] + ) + dst = tl.make_tensor_descriptor( + out_ptr, shape=[N], strides=[1], block_shape=[BLOCK_SIZE] + ) + dst.store([pid * BLOCK_SIZE], src.load([pid * BLOCK_SIZE])) + + +def test_mask_percentage_of_tensor_descriptor_accesses(_isolate_profiler_cfg): + """A descriptor load/store passes its bounds mask as a NumPy array, not a TensorHandle.""" + N = 100 + BLOCK_SIZE = 32 + num_blocks = triton.cdiv(N, BLOCK_SIZE) + x = torch.randn(N, dtype=torch.float32) + y = torch.empty(N, dtype=torch.float32) + + cfg.profiler_enable_block_sampling = False + cfg.profiler_disable_buffer_load_check = True + profiler = Profiler() + traced_kernel = tilelens.trace(profiler)(descriptor_copy_kernel) + traced_kernel[(num_blocks,)](x, y, N, BLOCK_SIZE) + + # 4 blocks of 32 elements; the last block has 4 in bounds, so 28 are masked. + assert profiler.load_mask_total_count == 128 + assert profiler.load_mask_false_count == 28 + assert profiler.store_mask_total_count == 128 + assert profiler.store_mask_false_count == 28 + # Bytes: 100 float32 elements in bounds out of 128 attempted. + assert profiler.load_bytes.total_bytes_attempted == 128 * 4 + assert profiler.load_bytes.total_bytes_true == 100 * 4 + assert profiler.store_bytes.total_bytes_attempted == 128 * 4 + assert profiler.store_bytes.total_bytes_true == 100 * 4 + + # ======== Block Sampling ======== @triton.jit def block_sampling_test_kernel( diff --git a/tilelens/clients/profiler/profiler.py b/tilelens/clients/profiler/profiler.py index e95723549..2f0828809 100644 --- a/tilelens/clients/profiler/profiler.py +++ b/tilelens/clients/profiler/profiler.py @@ -15,6 +15,15 @@ from dataclasses import dataclass, replace +def _mask_data(mask: TensorHandle | np.ndarray) -> np.ndarray: + """The boolean mask of a load or store. + + Tensor-descriptor and block-pointer accesses pass their bounds mask to the + masked load/store as a plain NumPy array, whose ``.data`` is a memoryview. + """ + return mask if isinstance(mask, np.ndarray) else mask.data + + @dataclass(frozen=False) class LoopInfo: length: int | None = None @@ -157,12 +166,13 @@ def grid_callback(self, grid: tuple[int, ...]): self.sampled_blocks = None def _report_load_store_bytes( - self, type, ptr: TensorHandle, mask: TensorHandle + self, type, ptr: TensorHandle, mask: TensorHandle | np.ndarray ): # internal methods assumed to be called under the lock dtype_tt = ptr.get_element_ty() dtype_np: np.dtype = _get_np_dtype(dtype_tt) - mask_true = np.count_nonzero(mask.data) - mask_false = np.count_nonzero(np.logical_not(mask.data)) + mask_data = _mask_data(mask) + mask_true = np.count_nonzero(mask_data) + mask_false = np.count_nonzero(np.logical_not(mask_data)) total_bytes_true = mask_true * dtype_np.itemsize total_bytes_attempted = (mask_true + mask_false) * dtype_np.itemsize if type == "load": @@ -200,17 +210,18 @@ def _check_32bit_range( self.potential_buffer_load_issue_found = True def register_op_callback(self, op_type: type[Op]) -> OpCallbacks: - def _get_mask_stats(mask: TensorHandle) -> tuple[int, int]: + def _get_mask_stats(mask: TensorHandle | np.ndarray) -> tuple[int, int]: """Get mask statistics: total count and false count. Args: - mask: TensorHandle containing boolean mask data + mask: TensorHandle or NumPy array containing boolean mask data Returns: Tuple of (total_count, false_count) """ - total_count = mask.data.size - false_count = np.count_nonzero(np.logical_not(mask.data)) + mask_data = _mask_data(mask) + total_count = mask_data.size + false_count = np.count_nonzero(np.logical_not(mask_data)) return total_count, false_count @self.lock_fn