Update masked_select and nonzero

benchmark/performance_utils.py

@@ -258,7 +258,11 @@ def get_latency(self, op, *args, **kwargs):
         if self.is_backward:
             out = fn()
             dout = torch.randn_like(out)
-            fn = lambda: out.backward(dout, retain_graph=True)
+            # fn = lambda: out.backward(dout, retain_graph=True)
+            xs = list(filter(lambda x: torch.is_tensor(x) and x.requires_grad, args))
+            fn = lambda: torch.autograd.grad(
+                (out,), xs, grad_outputs=(dout,), retain_graph=True
+            )
         if Config.cpu_mode:
             for i in range(Config.warm_up):
                 fn()
@@ -280,6 +284,7 @@ def get_latency(self, op, *args, **kwargs):
                 warmup=Config.warm_up,
                 rep=Config.repetition,
                 return_mode="median",
+                grad_to_none=xs if self.is_backward else None,
             )
         # average latency in ms
         return latency

benchmark/test_select_and_slice_perf.py

@@ -66,6 +66,7 @@ def index_select_gbps(bench_fn_args, latency):
     return io_amount * 1e-9 / (latency * 1e-3)
 
 
+@pytest.mark.index_select
 @pytest.mark.parametrize(
     "op_name, torch_op, input_fn, gbps_fn, dtypes",
     [
@@ -77,6 +78,23 @@ def index_select_gbps(bench_fn_args, latency):
             FLOAT_DTYPES,
             marks=pytest.mark.index_select,
         ),
+    ],
+)
+def test_perf_index_select(op_name, torch_op, input_fn, gbps_fn, dtypes):
+    bench = TensorSelectBenchmark(
+        input_fn=input_fn,
+        op_name=op_name,
+        torch_op=torch_op,
+        dtypes=dtypes,
+        get_gbps=gbps_fn,
+    )
+    bench.run()
+
+
+@pytest.mark.masked_select
+@pytest.mark.parametrize(
+    "op_name, torch_op, input_fn, gbps_fn, dtypes",
+    [
         pytest.param(
             "masked_select",
             torch.masked_select,
@@ -87,7 +105,7 @@ def index_select_gbps(bench_fn_args, latency):
         ),
     ],
 )
