dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[CCCL] 瘦身 + 补全: 移除 cudax/python/libcudacxx-tests 冗余文件, 新增 c2h 测试助手 + cmake 构建系统 + 8 个 CUDA thrust examples

Browse Source 
变更摘要:
- 删除: cudax/ (783 files, 7.2M) — 实验性组件，竞赛不需要
- 删除: python/ (226 files, 2.0M) — Python 绑定，竞赛不需要
- 删除: libcudacxx/{test,benchmarks,codegen,cmake,share} (4432 files, 31M)
  保留: libcudacxx/include/ (1463 headers, cuda::std 编译依赖)
- 新增: c2h/ (27 files) — CUB Catch2 测试辅助头文件，编译 243 个测试必需
- 新增: cmake/ (29 files) — CCCL 原生 CMake 构建系统
- 新增: thrust/examples/cuda/ (7 files) + cpp_integration/ (1 file)
  async_reduce, custom_temporary_allocation, explicit_cuda_stream,
  global_device_vector, range_view, unwrap_pointer, wrap_pointer, device

结果: cccl_upstream 从 74M→35M (瘦身 53%), 核心内容 100% 保留:
  27/27 tuning headers, 78 benchmarks, 243 tests,
  60 thrust examples, 18 CUB examples, 全部编译头文件
This commit is contained in:
muh-bot
2026-08-03 12:39:26 +00:00
parent a2a5dd8f00
commit 24ef6a91b5
5439 changed files with 0 additions and 719516 deletions
Download Patch File Download Diff File Expand all files Collapse all files

136
cccl_upstream/cudax/test/places/CMakeLists.txt
View File

@@ -1,136 +0,0 @@
set(
places_test_sources
include_only.cu
data_place_alloc.cu
data_place_vmm.cu
exec_place_scope.cu
placement.cu
stream_pool.cu
)
set(places_fail_tests exec_place_scope_data_place_fail.cu)
set(
places_unittested_headers
cuda/experimental/__places/cute_partition.cuh
cuda/experimental/__places/places.cuh
cuda/experimental/__places/exec/cuda_context.cuh
cuda/experimental/__places/exec/green_context.cuh
cuda/experimental/__places/partitions/blocked_partition.cuh
cuda/experimental/__places/partitions/cyclic_shape.cuh
cuda/experimental/__places/partitions/tiled_partition.cuh
)
cccl_get_cudatoolkit()
## cudax_add_places_test
#
# Add a places test executable and register it with ctest.
#
# target_name_var: Variable name to overwrite with the name of the test
# target. Useful for adding target information after creation.
# source: The source file for the test.
#
function(cudax_add_places_test target_name_var source)
get_filename_component(dir ${source} DIRECTORY)
get_filename_component(filename ${source} NAME_WE)
if (dir)
set(filename "${dir}/${filename}")
endif()
string(REPLACE "/" "." test_name "${filename}")
set(test_target cudax.test.places.${test_name})
cccl_add_executable(${test_target} SOURCES ${source} ADD_CTEST)
cudax_places_configure_target(${test_target})
target_link_libraries(${test_target} PRIVATE cudax.compiler_interface)
set(${target_name_var} ${test_target} PARENT_SCOPE)
endfunction()
## cudax_add_places_unittest_header
#
# Add a places unittested header executable and register it with ctest.
#
# Unittested headers contain a set of tests that are enabled by including
# `unittest.cuh` and defining `UNITTESTED_FILE`.
#
# target_name_var: Variable name to overwrite with the name of the test
# target. Useful for adding target information after creation.
# source: The source file for the test.
#
function(cudax_add_places_unittest_header target_name_var source)
get_filename_component(relative_path ${source} DIRECTORY)
get_filename_component(filename ${source} NAME_WE)
string(
REPLACE
"cuda/experimental/"
""
test_label
"${relative_path}/${filename}"
)
string(REPLACE "/" "." test_label "${test_label}")
set(test_target "cudax.test.places.unittest_headers.${test_label}")
get_filename_component(
source_full_path
../../../cudax/include/${source}
ABSOLUTE
)
set(source ${source_full_path})
set(ut_template "${cudax_SOURCE_DIR}/cmake/places_header_unittest.in.cu")
set(ut_source "${cudax_BINARY_DIR}/unittest_headers/${test_target}.cu")
configure_file(${ut_template} ${ut_source} @ONLY)
cccl_add_executable(${test_target} SOURCES ${ut_source} ADD_CTEST)
cudax_places_configure_target(${test_target})
target_link_libraries(${test_target} PRIVATE cudax.compiler_interface)
set(${target_name_var} ${test_target} PARENT_SCOPE)
endfunction()
# Basic tests:
foreach (source IN LISTS places_test_sources)
cudax_add_places_test(test_target "${source}")
endforeach()
# Unittested headers
foreach (source IN LISTS places_unittested_headers)
cudax_add_places_unittest_header(test_target "${source}")
endforeach()
## cudax_add_places_fail_test
#
# Adds an EXCLUDE_FROM_ALL build target for `source` and a ctest that
# verifies compilation fails with errors matching expected-error annotations
# in the source file.
#
function(cudax_add_places_fail_test target_name_var source)
get_filename_component(filename ${source} NAME_WE)
set(test_target cudax.test.places.error.${filename})
cccl_add_executable(
${test_target}
SOURCES ${source}
NO_METATARGETS
NO_CLANG_TIDY
)
cudax_places_configure_target(${test_target})
target_link_libraries(${test_target} PRIVATE cudax.compiler_interface)
cccl_add_xfail_compile_target_test(
${test_target}
SOURCE_FILE "${source}"
ERROR_REGEX_LABEL "expected-error"
)
set(${target_name_var} ${test_target} PARENT_SCOPE)
endfunction()
# Expected-failure compile tests (_fail.cu convention)
foreach (source IN LISTS places_fail_tests)
cudax_add_places_fail_test(test_target "${source}")
endforeach()

195
cccl_upstream/cudax/test/places/data_place_alloc.cu
View File

