feat(cccl): integrate missing CCCL directories — python/, ci/, .agent/, docs/, test/
Sparse-checkout from NVIDIA/cccl main branch to complete cccl_upstream: Added: - python/cuda_cccl/ (226 files) — Python bindings for device-level algorithms Critical for muh toolchain: cuda.compute.reduce_into, scan, radix_sort, etc. Includes 204 .py files with full test coverage for all 27 algorithms - ci/ (163 files) — Build/test infrastructure build_cub.sh, test_cub.sh, build_and_test_targets.sh, matrix.yaml Directly maps to our [INFRA-CI] and [INFRA-BUILD] items - .agent/skills/ (7 files) — NVIDIA's own agent skills for CCCL cccl-style/SKILL.md, cccl-test/SKILL.md, sass-diff/SKILL.md - docs/ (491 files) — Official CCCL documentation CI references, CMake guides, Python compute docs, libcudacxx PTX docs - test/ (12 files) — Top-level integration tests (cuda_smoke, stdpar) - Root configs: .clang-format, .clang-tidy, CONTRIBUTING.md, pyproject.toml - CLAUDE.md symlink → AGENTS.md (NVIDIA's standard) cccl_upstream now mirrors full NVIDIA/cccl structure: Before: 42M (cub + thrust + libcudacxx + cudax + c + examples + benchmarks) After: 53M (+python +ci +docs +.agent +test +configs) This completes the CCCL base needed for: - [muh-bench] items: ci/util/build_and_test_targets.sh for targeted builds - [CCCL-verify] items: python/cuda_cccl/tests/ as reference implementations - [CCCL-test] items: ci/test_cub.sh, ci/test_thrust.sh - Agent workflow: .agent/skills/ for consistent style and test patterns
This commit is contained in:
muh-bot
252
cccl_upstream/python/cuda_cccl/benchmarks/compute/utils.py
Normal file
252
cccl_upstream/python/cuda_cccl/benchmarks/compute/utils.py
Normal file
@@ -0,0 +1,252 @@
|
||||
import cupy as cp
|
||||
import numpy as np
|
||||
|
||||
import cuda.bench as bench
|
||||
|
||||
ALL_TYPES = {
|
||||
"I8": np.int8,
|
||||
"I16": np.int16,
|
||||
"I32": np.int32,
|
||||
"I64": np.int64,
|
||||
"U8": np.uint8,
|
||||
"U16": np.uint16,
|
||||
"U32": np.uint32,
|
||||
"U64": np.uint64,
|
||||
"F32": np.float32,
|
||||
"F64": np.float64,
|
||||
}
|
||||
|
||||
SIGNED_TYPES = {k: ALL_TYPES[k] for k in ("I8", "I16", "I32", "I64", "F32", "F64")}
|
||||
FLOAT_TYPES = {k: ALL_TYPES[k] for k in ("F32", "F64")}
|
||||
|
||||
# Matches C++ integral_types = {int8_t, int16_t, int32_t, int64_t}
|
||||
INTEGRAL_TYPES = {k: ALL_TYPES[k] for k in ("I8", "I16", "I32", "I64")}
|
||||
|
||||
# Matches C++ fundamental_types = {int8..int64, [int128,] float, double}
|
||||
# int128 is excluded because it is not supported by numpy/cupy.
|
||||
FUNDAMENTAL_TYPES = {k: ALL_TYPES[k] for k in ("I8", "I16", "I32", "I64", "F32", "F64")}
|
||||
|
||||
ENTROPY_TO_STEPS = {
|
||||
"1.000": 0,
|
||||
"0.811": 1,
|
||||
"0.544": 2,
|
||||
"0.337": 3,
|
||||
"0.201": 4,
|
||||
"0.000": 0,
|
||||
}
|
||||
|
||||
ENTROPY_TO_PROB = {
|
||||
"1.000": 1.0,
|
||||
"0.811": 0.811,
|
||||
"0.544": 0.544,
|
||||
"0.337": 0.337,
|
||||
"0.201": 0.201,
|
||||
"0.000": 0.0,
|
||||
}
|
||||
|
||||
|
||||
def as_cupy_stream(cs: bench.CudaStream) -> cp.cuda.Stream:
|
||||
"""Convert nvbench CudaStream to CuPy Stream."""
|
||||
return cp.cuda.ExternalStream(cs.addressof())
|
||||
|
||||
|
||||
def lerp_min_max(dtype, probability):
|
||||
"""Interpolate between min/max for dtype like nvbench_helper.cuh."""
|
||||
if probability == 1.0:
|
||||
if np.issubdtype(dtype, np.integer):
|
||||
return np.iinfo(dtype).max
|
||||
return np.finfo(dtype).max
|
||||
|
||||
if np.issubdtype(dtype, np.integer):
|
||||
min_val = float(np.iinfo(dtype).min)
|
||||
max_val = float(np.iinfo(dtype).max)
|
||||
else:
|
||||
min_val = float(np.finfo(dtype).min)
|
||||
max_val = float(np.finfo(dtype).max)
|
||||
|
||||
return dtype(min_val + probability * (max_val - min_val))
|
||||
|
||||
|
||||
def _bitwise_and(a, b, dtype):
|
||||
if np.issubdtype(dtype, np.floating):
|
||||
view_dtype = cp.uint32 if dtype == np.float32 else cp.uint64
|
||||
return (a.view(view_dtype) & b.view(view_dtype)).view(dtype)
|
||||
return a & b
|
||||
|
||||
|
||||
def _uniform_random(num_elements, dtype, min_val, max_val):
|
||||
rand = cp.random.random(num_elements)
|
||||
if np.issubdtype(dtype, np.floating):
|
||||
return ((float(max_val) - float(min_val)) * rand + float(min_val)).astype(dtype)
|
||||
min_f = float(min_val)
|
||||
max_f = float(max_val)
|
||||
return cp.floor((max_f - min_f + 1) * rand + min_f).astype(dtype)
|
||||
|
||||
|
||||
def generate_data_with_entropy(
|
||||
num_elements, dtype, entropy_str, stream, min_val=None, max_val=None
|
||||
):
|
||||
"""Generate data with nvbench_helper-style bit entropy."""
|
||||
if min_val is None or max_val is None:
|
||||
if np.issubdtype(dtype, np.integer):
|
||||
info = np.iinfo(dtype)
|
||||
default_min = info.min
|
||||
default_max = info.max
|
||||
else:
|
||||
info = np.finfo(dtype)
|
||||
default_min = info.tiny
|
||||
default_max = info.max
|
||||
min_val = default_min if min_val is None else min_val
|
||||
max_val = default_max if max_val is None else max_val
|
||||
|
||||
steps = ENTROPY_TO_STEPS[entropy_str]
|
||||
|
||||
with stream:
|
||||
if entropy_str == "0.000":
|
||||
scalar = _uniform_random(1, dtype, min_val, max_val)[0]
|
||||
data = cp.full(num_elements, scalar, dtype=dtype)
|
||||
else:
|
||||
data = _uniform_random(num_elements, dtype, min_val, max_val)
|
||||
for _ in range(steps):
|
||||
tmp = _uniform_random(num_elements, dtype, min_val, max_val)
|
||||
data = _bitwise_and(data, tmp, dtype)
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def generate_uniform_segment_offsets(num_elements, min_segment_size, max_segment_size):
|
||||
num_elements = int(num_elements)
|
||||
if min_segment_size <= 0:
|
||||
raise ValueError("min_segment_size must be positive")
|
||||
if max_segment_size < min_segment_size:
|
||||
raise ValueError("max_segment_size must be >= min_segment_size")
|
||||
|
||||
num_segments_est = int(np.ceil(num_elements / min_segment_size))
|
||||
if num_segments_est > np.iinfo(np.int32).max:
|
||||
raise MemoryError("Too many segments for int32 offsets")
|
||||
|
||||
sizes = cp.random.randint(
|
||||
min_segment_size,
|
||||
max_segment_size + 1,
|
||||
size=num_segments_est,
|
||||
dtype=cp.int64,
|
||||
)
|
||||
cumsum = cp.cumsum(sizes)
|
||||
cutoff = int(
|
||||
cp.searchsorted(
|
||||
cumsum, cp.asarray(num_elements, dtype=cp.int64), side="left"
|
||||
).item()
|
||||
)
|
||||
sizes = sizes[: cutoff + 1]
|
||||
prev = 0 if cutoff == 0 else int(cumsum[cutoff - 1].item())
|
||||
sizes[cutoff] = num_elements - prev
|
||||
|
||||
offsets = cp.empty(cutoff + 2, dtype=cp.int64)
|
||||
offsets[0] = 0
|
||||
offsets[1:] = cp.cumsum(sizes)
|
||||
offsets[-1] = num_elements
|
||||
return offsets
|
||||
|
||||
|
||||
def generate_power_law_offsets(num_elements, num_segments):
|
||||
if num_segments <= 0:
|
||||
return cp.asarray([0, num_elements], dtype=cp.int64)
|
||||
|
||||
# Mirror nvbench_helper power-law generation:
|
||||
# draw log-normal samples, normalize to total elements,
|
||||
# floor to integer segment sizes, then distribute remainder
|
||||
# across the first `diff` segments.
|
||||
samples = cp.random.lognormal(3.0, 1.2, size=num_segments)
|
||||
if int(cp.count_nonzero(samples).item()) == 0:
|
||||
samples = cp.ones(num_segments, dtype=cp.float64)
|
||||
|
||||
sample_sum = float(samples.sum().item())
|
||||
sizes = cp.floor(samples * num_elements / sample_sum).astype(cp.int64)
|
||||
|
||||
diff = int(num_elements - sizes.sum().item())
|
||||
if diff > 0:
|
||||
sizes[:diff] += 1
|
||||
|
||||
offsets = cp.empty(num_segments + 1, dtype=cp.int64)
|
||||
offsets[0] = 0
|
||||
offsets[1:] = cp.cumsum(sizes)
|
||||
return offsets
|
||||
|
||||
|
||||
def generate_fixed_segment_offsets(num_elements, segment_size, stream):
|
||||
num_segments = max(1, num_elements // segment_size)
|
||||
actual_elements = num_segments * segment_size
|
||||
|
||||
with stream:
|
||||
start_offsets = cp.arange(0, actual_elements, segment_size, dtype=np.int64)
|
||||
end_offsets = cp.arange(
|
||||
segment_size, actual_elements + 1, segment_size, dtype=np.int64
|
||||
)
|
||||
|
||||
return start_offsets, end_offsets, num_segments, actual_elements
|
||||
|
||||
|
||||
def generate_key_segments(
|
||||
num_elements, key_dtype, min_segment_size, max_segment_size, stream
|
||||
):
|
||||
"""Generate GPU key segments (runs of equal keys) matching C++ generate.uniform.key_segments.
|
||||
|
||||
All computation stays on GPU via CuPy. We avoid ``cp.repeat`` (which
|
||||
doesn't accept a device array for *repeats*) by building a segment-id
|
||||
array through ``cp.searchsorted`` on cumulative offsets instead.
|
||||
"""
|
||||
num_elements = int(num_elements)
|
||||
if min_segment_size <= 0:
|
||||
raise ValueError("min_segment_size must be positive")
|
||||
if max_segment_size < min_segment_size:
|
||||
raise ValueError("max_segment_size must be >= min_segment_size")
|
||||
|
||||
num_segments_est = int(np.ceil(num_elements / min_segment_size))
|
||||
if num_segments_est > np.iinfo(np.int32).max:
|
||||
raise MemoryError("Too many segments for int32 offsets")
|
||||
|
||||
with stream:
|
||||
sizes = cp.random.randint(
|
||||
min_segment_size,
|
||||
max_segment_size + 1,
|
||||
size=num_segments_est,
|
||||
dtype=cp.int64,
|
||||
)
|
||||
cumsum = cp.cumsum(sizes)
|
||||
|
||||
# Find how many full segments fit within num_elements
|
||||
cutoff = int(
|
||||
cp.searchsorted(
|
||||
cumsum, cp.asarray(num_elements, dtype=cp.int64), side="left"
|
||||
).item()
|
||||
)
|
||||
sizes = sizes[: cutoff + 1]
|
||||
prev = 0 if cutoff == 0 else int(cumsum[cutoff - 1].item())
|
||||
sizes[cutoff] = num_elements - prev
|
||||
|
||||
# Build cumulative offsets for the final segments
|
||||
offsets = cp.empty(cutoff + 2, dtype=cp.int64)
|
||||
offsets[0] = 0
|
||||
offsets[1:] = cp.cumsum(sizes)
|
||||
|
||||
# Instead of cp.repeat (which doesn't support device repeats),
|
||||
# use searchsorted to map each element index to its segment id.
|
||||
indices = cp.arange(num_elements, dtype=cp.int64)
|
||||
# searchsorted(offsets[1:], indices, side="right") gives the segment id
|
||||
segment_ids = cp.searchsorted(offsets[1:], indices, side="right")
|
||||
|
||||
# Map segment ids to key values, wrapping within dtype range
|
||||
if np.issubdtype(key_dtype, np.integer):
|
||||
info = np.iinfo(key_dtype)
|
||||
if np.dtype(key_dtype).itemsize < 8:
|
||||
range_size = int(info.max) - int(info.min) + 1
|
||||
keys = ((segment_ids % range_size) + int(info.min)).astype(
|
||||
key_dtype, copy=False
|
||||
)
|
||||
else:
|
||||
# For int64, avoid Python overflow: just cast directly
|
||||
keys = segment_ids.astype(key_dtype, copy=False)
|
||||
else:
|
||||
keys = segment_ids.astype(key_dtype, copy=False)
|
||||
|
||||
return keys
|
||||
Reference in New Issue
Block a user