-def test_generic_reduction_benchmark(op_name, torch_op, input_fn, gbps_fn, dtypes):
+def test_perf_masked_select(op_name, torch_op, input_fn, gbps_fn, dtypes):
     bench = TensorSelectBenchmark(
         input_fn=input_fn,
         op_name=op_name,

src/flag_gems/ops/masked_select.py

@@ -4,34 +4,125 @@
 import triton
 import triton.language as tl
 
-from flag_gems import runtime
 from flag_gems.runtime import torch_device_fn
 from flag_gems.utils import broadcastable, libentry
-from flag_gems.utils import triton_lang_extension as tle
+from flag_gems.utils.shape_utils import bracket_next_power_of_2
 
 logger = logging.getLogger(__name__)
 
 
 @libentry()
-@triton.autotune(configs=runtime.get_tuned_config("masked_select"), key=["n_elements"])
 @triton.jit
-def masked_select_kernel(
+def masked_select_single_pass_kernel(
+    inp_ptr, mask_ptr, out_ptr, N, BLOCK_SIZE: tl.constexpr
+):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    inp = tl.load(inp_ptr + offsets, mask=offsets < N)
+    mask = tl.load(mask_ptr + offsets, mask=offsets < N).to(tl.int1)
+    mask_ints = mask.to(tl.int32)
+    out_offsets = tl.cumsum(mask_ints, axis=0) - 1
+
+    tl.store(out_ptr + out_offsets, inp, mask=offsets < N and mask)
+
+
+def masked_select_single_pass(inp, mask, out, N):
+    BLOCK_SIZE = triton.next_power_of_2(N)
+    if BLOCK_SIZE <= 512:
+        num_warps = 4
+    elif BLOCK_SIZE <= 2048:
+        num_warps = 8
+    else:
+        num_warps = 16
+    masked_select_single_pass_kernel[(1,)](
+        inp, mask, out, N, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps
+    )
+    return out
+
+
+@libentry()
+@triton.jit(do_not_specialize=["N", "nr", "row_stride"])
+def mask_part_sum_kernel(
     inp_ptr,
-    select_mask_ptr,
-    prefix_sum_ptr,
+    mask_ptr,
+    part_sums_ptr,
+    counter_ptr,
+    N,
+    num_blocks,
+    num_blocks_per_row,
+    NP_BLOCK: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    start_block = row_id * num_blocks_per_row
+    offset = start_block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    acc = tl.zeros((BLOCK_SIZE,), dtype=part_sums_ptr.dtype.element_ty)
+
+    last_block_id = min(num_blocks - 1, start_block + num_blocks_per_row - 1)
+
+    for block_id in range(start_block, last_block_id):
+        select = tl.load(mask_ptr + offset)
+        select_ints = select.to(part_sums_ptr.dtype.element_ty)
+        acc += select_ints
+        offset += BLOCK_SIZE
+    # Peeled last block
+    select = tl.load(mask_ptr + offset, mask=offset < N, other=0)
+    select_ints = select.to(part_sums_ptr.dtype.element_ty)
+    acc += select_ints
+
+    part_sum = tl.sum(acc, axis=0)
+    tl.store(part_sums_ptr + row_id, part_sum)
+    # cumsum the part_sums
+    count = tl.atomic_add(counter_ptr, 1, sem="acq_rel")
+    np = tl.num_programs(0)
+    if count == np - 1:
+        mask = tl.arange(0, NP_BLOCK) < np
+        part_sums = tl.load(part_sums_ptr + tl.arange(0, NP_BLOCK), mask=mask)
+        final_sum = tl.sum(part_sums, axis=0)
+        pre_sums = tl.cumsum(part_sums, axis=0)
+        tl.store(
+            part_sums_ptr + tl.arange(0, NP_BLOCK), pre_sums - part_sums, mask=mask
+        )
+        tl.store(part_sums_ptr + np, final_sum)
+
+
+@libentry()
+@triton.jit(do_not_specialize=["N", "nr", "row_stride"])
+def write_back_kernel(
+    inp_ptr,
+    mask_ptr,
+    part_sums_ptr,
     out_ptr,
-    n_elements,
+    N,
+    num_blocks,
+    num_blocks_per_row,
+    NP_BLOCK: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
-    pid = tle.program_id(axis=0)
-    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    mask = offsets < n_elements
+    row_id = tl.program_id(0)
 
-    inp = tl.load(inp_ptr + offsets, mask=mask, other=0.0)
-    select_mask = tl.load(select_mask_ptr + offsets, mask=mask, other=0.0).to(tl.int1)
-    out_offset = tl.load(prefix_sum_ptr + offsets, mask=mask, other=0.0) - 1
+    start_block = row_id * num_blocks_per_row
+    offset = start_block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    advance = tl.load(part_sums_ptr + row_id)
 
-    tl.store(out_ptr + out_offset, inp, mask=(select_mask and mask))
+    last_block_id = min(num_blocks - 1, start_block + num_blocks_per_row - 1)
+
+    for block_id in range(start_block, last_block_id):
+        inp = tl.load(inp_ptr + offset)
+        select_mask = tl.load(mask_ptr + offset).to(tl.int1)
+        select_ints = select_mask.to(tl.constexpr(part_sums_ptr.dtype.element_ty))
+        out_ptr += advance
+        advance = tl.sum(select_ints, axis=0)
+        pre_sums = tl.cumsum(select_ints, axis=0) - 1
+        tl.store(out_ptr + pre_sums, inp, mask=select_mask)
+        offset += BLOCK_SIZE
+    # Peeled last block
+    inp = tl.load(inp_ptr + offset, mask=offset < N)
+    select_mask = tl.load(mask_ptr + offset, mask=offset < N, other=0).to(tl.int1)
+    select_ints = select_mask.to(tl.constexpr(part_sums_ptr.dtype.element_ty))
+    out_ptr += advance
+    pre_sums = tl.cumsum(select_ints, axis=0) - 1
+    tl.store(out_ptr + pre_sums, inp, mask=offset < N and select_mask)
 
 
 def masked_select(inp, mask):
@@ -48,13 +139,58 @@ def masked_select(inp, mask):
     inp = inp.contiguous()
     mask = mask.contiguous()
 
-    mask_flattened = mask.ravel()
+    N = inp.numel()
+    if N <= 4096:
+        out = torch.empty(mask.sum(), dtype=inp.dtype, device=inp.device)
+        return masked_select_single_pass(inp, mask, out, N)
+
+    # return mask_select(inp, mask)
+
+    BLOCK_SIZE = bracket_next_power_of_2(N, 128, 4096)
+    num_warps = min(16, BLOCK_SIZE // 32)
 
-    prefix_sum = mask_flattened.cumsum(axis=0)
-    out = torch.empty(prefix_sum[-1].item(), dtype=inp.dtype, device=inp.device)
+    # max degree of parallelism
+    np = torch_device_fn.get_device_properties(mask.device).multi_processor_count
+
+    # arranged as np rows of blocks
+    n_blocks = triton.cdiv(N, BLOCK_SIZE)
+    np = min(n_blocks, np)
+    n_blocks_per_row = triton.cdiv(n_blocks, np)
+    np = triton.cdiv(n_blocks, n_blocks_per_row)
+    NP_BLOCK = triton.next_power_of_2(np)
 
-    n_elements = inp.numel()
-    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
     with torch_device_fn.device(inp.device):
-        masked_select_kernel[grid](inp, mask_flattened, prefix_sum, out, n_elements)
+        # Compute per cta sums and cumulative sums across ctas
+        dtype = torch.int32 if N < 2**31 else torch.int64
+        part_sums = torch.empty(np + 1, dtype=dtype, device=mask.device)
+        barrier = torch.zeros([], dtype=torch.int, device=mask.device)
+        mask_part_sum_kernel[(np,)](
+            inp,
+            mask,
+            part_sums,
+            barrier,
+            N,
+            n_blocks,
+            n_blocks_per_row,
+            NP_BLOCK=NP_BLOCK,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+
+        # Write back selected data
+        out = torch.empty(part_sums[-1], dtype=inp.dtype, device=mask.device)
+        # write_offsets = pre_sums - part_sums
+        write_back_kernel[(np,)](
+            inp,
+            mask,
+            part_sums,
+            out,
+            N,
+            n_blocks,
+            n_blocks_per_row,
+            NP_BLOCK=triton.next_power_of_2(np),
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+
     return out

src/flag_gems/ops/nonzero.py

@@ -13,12 +13,7 @@
 
 
 @libentry()
-@triton.autotune(
-    configs=runtime.get_tuned_config("nonzero"),
-    key=[
-        "n_elements",
-    ],
-)
+@triton.heuristics(runtime.get_heuristic_config("elementwise_generic"))
 @triton.jit
 def nonzero_kernel(
     inp,
@@ -34,10 +29,10 @@ def nonzero_kernel(
     offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     mask = offset < n_elements
 
-    inp_vals = tl.load(inp + offset, mask=mask)
+    inp_vals = tl.load(inp + offset, mask=mask).to(tl.int1)
     out_offset = tl.load(prefix_sum + offset, mask=mask) - 1
 
-    nonzero_mask = mask and inp_vals == True  # noqa
+    nonzero_mask = mask and inp_vals  # noqa
 
     idx_flat = offset
     for dim in range(ndim - 1, -1, -1):

src/flag_gems/runtime/backend/_nvidia/heuristics_config_utils.py

@@ -2,6 +2,10 @@
 import triton
 
 
+def simple_elementwise_blocksize_heur(args):
+    return 1024
+
+
 def argmax_heur_block_m(args):
     return 4 if args["M"] < 4096 else 8
 
@@ -302,4 +306,8 @@ def vdot_heur_block_size(args):
     "vdot": {
         "BLOCK_SIZE": vdot_heur_block_size,
     },
+    "elementwise_generic": {
+        "BLOCK_SIZE": simple_elementwise_blocksize_heur,
+        "num_warps": lambda args: 8,
+    },
 }

src/flag_gems/utils/shape_utils.py

@@ -16,6 +16,10 @@
 Perm = Tuple[int]
 
 
+def bracket_next_power_of_2(N, lower, upper):
+    return min(max(triton.next_power_of_2(N), lower), upper)
+
+
 def broadcast(s1: Shape, s2: Shape) -> Shape:
     _s1, _s2 = s1, s2
     r1 = len(s1)