@@ -1,195 +0,0 @@
//===----------------------------------------------------------------------===//
//
// Part of CUDASTF in CUDA C++ Core Libraries,
// under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
*
* @brief Test that data_place can be used to allocate/deallocate memory
* directly without a CUDASTF context.
*
* This demonstrates how places can be used for raw memory allocation
* outside of the task-based programming model.
*/
#include <cuda/experimental/__places/places.cuh>
#include <cstdio>
using namespace cuda::experimental::places;
__global__ void init_kernel(int* ptr, int n, int value)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid < n)
{
ptr[tid] = value + tid;
}
}
__global__ void check_kernel(int* ptr, int n, int value, int* result)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid < n)
{
if (ptr[tid] != value + tid)
{
atomicExch(result, 1); // Set error flag
}
}
}
void test_host_allocation()
{
printf("Testing host allocation...\n");
const size_t n = 1024;
const size_t byte_size = n * sizeof(int);
// Allocate using data_place::host() - stream parameter is ignored for host allocations
auto place = data_place::host();
EXPECT(!place.allocation_is_stream_ordered()); // Host allocations are blocking
int* ptr = static_cast<int*>(place.allocate(byte_size));
EXPECT(ptr != nullptr);
// Initialize on host
for (size_t i = 0; i < n; i++)
{
ptr[i] = static_cast<int>(i * 2);
}
// Verify
for (size_t i = 0; i < n; i++)
{
EXPECT(ptr[i] == static_cast<int>(i * 2));
}
// Deallocate
place.deallocate(ptr, byte_size, nullptr);
printf(" Host allocation test PASSED\n");
}
void test_device_allocation()
{
printf("Testing device allocation...\n");
const size_t n = 1024;
const size_t byte_size = n * sizeof(int);
const int test_value = 42;
// Create a stream for the allocation
cudaStream_t stream;
cuda_try(cudaStreamCreate(&stream));
// Allocate using data_place::device(0)
auto place = data_place::device(0);
EXPECT(place.allocation_is_stream_ordered()); // Device allocations are stream-ordered
int* d_ptr = static_cast<int*>(place.allocate(byte_size, stream));
EXPECT(d_ptr != nullptr);
// Initialize on device
init_kernel<<<(n + 255) / 256, 256, 0, stream>>>(d_ptr, n, test_value);
// Allocate result flag on host for checking
int* d_result;
cuda_try(cudaMallocAsync(&d_result, sizeof(int), stream));
cuda_try(cudaMemsetAsync(d_result, 0, sizeof(int), stream));
// Check on device
check_kernel<<<(n + 255) / 256, 256, 0, stream>>>(d_ptr, n, test_value, d_result);
// Copy result back
int h_result = 0;
cuda_try(cudaMemcpyAsync(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost, stream));
cuda_try(cudaStreamSynchronize(stream));
EXPECT(h_result == 0); // No errors
// Cleanup
cuda_try(cudaFreeAsync(d_result, stream));
place.deallocate(d_ptr, byte_size, stream);
cuda_try(cudaStreamSynchronize(stream));
cuda_try(cudaStreamDestroy(stream));
printf(" Device allocation test PASSED\n");
}
void test_managed_allocation()
{
printf("Testing managed allocation...\n");
// Check if concurrent managed access is supported
int dev;
cuda_try(cudaGetDevice(&dev));
cudaDeviceProp prop;
cuda_try(cudaGetDeviceProperties(&prop, dev));
if (!prop.concurrentManagedAccess)
{
printf(" Concurrent CPU/GPU access not supported, skipping managed test.\n");
return;
}
const size_t n = 1024;
const size_t byte_size = n * sizeof(int);
const int test_value = 100;
cudaStream_t stream;
cuda_try(cudaStreamCreate(&stream));
// Allocate using data_place::managed()
auto place = data_place::managed();
EXPECT(!place.allocation_is_stream_ordered()); // Managed allocations are immediate (stream ignored)
int* ptr = static_cast<int*>(place.allocate(byte_size));
EXPECT(ptr != nullptr);
// Initialize on host (managed memory is accessible from both CPU and GPU)
for (size_t i = 0; i < n; i++)
{
ptr[i] = test_value + static_cast<int>(i);
}
// Read back on device and verify
int* d_result;
cuda_try(cudaMallocAsync(&d_result, sizeof(int), stream));
cuda_try(cudaMemsetAsync(d_result, 0, sizeof(int), stream));
check_kernel<<<(n + 255) / 256, 256, 0, stream>>>(ptr, n, test_value, d_result);
int h_result = 0;
cuda_try(cudaMemcpyAsync(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost, stream));
cuda_try(cudaStreamSynchronize(stream));
EXPECT(h_result == 0); // No errors
// Cleanup
cuda_try(cudaFreeAsync(d_result, stream));
cuda_try(cudaStreamSynchronize(stream));
place.deallocate(ptr, byte_size);
cuda_try(cudaStreamDestroy(stream));
printf(" Managed allocation test PASSED\n");
}
int main()
{
printf("=== Testing data_place direct allocation (no context) ===\n\n");
test_host_allocation();
test_device_allocation();
test_managed_allocation();
printf("\n=== All tests PASSED ===\n");
return 0;
}

395
cccl_upstream/cudax/test/places/data_place_vmm.cu
View File

