[INFRA] Import NVIDIA/CCCL upstream as optimization reference library

CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

Competition optimization vectors mapped to CCCL:
- Output TPS (83% weight): warp_reduce, block_reduce, device_topk
- Input TPS (14% weight): device_scan, block_load, prefetch
- Cache TPS (3% weight): prefix caching strategy patterns
- Memory (0.9 util): pooled/cached/buddy allocators

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0

cccl_upstream/python/cuda_cccl/benchmarks/compute/select/if.py

@@ -0,0 +1,119 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+"""
+Python benchmark for select_if using cuda.compute.
+
+C++ equivalent: cub/benchmarks/bench/select/if.cu
+
+Notes:
+- The C++ benchmark uses a `less_then_t<T>` predicate with threshold based on entropy
+- Entropy controls what fraction of elements are selected:
+  - 1.000 → selects ~100% (threshold = max value)
+  - 0.544 → selects ~54.4% (threshold at 54.4% of range)
+  - 0.000 → selects ~0% (threshold = min value)
+- InPlace axis controls whether output can alias input (not exposed in Python API)
+- Migration: Python cannot expose InPlace axis; output is sized to num_elements but metrics use actual selected count.
+"""
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+import cupy as cp
+import numpy as np
+from utils import (
+    ENTROPY_TO_PROB,
+    as_cupy_stream,
+    generate_data_with_entropy,
+    lerp_min_max,
+)
+from utils import (
+    FUNDAMENTAL_TYPES as TYPE_MAP,
+)
+
+import cuda.bench as bench
+from cuda.compute import make_select
+
+# Entropy values from C++ benchmark
+# These control the selection threshold and thus how many elements are selected
+
+
+def bench_select_if(state: bench.State):
+    type_str = state.get_string("T{ct}")
+    dtype = TYPE_MAP[type_str]
+    num_elements = int(state.get_int64("Elements{io}"))
+    entropy_str = state.get_string("Entropy")
+
+    probability = ENTROPY_TO_PROB[entropy_str]
+    threshold = lerp_min_max(dtype, probability)
+
+    alloc_stream = as_cupy_stream(state.get_stream())
+
+    # Match C++ benchmark: input data generation is independent of Entropy.
+    # Entropy only controls the selection threshold.
+    d_in = generate_data_with_entropy(num_elements, dtype, "1.000", alloc_stream)
+    with alloc_stream:
+        selected_elements = int(cp.count_nonzero(d_in < threshold).get())
+        d_out = cp.empty(selected_elements, dtype=dtype)
+
+        d_num_selected = cp.zeros(1, dtype=np.int64)
+
+    alloc_stream.synchronize()
+
+    # Create predicate: select elements less than threshold
+    # For numba device functions, we need to use the value directly in closure
+    thresh_val = threshold
+
+    def less_than_threshold(x):
+        return x < thresh_val
+
+    selector = make_select(
+        d_in=d_in,
+        d_out=d_out,
+        d_num_selected_out=d_num_selected,
+        cond=less_than_threshold,
+    )
+
+    temp_storage_bytes = selector(
+        temp_storage=None,
+        d_in=d_in,
+        d_out=d_out,
+        d_num_selected_out=d_num_selected,
+        cond=less_than_threshold,
+        num_items=num_elements,
+    )
+    with alloc_stream:
+        temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
+
+    state.add_element_count(num_elements)
+    state.add_global_memory_reads(num_elements * d_in.dtype.itemsize)
+    state.add_global_memory_writes(selected_elements * d_out.dtype.itemsize)
+    state.add_global_memory_writes(1 * d_num_selected.dtype.itemsize)
+
+    def launcher(launch: bench.Launch):
+        selector(
+            temp_storage=temp_storage,
+            d_in=d_in,
+            d_out=d_out,
+            d_num_selected_out=d_num_selected,
+            cond=less_than_threshold,
+            num_items=num_elements,
+            stream=launch.get_stream(),
+        )
+
+    state.exec(launcher, batched=False)
+
+
+if __name__ == "__main__":
+    b = bench.register(bench_select_if)
+    b.set_name("base")
+
+    b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
+    b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
+    b.add_string_axis("Entropy", ["1.000", "0.544", "0.000"])
+    # Note: InPlace axis is not exposed in Python API, so we skip it
+
+    bench.run_all_benchmarks(sys.argv)

