test: fix CUB compat tests + probe corex built-in CUB
v1 failures analyzed:
1. __shfl_down_sync: load_inline bug (not hardware issue) — fixed with file-based compile
2. CCCL 3.6: 'CUDA < 12 not supported' + variadic functions in device code
→ corex reports CUDA 10.2, corex clang rejects variadic in device code
v2 tests:
[1/4] __shfl_down_sync — does ivcore10 support warp shuffle?
[2/4] manual block reduce — SMEM + shuffle without CUB
[3/4] cub::BlockReduce using COREX's own CUB (/usr/local/corex/include/cub/)
[4/4] cub::BlockReduce using CCCL 3.6 with CCCL_IGNORE_DEPRECATED_CUDA_BELOW_12
probe_corex_cub.sh — discovers corex CUB version, compiles standalone test
Run on real machine:
bash qwen3_6_scripts/probe_corex_cub.sh
python3 qwen3_6_scripts/test_cub_compat_v2.py
This commit is contained in:
claude
73
qwen3_6_scripts/probe_corex_cub.sh
Executable file
73
qwen3_6_scripts/probe_corex_cub.sh
Executable file
@@ -0,0 +1,73 @@
|
||||
#!/bin/bash
|
||||
# probe_corex_cub.sh — Discover what CUB/CCCL corex ships
|
||||
# Run on real machine: bash qwen3_6_scripts/probe_corex_cub.sh
|
||||
|
||||
echo "=== corex include tree ==="
|
||||
ls /usr/local/corex/include/ 2>/dev/null
|
||||
echo ""
|
||||
|
||||
echo "=== CUB version ==="
|
||||
grep -r "CUB_VERSION\|CUB_MAJOR\|CUB_MINOR\|CUB_PATCH" /usr/local/corex/include/cub/version.cuh 2>/dev/null || \
|
||||
grep -r "CUB_VERSION" /usr/local/corex/include/cub/*.cuh 2>/dev/null | head -5
|
||||
echo ""
|
||||
|
||||
echo "=== Thrust version ==="
|
||||
grep "THRUST_VERSION" /usr/local/corex/include/thrust/version.h 2>/dev/null | head -3
|
||||
echo ""
|
||||
|
||||
echo "=== CUB block primitives ==="
|
||||
ls /usr/local/corex/include/cub/block/ 2>/dev/null
|
||||
echo ""
|
||||
|
||||
echo "=== CUB warp primitives ==="
|
||||
ls /usr/local/corex/include/cub/warp/ 2>/dev/null
|
||||
echo ""
|
||||
|
||||
echo "=== CUB device algorithms ==="
|
||||
ls /usr/local/corex/include/cub/device/ 2>/dev/null
|
||||
echo ""
|
||||
|
||||
echo "=== __shfl_down_sync in corex headers ==="
|
||||
grep -r "__shfl_down_sync\|__shfl_sync" /usr/local/corex/include/cub/util_ptx.cuh 2>/dev/null | head -5
|
||||
echo ""
|
||||
|
||||
echo "=== corex CUDA version ==="
|
||||
grep "CUDART_VERSION\|CUDA_VERSION" /usr/local/corex/include/cuda.h 2>/dev/null | head -3
|
||||
cat /usr/local/corex/version.txt 2>/dev/null
|
||||
echo ""
|
||||
|
||||
echo "=== Quick compile test: corex's own CUB ==="
|
||||
cat > /tmp/test_corex_cub.cu << 'EOF'
|
||||
#include <cub/block/block_reduce.cuh>
|
||||
#include <cstdio>
|
||||
|
||||
__global__ void test_kernel(float* out) {
|
||||
using BlockReduce = cub::BlockReduce<float, 256>;
|
||||
__shared__ typename BlockReduce::TempStorage temp_storage;
|
||||
float val = (float)threadIdx.x;
|
||||
float sum = BlockReduce(temp_storage).Sum(val);
|
||||
if (threadIdx.x == 0) out[0] = sum;
|
||||
}
|
||||
|
||||
int main() {
|
||||
float *d_out;
|
||||
cudaMalloc(&d_out, sizeof(float));
|
||||
test_kernel<<<1, 256>>>(d_out);
|
||||
float result;
|
||||
cudaMemcpy(&result, d_out, sizeof(float), cudaMemcpyDeviceToHost);
|
||||
printf("BlockReduce sum(0..255) = %.0f (expected 32640)\n", result);
|
||||
cudaFree(d_out);
|
||||
return 0;
|
||||
}
|
||||
EOF
|
||||
|
||||
/usr/local/corex/bin/clang++ --cuda-gpu-arch=ivcore10 --cuda-path=/usr/local/corex \
|
||||
-I/usr/local/corex/include -L/usr/local/corex/lib64 -lcudart \
|
||||
/tmp/test_corex_cub.cu -o /tmp/test_corex_cub 2>&1
|
||||
|
||||
if [ -f /tmp/test_corex_cub ]; then
|
||||
echo "Compile: SUCCESS"
|
||||
/tmp/test_corex_cub
|
||||
else
|
||||
echo "Compile: FAILED"
|
||||
fi
|
||||
259
qwen3_6_scripts/test_cub_compat_v2.py
Normal file
259
qwen3_6_scripts/test_cub_compat_v2.py
Normal file
@@ -0,0 +1,259 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
test_cub_compat_v2.py — Fixed CUB compatibility test for BI-V100
|
||||
|
||||
Changes from v1:
|
||||
- Write .cu to file + torch.utils.cpp_extension.load (not load_inline)
|
||||
- Test 3 uses corex's own CUB (not our CCCL 3.6 upstream)
|
||||
"""
|
||||
import torch
|
||||
import os
|
||||
import tempfile
|
||||
import shutil
|
||||
|
||||
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
def compile_and_run(name, cu_code, fn_name, extra_include=None):
|
||||
"""Write .cu to temp dir, compile with torch cpp_extension, return module."""
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
tmpdir = os.path.join(tempfile.gettempdir(), f"test_{name}")
|
||||
os.makedirs(tmpdir, exist_ok=True)
|
||||
cu_path = os.path.join(tmpdir, f"{name}.cu")
|
||||
with open(cu_path, "w") as f:
|
||||
f.write(cu_code)
|
||||
|
||||
extra_cflags = ["-std=c++17"]
|
||||
extra_include_paths = extra_include or []
|
||||
|
||||
mod = load(
|
||||
name=name,
|
||||
sources=[cu_path],
|
||||
extra_cflags=extra_cflags,
|
||||
extra_include_paths=extra_include_paths,
|
||||
verbose=True,
|
||||
)
|
||||
return mod
|
||||
|
||||
|
||||
def test_shfl():
|
||||
"""Test __shfl_down_sync on ivcore10."""
|
||||
code = r"""
|
||||
#include <torch/extension.h>
|
||||
|
||||
__global__ void shfl_kernel(float* out, int n) {
|
||||
int tid = threadIdx.x;
|
||||
float val = (float)tid;
|
||||
float shuffled = __shfl_down_sync(0xffffffff, val, 1);
|
||||
if (tid < n) out[tid] = shuffled;
|
||||
}
|
||||
|
||||
torch::Tensor test_shfl(int n) {
|
||||
auto out = torch::zeros({n}, torch::dtype(torch::kFloat32).device(torch::kCUDA));
|
||||
shfl_kernel<<<1, n>>>(out.data_ptr<float>(), n);
|
||||
return out;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("test_shfl", &test_shfl);
|
||||
}
|
||||
"""
|
||||
try:
|
||||
mod = compile_and_run("shfl_test", code, "test_shfl")
|
||||
result = mod.test_shfl(32)
|
||||
val = result[0].item()
|
||||
if abs(val - 1.0) < 0.01:
|
||||
print(f" ✓ __shfl_down_sync works: thread0 got {val}")
|
||||
return True
|
||||
else:
|
||||
print(f" ✗ __shfl_down_sync wrong: thread0={val} expected=1.0")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f" ✗ __shfl_down_sync failed: {str(e)[:300]}")
|
||||
return False
|
||||
|
||||
|
||||
def test_block_reduce_manual():
|
||||
"""Test SMEM + shuffle block reduce."""
|
||||
code = r"""
|
||||
#include <torch/extension.h>
|
||||
|
||||
__global__ void block_sum_kernel(const float* input, float* output, int n) {
|
||||
__shared__ float smem[32];
|
||||
int tid = threadIdx.x;
|
||||
int warp_id = tid / 32;
|
||||
int lane_id = tid % 32;
|
||||
|
||||
float val = (tid < n) ? input[tid] : 0.0f;
|
||||
|
||||
for (int offset = 16; offset > 0; offset >>= 1)
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
|
||||
if (lane_id == 0) smem[warp_id] = val;
|
||||
__syncthreads();
|
||||
|
||||
if (warp_id == 0) {
|
||||
val = (tid < blockDim.x / 32) ? smem[tid] : 0.0f;
|
||||
for (int offset = 16; offset > 0; offset >>= 1)
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
if (tid == 0) output[0] = val;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor block_sum(torch::Tensor input) {
|
||||
int n = input.size(0);
|
||||
auto output = torch::zeros({1}, input.options());
|
||||
int threads = ((n + 31) / 32) * 32;
|
||||
if (threads > 1024) threads = 1024;
|
||||
block_sum_kernel<<<1, threads>>>(input.data_ptr<float>(), output.data_ptr<float>(), n);
|
||||
return output;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("block_sum", &block_sum);
|
||||
}
|
||||
"""
|
||||
try:
|
||||
mod = compile_and_run("block_sum_test", code, "block_sum")
|
||||
x = torch.ones(256, dtype=torch.float32, device="cuda")
|
||||
result = mod.block_sum(x)
|
||||
val = result[0].item()
|
||||
if abs(val - 256.0) < 0.01:
|
||||
print(f" ✓ block reduce works: sum(ones[256]) = {val}")
|
||||
return True
|
||||
else:
|
||||
print(f" ✗ block reduce wrong: {val} expected=256.0")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f" ✗ block reduce failed: {str(e)[:300]}")
|
||||
return False
|
||||
|
||||
|
||||
def test_corex_cub():
|
||||
"""Test cub::BlockReduce using corex's OWN CUB headers (not our CCCL 3.6)."""
|
||||
code = r"""
|
||||
#include <torch/extension.h>
|
||||
#include <cub/block/block_reduce.cuh>
|
||||
|
||||
__global__ void cub_reduce_kernel(const float* input, float* output, int n) {
|
||||
using BlockReduce = cub::BlockReduce<float, 256>;
|
||||
__shared__ typename BlockReduce::TempStorage temp_storage;
|
||||
int tid = threadIdx.x;
|
||||
float val = (tid < n) ? input[tid] : 0.0f;
|
||||
float aggregate = BlockReduce(temp_storage).Sum(val);
|
||||
if (tid == 0) output[0] = aggregate;
|
||||
}
|
||||
|
||||
torch::Tensor cub_reduce(torch::Tensor input) {
|
||||
int n = input.size(0);
|
||||
auto output = torch::zeros({1}, input.options());
|
||||
cub_reduce_kernel<<<1, 256>>>(input.data_ptr<float>(), output.data_ptr<float>(), n);
|
||||
return output;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("cub_reduce", &cub_reduce);
|
||||
}
|
||||
"""
|
||||
# Use corex's own CUB, NOT our cccl_upstream
|
||||
corex_include = "/usr/local/corex/include"
|
||||
extra = [corex_include] if os.path.isdir(corex_include) else []
|
||||
|
||||
try:
|
||||
mod = compile_and_run("corex_cub_test", code, "cub_reduce", extra_include=extra)
|
||||
x = torch.arange(256, dtype=torch.float32, device="cuda")
|
||||
result = mod.cub_reduce(x)
|
||||
val = result[0].item()
|
||||
expected = 256 * 255 / 2 # 32640
|
||||
if abs(val - expected) < 1.0:
|
||||
print(f" ✓ cub::BlockReduce (corex) works: sum(0..255) = {val}")
|
||||
return True
|
||||
else:
|
||||
print(f" ✗ cub::BlockReduce wrong: {val} expected={expected}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f" ✗ cub::BlockReduce (corex) failed: {str(e)[:300]}")
|
||||
return False
|
||||
|
||||
|
||||
def test_cccl36_define_bypass():
|
||||
"""Test CCCL 3.6 with CCCL_IGNORE_DEPRECATED_CUDA_BELOW_12 define."""
|
||||
cccl_cub = os.path.join(os.path.dirname(SCRIPT_DIR), "cccl_upstream", "cub")
|
||||
cccl_libcudacxx = os.path.join(os.path.dirname(SCRIPT_DIR), "cccl_upstream", "libcudacxx", "include")
|
||||
|
||||
if not os.path.isdir(cccl_cub):
|
||||
print(" ⊘ SKIP: cccl_upstream not found")
|
||||
return None
|
||||
|
||||
code = r"""
|
||||
#define CCCL_IGNORE_DEPRECATED_CUDA_BELOW_12
|
||||
#include <torch/extension.h>
|
||||
#include <cub/block/block_reduce.cuh>
|
||||
|
||||
__global__ void cub36_kernel(const float* input, float* output, int n) {
|
||||
using BlockReduce = cub::BlockReduce<float, 256>;
|
||||
__shared__ typename BlockReduce::TempStorage temp_storage;
|
||||
int tid = threadIdx.x;
|
||||
float val = (tid < n) ? input[tid] : 0.0f;
|
||||
float aggregate = BlockReduce(temp_storage).Sum(val);
|
||||
if (tid == 0) output[0] = aggregate;
|
||||
}
|
||||
|
||||
torch::Tensor cub36_reduce(torch::Tensor input) {
|
||||
int n = input.size(0);
|
||||
auto output = torch::zeros({1}, input.options());
|
||||
cub36_kernel<<<1, 256>>>(input.data_ptr<float>(), output.data_ptr<float>(), n);
|
||||
return output;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("cub36_reduce", &cub36_reduce);
|
||||
}
|
||||
"""
|
||||
try:
|
||||
mod = compile_and_run("cccl36_cub_test", code, "cub36_reduce",
|
||||
extra_include=[cccl_cub, cccl_libcudacxx])
|
||||
x = torch.arange(256, dtype=torch.float32, device="cuda")
|
||||
result = mod.cub36_reduce(x)
|
||||
val = result[0].item()
|
||||
expected = 32640
|
||||
if abs(val - expected) < 1.0:
|
||||
print(f" ✓ CCCL 3.6 BlockReduce (with bypass) works: {val}")
|
||||
return True
|
||||
else:
|
||||
print(f" ✗ CCCL 3.6 BlockReduce wrong: {val}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f" ✗ CCCL 3.6 (with bypass) failed: {str(e)[:300]}")
|
||||
return False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if not torch.cuda.is_available():
|
||||
print("No CUDA"); exit(1)
|
||||
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
||||
print()
|
||||
|
||||
print("[1/4] __shfl_down_sync (warp shuffle)")
|
||||
r1 = test_shfl()
|
||||
|
||||
print("[2/4] manual block reduce (SMEM + shuffle)")
|
||||
r2 = test_block_reduce_manual()
|
||||
|
||||
print("[3/4] cub::BlockReduce (corex's own CUB)")
|
||||
r3 = test_corex_cub()
|
||||
|
||||
print("[4/4] cub::BlockReduce (CCCL 3.6 with CUDA<12 bypass)")
|
||||
r4 = test_cccl36_define_bypass()
|
||||
|
||||
print()
|
||||
results = [r for r in [r1, r2, r3, r4] if r is not None]
|
||||
passed = sum(results)
|
||||
total = len(results)
|
||||
print(f"{'='*50}")
|
||||
print(f" {passed}/{total} passed")
|
||||
if r1: print(" → warp shuffle works on ivcore10")
|
||||
if r3: print(" → corex CUB is usable for block primitives")
|
||||
if r4: print(" → CCCL 3.6 CAN compile with bypass define")
|
||||
elif r4 is False: print(" → CCCL 3.6 variadic function issue remains — use corex CUB instead")
|
||||
print(f"{'='*50}")
|
||||
Reference in New Issue
Block a user