@@ -1,395 +0,0 @@
//===----------------------------------------------------------------------===//
//
// Part of CUDASTF in CUDA C++ Core Libraries,
// under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
*
* @brief Test that data_place::mem_create() can be used to create VMM-based
* physical memory allocations.
*
* This tests the low-level VMM allocation interface used by localized arrays
* (composite_slice) for creating physical memory segments that are mapped
* into a contiguous virtual address space.
*/
#include <cuda/experimental/__places/places.cuh>
#if _CCCL_CTK_AT_LEAST(12, 4)
# include <cuda/experimental/__places/exec/green_context.cuh>
#endif // _CCCL_CTK_AT_LEAST(12, 4)
#include <cstdio>
using namespace cuda::experimental::places;
__global__ void init_kernel(int* ptr, int n, int value)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid < n)
{
ptr[tid] = value + tid;
}
}
__global__ void check_kernel(int* ptr, int n, int value, int* result)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid < n)
{
if (ptr[tid] != value + tid)
{
atomicExch(result, 1); // Set error flag
}
}
}
// Check if VMM is supported on the current device
bool vmm_supported(int dev_id = 0)
{
CUdevice dev;
cuda_try(cuDeviceGet(&dev, dev_id));
int supportsVMM;
cuda_try(cuDeviceGetAttribute(&supportsVMM, CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED, dev));
return supportsVMM == 1;
}
// Get allocation granularity for VMM
size_t get_granularity(int dev_id)
{
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.location.id = dev_id;
size_t granularity;
cuda_try(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
return granularity;
}
void test_device_vmm_allocation()
{
printf("Testing device VMM allocation (mem_create)...\n");
int dev_id = 0;
cuda_try(cudaSetDevice(dev_id));
// Get allocation granularity - VMM allocations must be aligned to this
size_t granularity = get_granularity(dev_id);
printf(" Allocation granularity: %zu bytes\n", granularity);
// Allocate at least one granularity unit
const size_t alloc_size = granularity;
const size_t n = alloc_size / sizeof(int);
const int test_value = 42;
// Create physical memory using data_place::mem_create
auto place = data_place::device(dev_id);
CUmemGenericAllocationHandle handle;
CUresult result = place.mem_create(&handle, alloc_size);
EXPECT(result == CUDA_SUCCESS);
// Reserve virtual address space
CUdeviceptr va_ptr;
cuda_try(cuMemAddressReserve(&va_ptr, alloc_size, 0, 0, 0));
// Map the physical allocation to the virtual address
cuda_try(cuMemMap(va_ptr, alloc_size, 0, handle, 0));
// Set access permissions for the current device
CUmemAccessDesc accessDesc = {};
accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
accessDesc.location.id = dev_id;
accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
cuda_try(cuMemSetAccess(va_ptr, alloc_size, &accessDesc, 1));
// Now we can use the memory!
int* d_ptr = reinterpret_cast<int*>(va_ptr); // NOLINT(performance-no-int-to-ptr)
// Create a stream for operations
cudaStream_t stream;
cuda_try(cudaStreamCreate(&stream));
// Initialize on device
init_kernel<<<(n + 255) / 256, 256, 0, stream>>>(d_ptr, n, test_value);
cuda_try(cudaGetLastError());
// Allocate result flag for checking
int* d_result;
cuda_try(cudaMallocAsync(&d_result, sizeof(int), stream));
cuda_try(cudaMemsetAsync(d_result, 0, sizeof(int), stream));
// Check on device
check_kernel<<<(n + 255) / 256, 256, 0, stream>>>(d_ptr, n, test_value, d_result);
cuda_try(cudaGetLastError());
// Copy result back
int h_result = 0;
cuda_try(cudaMemcpyAsync(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost, stream));
cuda_try(cudaStreamSynchronize(stream));
EXPECT(h_result == 0); // No errors
// Cleanup
cuda_try(cudaFreeAsync(d_result, stream));
cuda_try(cudaStreamSynchronize(stream));
cuda_try(cudaStreamDestroy(stream));
// Unmap and release VMM resources
cuda_try(cuMemUnmap(va_ptr, alloc_size));
cuda_try(cuMemRelease(handle));
cuda_try(cuMemAddressFree(va_ptr, alloc_size));
printf(" Device VMM allocation test PASSED\n");
}
// Host VMM requires CU_MEM_LOCATION_TYPE_HOST which is only available in CUDA 12.2+
#if _CCCL_CTK_AT_LEAST(12, 2)
void test_host_vmm_allocation()
{
printf("Testing host VMM allocation (mem_create)...\n");
// Host VMM allocations use CU_MEM_LOCATION_TYPE_HOST
// First check if host VMM is supported (requires appropriate driver/hardware)
// Get host allocation granularity
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_HOST;
prop.location.id = 0;
size_t granularity;
CUresult gran_result = cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM);
if (gran_result != CUDA_SUCCESS)
{
printf(" Host VMM not supported on this system, skipping.\n");
return;
}
printf(" Host allocation granularity: %zu bytes\n", granularity);
const size_t alloc_size = granularity;
const size_t n = alloc_size / sizeof(int);
// Create physical memory using data_place::mem_create with host place
auto place = data_place::host();
CUmemGenericAllocationHandle handle;
CUresult result = place.mem_create(&handle, alloc_size);
if (result != CUDA_SUCCESS)
{
printf(" Host mem_create not supported (error %d), skipping.\n", result);
return;
}
// Reserve virtual address space
CUdeviceptr va_ptr;
cuda_try(cuMemAddressReserve(&va_ptr, alloc_size, 0, 0, 0));
// Map the physical allocation to the virtual address
cuda_try(cuMemMap(va_ptr, alloc_size, 0, handle, 0));
// Set access permissions for the host
CUmemAccessDesc accessDesc = {};
accessDesc.location.type = CU_MEM_LOCATION_TYPE_HOST;
accessDesc.location.id = 0;
accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
cuda_try(cuMemSetAccess(va_ptr, alloc_size, &accessDesc, 1));
// Use the memory from the host
int* ptr = reinterpret_cast<int*>(va_ptr); // NOLINT(performance-no-int-to-ptr)
// Initialize on host
for (size_t i = 0; i < n; i++)
{
ptr[i] = static_cast<int>(i * 2);
}
// Verify on host
for (size_t i = 0; i < n; i++)
{
EXPECT(ptr[i] == static_cast<int>(i * 2));
}
// Cleanup VMM resources
cuda_try(cuMemUnmap(va_ptr, alloc_size));
cuda_try(cuMemRelease(handle));
cuda_try(cuMemAddressFree(va_ptr, alloc_size));
printf(" Host VMM allocation test PASSED\n");
}
#endif // _CCCL_CTK_AT_LEAST(12, 2)
void test_multi_segment_vmm()
{
printf("Testing multi-segment VMM allocation...\n");
int dev_id = 0;
cuda_try(cudaSetDevice(dev_id));
size_t granularity = get_granularity(dev_id);
// Create two segments and map them contiguously
const size_t segment_size = granularity;
const size_t total_size = 2 * segment_size;
const size_t n = total_size / sizeof(int);
const int test_value = 100;
auto place = data_place::device(dev_id);
// Create two physical allocations
CUmemGenericAllocationHandle handle1, handle2;
cuda_try(place.mem_create(&handle1, segment_size));
cuda_try(place.mem_create(&handle2, segment_size));
// Reserve contiguous virtual address space for both
CUdeviceptr va_ptr;
cuda_try(cuMemAddressReserve(&va_ptr, total_size, 0, 0, 0));
// Map both segments contiguously
cuda_try(cuMemMap(va_ptr, segment_size, 0, handle1, 0));
cuda_try(cuMemMap(va_ptr + segment_size, segment_size, 0, handle2, 0));
// Set access for the entire range
CUmemAccessDesc accessDesc = {};
accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
accessDesc.location.id = dev_id;
accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
cuda_try(cuMemSetAccess(va_ptr, total_size, &accessDesc, 1));
int* d_ptr = reinterpret_cast<int*>(va_ptr); // NOLINT(performance-no-int-to-ptr)
cudaStream_t stream;
cuda_try(cudaStreamCreate(&stream));
// Initialize the entire contiguous range
init_kernel<<<(n + 255) / 256, 256, 0, stream>>>(d_ptr, n, test_value);
cuda_try(cudaGetLastError());
// Check the entire range
int* d_result;
cuda_try(cudaMallocAsync(&d_result, sizeof(int), stream));
cuda_try(cudaMemsetAsync(d_result, 0, sizeof(int), stream));
check_kernel<<<(n + 255) / 256, 256, 0, stream>>>(d_ptr, n, test_value, d_result);
cuda_try(cudaGetLastError());
int h_result = 0;
cuda_try(cudaMemcpyAsync(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost, stream));
cuda_try(cudaStreamSynchronize(stream));
EXPECT(h_result == 0);
// Cleanup
cuda_try(cudaFreeAsync(d_result, stream));
cuda_try(cudaStreamSynchronize(stream));
cuda_try(cudaStreamDestroy(stream));
cuda_try(cuMemUnmap(va_ptr, segment_size));
cuda_try(cuMemUnmap(va_ptr + segment_size, segment_size));
cuda_try(cuMemRelease(handle1));
cuda_try(cuMemRelease(handle2));
cuda_try(cuMemAddressFree(va_ptr, total_size));
printf(" Multi-segment VMM allocation test PASSED\n");
}
#if _CCCL_CTK_AT_LEAST(12, 4)
// Green context data places override mem_create() to allocate device-pinned
// physical memory bound to the green context's device. This exercises that
// override end-to-end through the VMM map/access/kernel path.
void test_green_ctx_vmm_allocation()
{
int dev_id = 0;
cuda_try(cudaSetDevice(dev_id));
// Split the device into green contexts (8 SMs each).
green_context_helper gc_helper(8, dev_id);
auto gc_view = gc_helper.get_view(0);
auto place = data_place::green_ctx(gc_view);
size_t granularity = get_granularity(dev_id);
const size_t alloc_size = granularity;
const size_t n = alloc_size / sizeof(int);
const int test_value = 55;
// Create physical memory using the green context data place's mem_create().
CUmemGenericAllocationHandle handle;
CUresult result = place.mem_create(&handle, alloc_size);
EXPECT(result == CUDA_SUCCESS);
// Reserve virtual address space and map the physical allocation into it.
CUdeviceptr va_ptr;
cuda_try(cuMemAddressReserve(&va_ptr, alloc_size, 0, 0, 0));
cuda_try(cuMemMap(va_ptr, alloc_size, 0, handle, 0));
// Green context memory is pinned on the underlying device, so grant device
// access accordingly.
CUmemAccessDesc accessDesc = {};
accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
accessDesc.location.id = dev_id;
accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
cuda_try(cuMemSetAccess(va_ptr, alloc_size, &accessDesc, 1));
int* d_ptr = reinterpret_cast<int*>(va_ptr); // NOLINT(performance-no-int-to-ptr)
cudaStream_t stream;
cuda_try(cudaStreamCreate(&stream));
init_kernel<<<(n + 255) / 256, 256, 0, stream>>>(d_ptr, n, test_value);
cuda_try(cudaGetLastError());
int* d_result;
cuda_try(cudaMallocAsync(&d_result, sizeof(int), stream));
cuda_try(cudaMemsetAsync(d_result, 0, sizeof(int), stream));
check_kernel<<<(n + 255) / 256, 256, 0, stream>>>(d_ptr, n, test_value, d_result);
cuda_try(cudaGetLastError());
int h_result = 0;
cuda_try(cudaMemcpyAsync(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost, stream));
cuda_try(cudaStreamSynchronize(stream));
EXPECT(h_result == 0);
// Cleanup
cuda_try(cudaFreeAsync(d_result, stream));
cuda_try(cudaStreamSynchronize(stream));
cuda_try(cudaStreamDestroy(stream));
cuda_try(cuMemUnmap(va_ptr, alloc_size));
cuda_try(cuMemRelease(handle));
cuda_try(cuMemAddressFree(va_ptr, alloc_size));
}
#endif // _CCCL_CTK_AT_LEAST(12, 4)
int main()
{
printf("=== Testing data_place VMM allocation (mem_create) ===\n\n");
// Initialize CUDA driver API
cuda_try(cuInit(0));
// Check VMM support
if (!vmm_supported())
{
printf("VMM not supported on this device, skipping tests.\n");
return 0;
}
test_device_vmm_allocation();
#if _CCCL_CTK_AT_LEAST(12, 2)
test_host_vmm_allocation();
#endif // _CCCL_CTK_AT_LEAST(12, 2)
test_multi_segment_vmm();
#if _CCCL_CTK_AT_LEAST(12, 4)
test_green_ctx_vmm_allocation();
#endif // _CCCL_CTK_AT_LEAST(12, 4)
printf("\n=== All VMM tests PASSED ===\n");
return 0;
}