cccl_upstream/python/cuda_cccl/benchmarks/compute/select/unique_by_key.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+"""
+Python benchmark for unique_by_key using cuda.compute.
+
+C++ equivalent: cub/benchmarks/bench/select/unique_by_key.cu
+
+Notes:
+- The C++ benchmark uses MaxSegSize axis to control segment sizes
+- Uses equal_to comparison operator for key equality
+- Generates key segments with sizes between 1 and MaxSegSize
+- Both keys and values are processed
+- Migration: Python fixes offsets and generates key segments on GPU to mirror C++.
+- OffsetT axis is omitted because the Python API does not expose offset type.
+"""
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+import cupy as cp
+import numpy as np
+from utils import INTEGRAL_TYPES, SIGNED_TYPES, as_cupy_stream, generate_key_segments
+
+import cuda.bench as bench
+from cuda.compute import OpKind, make_unique_by_key
+
+KEY_TYPE_MAP = INTEGRAL_TYPES
+VALUE_TYPE_MAP = {**SIGNED_TYPES, "C32": np.complex64}
+
+
+def bench_unique_by_key(state: bench.State):
+    key_type_str = state.get_string("KeyT{ct}")
+    value_type_str = state.get_string("ValueT{ct}")
+    key_dtype = KEY_TYPE_MAP[key_type_str]
+    value_dtype = VALUE_TYPE_MAP[value_type_str]
+    num_elements = int(state.get_int64("Elements{io}"))
+    max_seg_size = int(state.get_int64("MaxSegSize"))
+
+    if num_elements > np.iinfo(np.int32).max:
+        state.skip("Skipping: num_elements exceeds int32 limits")
+        return
+
+    alloc_stream = as_cupy_stream(state.get_stream())
+
+    d_in_keys = generate_key_segments(
+        num_elements,
+        key_dtype,
+        min_segment_size=1,
+        max_segment_size=max_seg_size,
+        stream=alloc_stream,
+    )
+
+    with alloc_stream:
+        d_in_values = cp.zeros(num_elements, dtype=value_dtype)
+
+        d_out_keys = cp.empty(num_elements, dtype=key_dtype)
+        d_out_values = cp.empty(num_elements, dtype=value_dtype)
+        d_num_selected = cp.empty(1, dtype=np.int32)
+
+    alloc_stream.synchronize()
+
+    uniquer = make_unique_by_key(
+        d_in_keys=d_in_keys,
+        d_in_items=d_in_values,
+        d_out_keys=d_out_keys,
+        d_out_items=d_out_values,
+        d_out_num_selected=d_num_selected,
+        op=OpKind.EQUAL_TO,
+    )
+
+    temp_storage_bytes = uniquer(
+        temp_storage=None,
+        d_in_keys=d_in_keys,
+        d_in_items=d_in_values,
+        d_out_keys=d_out_keys,
+        d_out_items=d_out_values,
+        d_out_num_selected=d_num_selected,
+        op=OpKind.EQUAL_TO,
+        num_items=num_elements,
+    )
+    with alloc_stream:
+        temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
+
+    # Run once before timing to materialize the number of selected runs,
+    # matching the C++ metric accounting flow.
+    uniquer(
+        temp_storage=temp_storage,
+        d_in_keys=d_in_keys,
+        d_in_items=d_in_values,
+        d_out_keys=d_out_keys,
+        d_out_items=d_out_values,
+        d_out_num_selected=d_num_selected,
+        op=OpKind.EQUAL_TO,
+        num_items=num_elements,
+        stream=alloc_stream,
+    )
+    alloc_stream.synchronize()
+    num_runs = int(d_num_selected.get()[0])
+
+    state.add_element_count(num_elements)
+    state.add_global_memory_reads(int(num_elements * d_in_keys.dtype.itemsize))
+    state.add_global_memory_reads(int(num_elements * d_in_values.dtype.itemsize))
+    state.add_global_memory_writes(int(num_runs * d_out_keys.dtype.itemsize))
+    state.add_global_memory_writes(int(num_runs * d_out_values.dtype.itemsize))
+    state.add_global_memory_writes(int(d_num_selected.dtype.itemsize))
+
+    def launcher(launch: bench.Launch):
+        uniquer(
+            temp_storage=temp_storage,
+            d_in_keys=d_in_keys,
+            d_in_items=d_in_values,
+            d_out_keys=d_out_keys,
+            d_out_items=d_out_values,
+            d_out_num_selected=d_num_selected,
+            op=OpKind.EQUAL_TO,
+            num_items=num_elements,
+            stream=launch.get_stream(),
+        )
+
+    state.exec(launcher, batched=False)
+
+
+if __name__ == "__main__":
+    b = bench.register(bench_unique_by_key)
+    b.set_name("base")
+
+    b.add_string_axis("KeyT{ct}", list(KEY_TYPE_MAP.keys()))
+    b.add_string_axis("ValueT{ct}", list(VALUE_TYPE_MAP.keys()))
+    b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
+    b.add_int64_power_of_two_axis("MaxSegSize", [1, 4, 8])
+    # Note: OffsetT axis from C++ is not exposed in Python API
+
+    bench.run_all_benchmarks(sys.argv)