[FIX] Accept NumPy-array masks in the profiler's load/store statistics

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(

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