283
cccl_upstream/cudax/test/places/exec_place_scope.cu
View File

@@ -1,283 +0,0 @@
//===----------------------------------------------------------------------===//
//
// Part of CUDASTF in CUDA C++ Core Libraries,
// under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief Unit tests for exec_place_scope RAII helper
*/
#include <cuda/experimental/__places/places.cuh>
#include <thread>
#include <vector>
using namespace cuda::experimental::places;
// Test 1: Basic scope functionality - single device switch
void test_basic_scope(int ndevs)
{
if (ndevs < 2)
{
fprintf(stderr, "test_basic_scope: skipping (need at least 2 devices).\n");
return;
}
// Start on device 0
cuda_try(cudaSetDevice(0));
int dev_before = -1;
cuda_try(cudaGetDevice(&dev_before));
EXPECT(dev_before == 0);
// Use scope to switch to device 1
{
exec_place_scope scope(exec_place::device(1));
int dev_inside = -1;
cuda_try(cudaGetDevice(&dev_inside));
EXPECT(dev_inside == 1);
}
// After scope destruction, should be back to device 0
int dev_after = -1;
cuda_try(cudaGetDevice(&dev_after));
EXPECT(dev_after == 0);
}
// Test 2: Nested scopes
void test_nested_scopes(int ndevs)
{
if (ndevs < 3)
{
fprintf(stderr, "test_nested_scopes: skipping (need at least 3 devices).\n");
return;
}
cuda_try(cudaSetDevice(0));
{
exec_place_scope scope1(exec_place::device(1));
int dev = -1;
cuda_try(cudaGetDevice(&dev));
EXPECT(dev == 1);
{
exec_place_scope scope2(exec_place::device(2));
cuda_try(cudaGetDevice(&dev));
EXPECT(dev == 2);
}
// After inner scope destruction, should be back to device 1
cuda_try(cudaGetDevice(&dev));
EXPECT(dev == 1);
}
// After outer scope destruction, should be back to device 0
int dev = -1;
cuda_try(cudaGetDevice(&dev));
EXPECT(dev == 0);
}
// Test 3: Scope with host execution place (should be no-op for device)
void test_host_place_scope(int ndevs)
{
if (ndevs < 1)
{
fprintf(stderr, "test_host_place_scope: skipping (need at least 1 device).\n");
return;
}
cuda_try(cudaSetDevice(0));
int dev_before = -1;
cuda_try(cudaGetDevice(&dev_before));
{
exec_place_scope scope(exec_place::host());
// Device should remain unchanged when using host place
int dev_inside = -1;
cuda_try(cudaGetDevice(&dev_inside));
EXPECT(dev_inside == dev_before);
}
int dev_after = -1;
cuda_try(cudaGetDevice(&dev_after));
EXPECT(dev_after == dev_before);
}
// Test 4: Scope with same device (should be efficient no-op)
void test_same_device_scope(int ndevs)
{
if (ndevs < 1)
{
fprintf(stderr, "test_same_device_scope: skipping (need at least 1 device).\n");
return;
}
cuda_try(cudaSetDevice(0));
{
exec_place_scope scope(exec_place::device(0));
int dev = -1;
cuda_try(cudaGetDevice(&dev));
EXPECT(dev == 0);
}
int dev = -1;
cuda_try(cudaGetDevice(&dev));
EXPECT(dev == 0);
}
// Test 5: Stream creation within scope
void test_stream_creation_in_scope(int ndevs)
{
if (ndevs < 2)
{
fprintf(stderr, "test_stream_creation_in_scope: skipping (need at least 2 devices).\n");
return;
}
cuda_try(cudaSetDevice(0));
cudaStream_t stream;
{
exec_place_scope scope(exec_place::device(1));
cuda_try(cudaStreamCreate(&stream));
}
#if _CCCL_CTK_AT_LEAST(12, 8)
// Verify stream was created on device 1 (cudaStreamGetDevice requires CUDA 12.8+)
int stream_dev = -1;
cuda_try(cudaStreamGetDevice(stream, &stream_dev));
EXPECT(stream_dev == 1);
#endif // _CCCL_CTK_AT_LEAST(12, 8)
// Clean up (need to be on correct device for some operations)
{
exec_place_scope scope(exec_place::device(1));
cuda_try(cudaStreamDestroy(stream));
}
}
// Test 6: Multiple threads with scopes
void test_multithreaded_scopes(int ndevs)
{
if (ndevs < 2)
{
fprintf(stderr, "test_multithreaded_scopes: skipping (need at least 2 devices).\n");
return;
}
const int num_threads = ::std::min(ndevs, 4);
::std::vector<bool> results(num_threads, false);
::std::vector<::std::thread> threads;
threads.reserve(num_threads);
for (int i = 0; i < num_threads; ++i)
{
threads.emplace_back([&results, i, ndevs]() {
// Each thread starts on device 0
cuda_try(cudaSetDevice(0));
int target_dev = i % ndevs;
{
exec_place_scope scope(exec_place::device(target_dev));
int dev = -1;
cuda_try(cudaGetDevice(&dev));
if (dev != target_dev)
{
return;
}
// Do some work
cudaStream_t stream;
cuda_try(cudaStreamCreate(&stream));
cuda_try(cudaStreamSynchronize(stream));
cuda_try(cudaStreamDestroy(stream));
}
// Verify restoration
int dev_after = -1;
cuda_try(cudaGetDevice(&dev_after));
if (dev_after != 0)
{
return;
}
results[i] = true;
});
}
for (auto& th : threads)
{
th.join();
}
for (int i = 0; i < num_threads; ++i)
{
EXPECT(results[i]);
}
}
// Test 7: Stress test with multiple iterations
void test_stress_iterations(int ndevs)
{
if (ndevs < 2)
{
fprintf(stderr, "test_stress_iterations: skipping (need at least 2 devices).\n");
return;
}
cuda_try(cudaSetDevice(0));
const int iterations = 100;
for (int iter = 0; iter < iterations; ++iter)
{
int target_dev = iter % ndevs;
{
exec_place_scope scope(exec_place::device(target_dev));
int dev = -1;
cuda_try(cudaGetDevice(&dev));
EXPECT(dev == target_dev);
}
int dev_after = -1;
cuda_try(cudaGetDevice(&dev_after));
EXPECT(dev_after == 0);
}
}
int main()
{
// Initialize CUDA
cuda_try(cudaFree(nullptr));
int ndevs;
cuda_try(cudaGetDeviceCount(&ndevs));
test_basic_scope(ndevs);
test_nested_scopes(ndevs);
test_host_place_scope(ndevs);
test_same_device_scope(ndevs);
test_stream_creation_in_scope(ndevs);
test_multithreaded_scopes(ndevs);
test_stress_iterations(ndevs);
return 0;
}

26
cccl_upstream/cudax/test/places/exec_place_scope_data_place_fail.cu
View File

@@ -1,26 +0,0 @@
//===----------------------------------------------------------------------===//
//
// Part of CUDASTF in CUDA C++ Core Libraries,
// under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief Static error check: exec_place_scope should not accept data_place
*/
#include <cuda/experimental/__places/places.cuh>
using namespace cuda::experimental::places;
int main()
{
// expected-error {{"exec_place_scope cannot be constructed from data_place; use data_place::affine_exec_place() to
// get the exec_place first"}}
exec_place_scope scope(data_place::device(0));
return EXIT_FAILURE;
}

31
cccl_upstream/cudax/test/places/include_only.cu
View File

@@ -1,31 +0,0 @@
//===----------------------------------------------------------------------===//
//
// Part of CUDASTF in CUDA C++ Core Libraries,
// under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
// Single umbrella include: must pull in every public places surface (including
// partition strategies) without requiring clients to list internal headers.
#include <cuda/experimental/places.cuh>
using namespace cuda::experimental::places;
int main()
{
auto host_place = data_place::host();
auto dev0_place = data_place::device(0);
auto exec_host = exec_place::host();
auto exec_dev0 = exec_place::device(0);
(void) host_place;
(void) dev0_place;
(void) exec_host;
(void) exec_dev0;
return 0;
}

379
cccl_upstream/cudax/test/places/placement.cu
View File

@@ -1,379 +0,0 @@
//===----------------------------------------------------------------------===//
//
// Part of CUDASTF in CUDA C++ Core Libraries,
// under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
*
* @brief Test placement evaluation and geometry-aware (shaped) composite
* allocation: evaluate_localized_placement(), the
* allocate_nd(data_dims, elemsize) data_place interface, and
* cute_partition-backed composite places.
*
* Runs on a single GPU (all places on device 0); with two or more GPUs it
* additionally asserts physical residency of the allocated blocks and the
* peer-access path.
*/
#include <cuda/experimental/__places/cute_partition.cuh>
#include <cuda/experimental/__places/partitions/blocked_partition.cuh>
#include <cuda/experimental/__places/places.cuh>
#include <cstdio>
#include <limits>
using namespace cuda::experimental::places;
namespace
{
__global__ void init_kernel(int* ptr, size_t n, int value)
{
size_t tid = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (tid < n)
{
ptr[tid] = value + static_cast<int>(tid % 1024);
}
}
__global__ void check_kernel(const int* ptr, size_t n, int value, int* result)
{
size_t tid = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (tid < n)
{
if (ptr[tid] != value + static_cast<int>(tid % 1024))
{
atomicExch(result, 1); // Set error flag
}
}
}
// Check if VMM is supported on the current device
bool vmm_supported(int dev_id = 0)
{
CUdevice dev;
cuda_try(cuDeviceGet(&dev, dev_id));
int supportsVMM;
cuda_try(cuDeviceGetAttribute(&supportsVMM, CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED, dev));
return supportsVMM == 1;
}
exec_place make_device_grid(int ndevs, size_t nplaces)
{
::std::vector<exec_place> places;
for (size_t i = 0; i < nplaces; i++)
{
places.push_back(exec_place::device(static_cast<int>(i % static_cast<size_t>(ndevs))));
}
return make_grid(mv(places));
}
// Write to the buffer through a kernel and verify the content
void write_and_check(int* ptr, size_t n, int value)
{
const int nthreads = 256;
const int nblocks = static_cast<int>((n + nthreads - 1) / nthreads);
init_kernel<<<nblocks, nthreads>>>(ptr, n, value);
cuda_try(cudaGetLastError());
int* d_result;
cuda_try(cudaMalloc(&d_result, sizeof(int)));
cuda_try(cudaMemset(d_result, 0, sizeof(int)));
check_kernel<<<nblocks, nthreads>>>(ptr, n, value, d_result);
cuda_try(cudaGetLastError());
int h_result = -1;
cuda_try(cudaMemcpy(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost));
cuda_try(cudaFree(d_result));
EXPECT(h_result == 0, "kernel readback mismatch");
}
void test_evaluate_blocked_even()
{
printf("Testing evaluate_localized_placement (even blocked split)...\n");
const size_t block_size = 2 * 1024 * 1024;
const size_t n = 2 * block_size; // bytes, elemsize 1: exactly 2 blocks
const dim4 data_dims(n);
auto grid = make_device_grid(1, 2);
auto stats =
evaluate_localized_placement(grid, &blocked_partition_custom<0>::get_executor, data_dims, 1, 10, block_size);
EXPECT(stats.total_bytes == n);
EXPECT(stats.vm_bytes == n);
EXPECT(stats.block_size == block_size);
EXPECT(stats.nblocks == 2);
// One block per place, and the split is block-aligned: every probe agrees
EXPECT(stats.nallocs == 2);
EXPECT(stats.accuracy() == 1.0);
size_t total = 0;
for (const auto& entry : stats.bytes_per_place)
{
total += entry.second;
}
EXPECT(total == n);
printf(" evaluate (even blocked) test PASSED\n");
}
void test_evaluate_straddling_block()
{
printf("Testing evaluate_localized_placement (majority tie-breaking)...\n");
// 2.5 MiB of one-byte elements blocked over 2 places: each place owns
// 1.25 MiB, so block 0 straddles the two owners 62.5%/37.5%. The majority
// vote must keep it on place 0, and the accuracy must reflect the
// straddling (structurally: the minority share is large enough that all-64
// probes landing in the majority has probability ~1e-13 for any uniform
// sampler, so the checks do not depend on the stdlib's distribution).
const size_t block_size = 2 * 1024 * 1024;
const size_t n = 5 * 1024 * 1024 / 2;
const dim4 data_dims(n);
auto grid = make_device_grid(1, 2);
auto stats =
evaluate_localized_placement(grid, &blocked_partition_custom<0>::get_executor, data_dims, 1, 64, block_size);
EXPECT(stats.nblocks == 2);
EXPECT(stats.nallocs == 2); // majority breaks the tie: block 0 and block 1 differ
EXPECT(stats.accuracy() < 1.0);
EXPECT(stats.accuracy() > 0.7); // expected ~0.8 (block 0 ~62.5% local, block 1 fully local)
// The decision procedure is seeded: evaluating twice gives the same stats
auto stats2 =
evaluate_localized_placement(grid, &blocked_partition_custom<0>::get_executor, data_dims, 1, 64, block_size);
EXPECT(stats.matching_samples == stats2.matching_samples);
EXPECT(stats.total_samples == stats2.total_samples);
EXPECT(stats.bytes_per_place == stats2.bytes_per_place);
printf(" evaluate (majority tie-breaking) test PASSED\n");
}
void test_evaluate_cute_matches_mapper()
{
printf("Testing evaluate_localized_placement (cute partition vs mapper)...\n");
const size_t block_size = 2 * 1024 * 1024;
const size_t n = 4 * block_size;
const dim4 data_dims(n);
auto grid = make_device_grid(1, 2);
auto part = make_partition(data_dims, partition_spec{blocked<0>}, grid.get_dims());
auto stats_mapper =
evaluate_localized_placement(grid, &blocked_partition_custom<0>::get_executor, data_dims, 1, 10, block_size);
auto stats_cute = evaluate_localized_placement(grid, part, 1, 10, block_size);
EXPECT(stats_mapper.nblocks == stats_cute.nblocks);
EXPECT(stats_mapper.nallocs == stats_cute.nallocs);
EXPECT(stats_mapper.bytes_per_place == stats_cute.bytes_per_place);
EXPECT(stats_mapper.matching_samples == stats_cute.matching_samples);
printf(" evaluate (cute vs mapper) test PASSED\n");
}
void test_shaped_alloc_callback_composite(int ndevs)
{
printf("Testing shaped allocation on a partition_fn_t composite place...\n");
const size_t n = 1024 * 1024; // ints
const dim4 data_dims(n);
auto grid = make_device_grid(ndevs, 2);
data_place dp = data_place::composite(blocked_partition_custom<0>{}, grid);
// The byte-count allocate cannot know the tensor geometry: it must throw
bool thrown = false;
try
{
dp.allocate(static_cast<::std::ptrdiff_t>(n * sizeof(int)));
}
catch (const ::std::runtime_error&)
{
thrown = true;
}
EXPECT(thrown, "byte-count allocate on a composite place must throw");
void* ptr = dp.allocate_nd(data_dims, sizeof(int));
EXPECT(ptr != nullptr);
write_and_check(static_cast<int*>(ptr), n, 17);
dp.deallocate(ptr, n * sizeof(int));
printf(" shaped allocation (callback composite) test PASSED\n");
}
void test_shaped_alloc_cute_composite(int ndevs)
{
printf("Testing shaped allocation on a cute_partition composite place...\n");
const size_t n = 1024 * 1024; // ints
const dim4 data_dims(n);
auto grid = make_device_grid(ndevs, 2);
auto part = make_partition(data_dims, partition_spec{blocked<0>}, grid.get_dims());
data_place dp = make_composite_data_place(grid, part);
// The partition is specific to one tensor: other extents must be rejected
bool thrown = false;
try
{
dp.allocate_nd(dim4(n / 2), sizeof(int));
}
catch (const ::std::invalid_argument&)
{
thrown = true;
}
EXPECT(thrown, "extent mismatch with the partition must throw");
void* ptr = dp.allocate_nd(data_dims, sizeof(int));
EXPECT(ptr != nullptr);
write_and_check(static_cast<int*>(ptr), n, 41);
dp.deallocate(ptr, n * sizeof(int));
printf(" shaped allocation (cute composite) test PASSED\n");
}
void test_shaped_alloc_overflow(int ndevs)
{
const size_t huge = ::std::numeric_limits<size_t>::max();
// (2^64-1)^2 wraps to 1: an unchecked size computation would hand back a
// one-byte allocation for an astronomically large tensor
const dim4 wrapping_dims(huge, huge, 1, 1);
auto expect_invalid = [](const data_place& dp, dim4 dims, size_t elemsize) {
bool thrown = false;
try
{
dp.allocate_nd(dims, elemsize);
}
catch (const ::std::invalid_argument&)
{
thrown = true;
}
EXPECT(thrown, "overflowing geometry must throw invalid_argument");
};
data_place dev = data_place::device(0);
expect_invalid(dev, wrapping_dims, 1);
// elemsize participates in the product too
expect_invalid(dev, dim4(huge, 1, 1, 1), 2);
// A representable product that exceeds PTRDIFF_MAX must also be rejected
expect_invalid(dev, dim4(size_t{1} << 62, 2, 1, 1), 1);
// On a composite place the wrapped geometry used to reach the blocked
// partitioner with a zero part_size and kill the process with SIGFPE
auto grid = make_device_grid(ndevs, 2);
data_place c = data_place::composite(blocked_partition_custom<1>{}, grid);
expect_invalid(c, wrapping_dims, 1);
}
void test_multi_gpu_residency(int ndevs)
{
if (ndevs < 2)
{
printf("Skipping multi-GPU residency test (requires 2+ devices)\n");
return;
}
if (!vmm_supported(1))
{
printf("Skipping multi-GPU residency test (device 1 lacks VMM support)\n");
return;
}
printf("Testing multi-GPU residency of a blocked shaped allocation...\n");
// Query the allocation granularity so each place owns a whole number of blocks
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.location.id = 0;
size_t granularity = cuda_try<cuMemGetAllocationGranularity>(&prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM);
const size_t n = 2 * granularity / sizeof(int); // one block per place
const dim4 data_dims(n);
::std::vector<exec_place> places;
places.push_back(exec_place::device(0));
places.push_back(exec_place::device(1));
auto grid = make_grid(mv(places));
data_place dp = data_place::composite(blocked_partition_custom<0>{}, grid);
void* ptr = dp.allocate_nd(data_dims, sizeof(int));
EXPECT(ptr != nullptr);
// Each half of the range must be physically backed by its owner
for (int half = 0; half < 2; half++)
{
int ordinal = -1;
CUdeviceptr probe_ptr = reinterpret_cast<CUdeviceptr>(ptr) + static_cast<size_t>(half) * granularity;
cuda_try(cuPointerGetAttribute(&ordinal, CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL, probe_ptr));
EXPECT(ordinal == half, "block is not resident on the place that owns it");
}
// Peer path: touch the whole range (including device-0-owned blocks) from
// device 1, which exercises the cuMemSetAccess mappings
int peer_01 = cuda_try<cudaDeviceCanAccessPeer>(0, 1);
int peer_10 = cuda_try<cudaDeviceCanAccessPeer>(1, 0);
if (peer_01 && peer_10)
{
cuda_try(cudaSetDevice(1));
write_and_check(static_cast<int*>(ptr), n, 73);
cuda_try(cudaSetDevice(0));
}
else
{
printf(" (peer access unavailable between devices 0 and 1: cross-device touch skipped)\n");
}
dp.deallocate(ptr, n * sizeof(int));
printf(" multi-GPU residency test PASSED\n");
}
} // namespace
int main()
{
int ndevs = 0;
if (cudaGetDeviceCount(&ndevs) != cudaSuccess || ndevs == 0)
{
printf("Skipping placement tests: no CUDA device\n");
return 0;
}
cuda_try(cudaSetDevice(0));
cuda_try(cudaFree(nullptr));
if (!vmm_supported())
{
printf("Skipping placement tests: VMM is not supported on this machine\n");
return 0;
}
printf("=== Testing placement evaluation and shaped allocation ===\n\n");
test_evaluate_blocked_even();
test_evaluate_straddling_block();
test_evaluate_cute_matches_mapper();
test_shaped_alloc_callback_composite(ndevs);
test_shaped_alloc_cute_composite(ndevs);
test_shaped_alloc_overflow(ndevs);
test_multi_gpu_residency(ndevs);
printf("\n=== All placement tests PASSED ===\n");
return 0;
}

181
cccl_upstream/cudax/test/places/stream_pool.cu
View File

@@ -1,181 +0,0 @@
//===----------------------------------------------------------------------===//
//
// Part of CUDASTF in CUDA C++ Core Libraries,
// under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief Tests for the standalone stream pool functionality in exec_place.
*
* Verifies that exec_place::pick_stream(resources) works without a CUDASTF
* context, returning valid CUDA streams from the per-place stream pool
* lazily created inside an `exec_place_resources` registry.
*/
#include <cuda/experimental/__places/places.cuh>
using namespace cuda::experimental::places;
__global__ void increment_kernel(int* data, int n)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid < n)
{
data[tid] += 1;
}
}
// Streams returned by pick_stream(resources) are owned by the supplied
// `exec_place_resources` registry (round-robin, lazily created). Callers
// must NOT destroy them; their lifetime ends with the registry.
void test_basic_pick_stream()
{
exec_place_resources resources;
exec_place place = exec_place::current_device();
cudaStream_t stream = place.pick_stream(resources);
_CCCL_ASSERT(stream != nullptr, "pick_stream must return a valid stream");
int current_device;
cuda_try(cudaGetDevice(&current_device));
_CCCL_ASSERT(get_device_from_stream(stream) == current_device, "stream must belong to the current device");
fprintf(stderr, "test_basic_pick_stream: PASSED\n");
}
void test_pick_stream_computation_hint()
{
exec_place_resources resources;
exec_place place = exec_place::current_device();
cudaStream_t compute_stream = place.pick_stream(resources, true);
cudaStream_t transfer_stream = place.pick_stream(resources, false);
_CCCL_ASSERT(compute_stream != nullptr, "compute stream must be valid");
_CCCL_ASSERT(transfer_stream != nullptr, "transfer stream must be valid");
fprintf(stderr, "test_pick_stream_computation_hint: PASSED\n");
}
void test_pick_stream_specific_device(int ndevs)
{
if (ndevs < 2)
{
fprintf(stderr, "test_pick_stream_specific_device: skipped (need >= 2 devices)\n");
return;
}
exec_place_resources resources;
for (int d = 0; d < ndevs && d < 2; d++)
{
exec_place dev = exec_place::device(d);
cudaStream_t stream = dev.pick_stream(resources);
_CCCL_ASSERT(stream != nullptr, "stream must be valid");
_CCCL_ASSERT(get_device_from_stream(stream) == d, "stream must belong to the requested device");
}
fprintf(stderr, "test_pick_stream_specific_device: PASSED\n");
}
void test_launch_kernel_on_picked_stream()
{
exec_place_resources resources;
exec_place place = exec_place::current_device();
cudaStream_t stream = place.pick_stream(resources);
constexpr int N = 256;
int* d_data;
cuda_try(cudaMallocAsync(&d_data, N * sizeof(int), stream));
cuda_try(cudaMemsetAsync(d_data, 0, N * sizeof(int), stream));
increment_kernel<<<1, N, 0, stream>>>(d_data, N);
int h_data[N];
cuda_try(cudaMemcpyAsync(h_data, d_data, N * sizeof(int), cudaMemcpyDeviceToHost, stream));
cuda_try(cudaStreamSynchronize(stream));
for (const auto& v : h_data)
{
_CCCL_ASSERT(v == 1, "kernel result mismatch");
}
cuda_try(cudaFreeAsync(d_data, stream));
cuda_try(cudaStreamSynchronize(stream));
fprintf(stderr, "test_launch_kernel_on_picked_stream: PASSED\n");
}
void test_round_robin_streams()
{
exec_place_resources resources;
exec_place place = exec_place::current_device();
cudaStream_t first = place.pick_stream(resources);
cudaStream_t second = place.pick_stream(resources);
_CCCL_ASSERT(first != nullptr, "first stream must be valid");
_CCCL_ASSERT(second != nullptr, "second stream must be valid");
fprintf(stderr, "test_round_robin_streams: PASSED\n");
}
// Two independent registries must hand out independent streams for the same
// place: this is the property that lets multiple STF contexts (or multiple
// threads with their own `async_resources_handle`) share a device without
// touching each other's stream pools.
void test_two_handles_isolation()
{
exec_place_resources r1;
exec_place_resources r2;
exec_place place = exec_place::current_device();
cudaStream_t s1 = place.pick_stream(r1);
cudaStream_t s2 = place.pick_stream(r2);
_CCCL_ASSERT(s1 != nullptr && s2 != nullptr, "streams must be valid");
_CCCL_ASSERT(s1 != s2, "different registries must own different streams");
_CCCL_ASSERT(r1.size() == 1 && r2.size() == 1, "each registry should hold exactly one entry");
fprintf(stderr, "test_two_handles_isolation: PASSED\n");
}
// A registry destroyed before another is created must release its CUDA
// streams; subsequent device-reset followed by a fresh registry must not
// observe any stale handles. This is the property that lets pytest sessions
// survive `cuda.bindings.driver.cuDevicePrimaryCtxReset` between tests.
void test_reset_survives_with_fresh_registry()
{
{
exec_place_resources resources;
cudaStream_t stream = exec_place::current_device().pick_stream(resources);
cuda_try(cudaStreamSynchronize(stream));
}
// Old registry destroyed -> its cached streams are gone -> reset is safe.
cuda_try(cudaDeviceReset());
exec_place_resources resources;
cudaStream_t stream = exec_place::current_device().pick_stream(resources);
_CCCL_ASSERT(stream != nullptr, "fresh registry must produce a valid stream after reset");
cuda_try(cudaStreamSynchronize(stream));
fprintf(stderr, "test_reset_survives_with_fresh_registry: PASSED\n");
}
int main()
{
int ndevs;
cuda_try(cudaGetDeviceCount(&ndevs));
test_basic_pick_stream();
test_pick_stream_computation_hint();
test_pick_stream_specific_device(ndevs);
test_launch_kernel_on_picked_stream();
test_round_robin_streams();
test_two_handles_isolation();
test_reset_survives_with_fresh_registry();
}