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

[CCCL] Add missing CCCL components: c2h, nvbench_helper, cmake, cudax, AGENTS.md

Browse Source 
Added 863 files from NVIDIA/cccl sparse checkout:
- c2h/ (27 files): Catch2 test helpers — generators, validators, runner
- nvbench_helper/ (10 files): Benchmark harness utilities
- cmake/ (29 files): CMake presets and build helpers
- cudax/ (794 files): Experimental CUDA extensions
- AGENTS.md: NVIDIA's official AI agent instructions for CCCL
- CMakePresets.json: Standardized build configurations
- cccl-version.json: Version tracking

Also added CCCL_ASSET_MAP.md mapping all 4295 CCCL files to
competition value and PRD items.

cccl_upstream now covers 100% of competition-critical assets:
- 27 tuning headers (SM80/90/100 benchmark data)
- 32 dispatch headers (algorithm implementations)
- 60 Thrust examples (correctness verification)
- 217 CUB Catch2 tests (regression matrix)
- 153 CUB benchmarks (parameter space search)
- 18 CUB examples (API verification)
- 27 test helpers + benchmark harness
- 794 cudax experimental extensions
This commit is contained in:
muh-bot
2026-08-06 02:14:18 +00:00
parent b0d597363a
commit dedf08166a
864 changed files with 174321 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
cccl_upstream/cudax/test/stf/places/affinity_gc.cu Normal file
View File

@@ -0,0 +1,42 @@
//===----------------------------------------------------------------------===//
//
// 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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
#include <cuda/experimental/__places/place_partition.cuh>
#include <cuda/experimental/__stf/internal/stf_places_partition_into_stf.cuh>
#include <cuda/experimental/stf.cuh>
using namespace cuda::experimental::stf;
#if _CCCL_CTK_AT_LEAST(12, 4)
/**
* @brief Test green context partition and affinity: partition by green_context, push/pop affinity per subplace.
*/
void test_green_ctx_affinity()
{
async_resources_handle handle;
for (auto p : place_partition(exec_place::current_device(), handle, place_partition_scope::green_context))
{
handle.push_affinity(::std::make_shared<exec_place>(p));
_CCCL_ASSERT(handle.current_affinity().size() == 1, "invalid value");
handle.pop_affinity();
}
}
#endif // _CCCL_CTK_AT_LEAST(12, 4)
int main()
{
#if _CCCL_CTK_BELOW(12, 4)
fprintf(stderr, "Green contexts are not supported by this version of CUDA: skipping test.\n");
return 0;
#else // ^^^ _CCCL_CTK_BELOW(12, 4) ^^^ / vvv _CCCL_CTK_AT_LEAST(12, 4) vvv
test_green_ctx_affinity();
return 0;
#endif // ^^^ _CCCL_CTK_AT_LEAST(12, 4) ^^^
}

75
cccl_upstream/cudax/test/stf/places/cuda_stream_place.cu Normal file
View File

@@ -0,0 +1,75 @@
//===----------------------------------------------------------------------===//
//
// 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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
*
* @brief An AXPY kernel using an exec place attached to a specific CUDA stream
*
*/
#include <cuda/experimental/stf.cuh>
#include "nvtx3/nvToolsExtCudaRt.h"
using namespace cuda::experimental::stf;
double X0(size_t i)
{
return sin((double) i);
}
double Y0(size_t i)
{
return cos((double) i);
}
int main()
{
cudaStream_t stream;
cuda_safe_call(cudaStreamCreate(&stream));
nvtxNameCudaStreamA(stream, "user stream");
// context ctx;
stream_ctx ctx;
const size_t N = 16;
double X[N], Y[N];
for (size_t i = 0; i < N; i++)
{
X[i] = X0(i);
Y[i] = Y0(i);
}
double alpha = 3.14;
auto lX = ctx.logical_data(X);
auto lY = ctx.logical_data(Y);
/* Compute Y = Y + alpha X on the user stream */
auto where = exec_place::cuda_stream(stream);
for (size_t iter = 0; iter < 20; iter++)
{
ctx.parallel_for(where, lX.shape(), lX.read(), lY.rw())->*[alpha] __device__(size_t i, auto x, auto y) {
y(i) += alpha * x(i);
};
}
ctx.finalize();
for (size_t i = 0; i < N; i++)
{
assert(fabs(Y[i] - (Y0(i) + 2 * 10.0 * alpha * X0(i))) < 0.0001);
assert(fabs(X[i] - X0(i)) < 0.0001);
}
cuda_safe_call(cudaStreamDestroy(stream));
}

230
cccl_upstream/cudax/test/stf/places/cute_block_owners.cu Normal file
View File

@@ -0,0 +1,230 @@
//===----------------------------------------------------------------------===//
//
// 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 cute_partition_descriptor::try_block_owners — the analytic page
* plan — cross-validated against a brute-force per-element owner scan, and
* exercised end-to-end through a localized_array allocation.
*/
#include <cuda/experimental/__stf/localization/composite_slice.cuh>
#include <cuda/experimental/__stf/utility/unittest.cuh>
#include <map>
using namespace cuda::experimental::stf;
using namespace cuda::experimental::places;
namespace
{
// Brute force: owner of every byte via per-element owner() queries.
::std::vector<pos4> brute_block_owners(
const cute_partition_descriptor& part,
size_t block_size_bytes,
size_t elemsize,
size_t total_elems,
dim4 data_dims,
size_t* misplaced_bytes)
{
const size_t total_bytes = total_elems * elemsize;
const size_t nblocks = (total_bytes + block_size_bytes - 1) / block_size_bytes;
*misplaced_bytes = 0;
::std::vector<pos4> owners;
for (size_t b = 0; b < nblocks; b++)
{
::std::map<::std::array<ssize_t, 4>, size_t> census;
const size_t lo = b * block_size_bytes;
const size_t hi = ::std::min((b + 1) * block_size_bytes, total_bytes);
for (size_t byte = lo; byte < hi; byte++)
{
const pos4 o = part.owner(data_dims.index_to_pos(byte / elemsize));
census[{o.x, o.y, o.z, o.t}]++;
}
size_t best = 0;
::std::array<ssize_t, 4> bp{};
for (const auto& e : census)
{
if (e.second > best)
{
best = e.second;
bp = e.first;
}
}
owners.push_back(pos4(bp[0], bp[1], bp[2], bp[3]));
*misplaced_bytes += (hi - lo) - best;
}
return owners;
}
void check_case(dim4 data_dims, const ::std::vector<dim_spec>& spec, dim4 grid_dims, size_t elemsize, size_t block)
{
const auto part = make_partition_descriptor(data_dims, spec, grid_dims);
const size_t total_elems = data_dims.size();
size_t misplaced = 0;
auto plan = part.try_block_owners(block, elemsize, &misplaced);
if (!plan)
{
return; // dense: sampled tier takes over (covered elsewhere)
}
size_t brute_misplaced = 0;
const auto brute = brute_block_owners(part, block, elemsize, total_elems, data_dims, &brute_misplaced);
EXPECT(plan->size() == brute.size());
EXPECT(misplaced == brute_misplaced);
// The runs path must agree with the brute-force census wherever the
// strict quotient exists: zero misplacement, and every block of every run
// owned by the brute owner. Where it declines, the census must have found
// at least one straddled (misplaced) block or an in-block boundary.
if (const auto runs = part.try_block_runs(block, elemsize))
{
EXPECT(misplaced == 0);
size_t covered = 0;
for (const auto& r : *runs)
{
// strict tiling: each run starts where the previous ended and stays
// in bounds (covered == total alone would accept overlap/gap pairs)
EXPECT(r.first_block == covered);
EXPECT(r.num_blocks > 0);
EXPECT(r.num_blocks <= brute.size() - covered);
for (size_t b = r.first_block; b < r.first_block + r.num_blocks; b++)
{
EXPECT(brute[b] == r.owner);
}
covered += r.num_blocks;
}
EXPECT(covered == brute.size());
}
for (size_t b = 0; b < brute.size(); b++)
{
// majority may tie: accept the analytic owner iff its byte count ties the
// brute majority; equality of misplaced bytes above already pins that.
if (!((*plan)[b] == brute[b]))
{
// re-census this block for the analytic owner's count
size_t analytic_bytes = 0, brute_bytes = 0;
const size_t lo = b * block;
const size_t hi = ::std::min((b + 1) * block, total_elems * elemsize);
for (size_t byte = lo; byte < hi; byte++)
{
const pos4 o = part.owner(data_dims.index_to_pos(byte / elemsize));
analytic_bytes += (o == (*plan)[b]);
brute_bytes += (o == brute[b]);
}
EXPECT(analytic_bytes == brute_bytes); // a genuine tie
}
}
}
void property_suite()
{
const ::std::vector<size_t> grids = {1, 2, 3, 4, 6, 8};
for (size_t g : grids)
{
for (size_t elemsize : {2, 4})
{
for (size_t block : {16, 64, 256})
{
// 1-D blocked / cyclic / block_cyclic
check_case(dim4(13), {{dim_policy::blocked, 0, 0}}, dim4(g), elemsize, block);
check_case(dim4(48), {{dim_policy::blocked, 0, 0}}, dim4(g), elemsize, block);
check_case(dim4(64), {{dim_policy::cyclic, 0, 0}}, dim4(g), elemsize, block);
check_case(dim4(64), {{dim_policy::block_cyclic, 0, 4}}, dim4(g), elemsize, block);
// 2-D: outer blocked, inner whole (expert-major shape)
check_case(dim4(12, 16), {{dim_policy::blocked, 0, 0}, {}}, dim4(g), elemsize, block);
// 2-D: inner cyclic
check_case(dim4(12, 16), {{}, {dim_policy::cyclic, 0, 0}}, dim4(g), elemsize, block);
}
}
}
// tiled 2-D on a (2,2) grid + 3-D with a middle whole dim
for (size_t elemsize : {2, 4})
{
for (size_t block : {16, 64, 256})
{
check_case(dim4(12, 16), {{dim_policy::blocked, 0, 0}, {dim_policy::blocked, 1, 0}}, dim4(2, 2), elemsize, block);
check_case(
dim4(6, 8, 4), {{dim_policy::blocked, 0, 0}, {}, {dim_policy::blocked, 1, 0}}, dim4(2, 2), elemsize, block);
}
}
// dense detection: element-cyclic far below the block size must decline
const auto dense = make_partition_descriptor(dim4(1 << 22), {{dim_policy::cyclic, 0, 0}}, dim4(2));
size_t mis = 0;
EXPECT(!dense.try_block_owners(2 * 1024 * 1024, 4, &mis).has_value());
}
void end_to_end_allocation()
{
// Exercise the provider path through a real VMM allocation: a 2-place
// grid on the current device (placement plumbing, merge, stats).
const auto d0 = exec_place::device(cuda_try<cudaGetDevice>());
::std::vector<exec_place> places{d0, d0};
const auto grid = make_grid(mv(places));
const size_t n = (6 * 1024 * 1024) + 512; // forces one straddle block
const dim4 data_dims(n);
const auto part = make_partition_descriptor(data_dims, {{dim_policy::blocked, 0, 0}}, grid.get_dims());
localized_array arr(
grid, make_partition_placement_provider(part, data_dims, data_dims.size(), sizeof(int)), n, sizeof(int), data_dims);
const auto& st = arr.get_stats();
// exact plan: sample counters hold byte counts
EXPECT(st.total_samples == n * sizeof(int));
size_t mis = 0;
const auto owners = part.try_block_owners(st.block_size, sizeof(int), &mis);
EXPECT(owners.has_value() == true);
EXPECT(st.matching_samples == st.total_samples - mis);
EXPECT(st.nallocs <= 3); // two shards + at most one straddle merge break
}
void malformed_providers_throw()
{
const auto d0 = exec_place::device(cuda_try<cudaGetDevice>());
::std::vector<exec_place> places{d0, d0};
const auto grid = make_grid(mv(places));
const size_t n = 4 * 1024 * 1024; // 16 MB of int: 8 blocks at 2 MB
const dim4 dims(n);
auto expect_throw = [&](auto&& provider) {
bool thrown = false;
try
{
localized_array arr(grid, provider, n, sizeof(int), dims);
}
catch (const ::std::invalid_argument&)
{
thrown = true;
}
EXPECT(thrown);
};
// gap: second run skips a block
expect_throw([](size_t, size_t nblocks, localized_stats&) {
return ::std::vector<block_run>{{pos4(0), 0, 1}, {pos4(1), 2, nblocks - 2}};
});
// short: does not cover the final block
expect_throw([](size_t, size_t nblocks, localized_stats&) {
return ::std::vector<block_run>{{pos4(0), 0, nblocks - 1}};
});
// zero-length run
expect_throw([](size_t, size_t nblocks, localized_stats&) {
return ::std::vector<block_run>{{pos4(0), 0, 0}, {pos4(1), 0, nblocks}};
});
}
} // namespace
int main()
{
property_suite();
end_to_end_allocation();
malformed_providers_throw();
printf("cute_block_owners: all checks passed\n");
return 0;
}

294
cccl_upstream/cudax/test/stf/places/cute_parallel_for.cu Normal file
View File

@@ -0,0 +1,294 @@
//===----------------------------------------------------------------------===//
//
// 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 parallel_for over a grid driven by a cute_partition instance
*
* The same user-facing parallel_for entry point accepts value-defined
* partitioners: the partition decides both the kernel decomposition
* (per-place sub-shapes) and the data placement (composite data place backed
* by the partition).
*/
#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
#include <cuda/experimental/__stf/localization/composite_slice.cuh>
#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
using namespace cuda::experimental::stf;
namespace
{
void test_cute_composite_cache(const exec_place& grid)
{
const size_t n = 4096;
const dim4 data_dims(n);
const auto part = make_partition(data_dims, partition_spec{blocked<0>}, grid.get_dims());
const auto place = cuda::experimental::places::make_composite_data_place(grid, part);
const auto delinearize = [data_dims](size_t ind) {
return data_dims.index_to_pos(ind);
};
reserved::composite_slice_cache cache;
// Equal element counts are insufficient: a different tensor shape changes
// delinearization and therefore ownership.
const dim4 mismatched_dims(n / 2, 2);
const auto mismatched_delinearize = [mismatched_dims](size_t ind) {
return mismatched_dims.index_to_pos(ind);
};
bool mismatch_thrown = false;
try
{
(void) cache.get(place, mismatched_delinearize, n, sizeof(size_t), mismatched_dims);
}
catch (const ::std::invalid_argument&)
{
mismatch_thrown = true;
}
EXPECT(mismatch_thrown);
auto [first, first_prereqs] = cache.get(place, delinearize, n, sizeof(size_t), data_dims);
EXPECT(first_prereqs.empty());
const auto first_base = first->get_base_ptr();
cache.put(place, mv(first), first_prereqs, n, sizeof(size_t), data_dims);
// A separately constructed but equivalent place must find the same cached
// VMM allocation through the value-keyed CuTe pool.
const auto equivalent_part = make_partition(data_dims, partition_spec{blocked<0>}, grid.get_dims());
const auto equivalent_place = cuda::experimental::places::make_composite_data_place(grid, equivalent_part);
auto [second, second_prereqs] = cache.get(equivalent_place, delinearize, n, sizeof(size_t), data_dims);
EXPECT(second_prereqs.empty());
EXPECT(second->get_base_ptr() == first_base);
cache.put(equivalent_place, mv(second), second_prereqs, n, sizeof(size_t), data_dims);
EXPECT(cache.deinit().empty());
}
void test_static_codegen_parity(stream_ctx& ctx, const exec_place& grid)
{
const size_t nx = 64;
const size_t ny = 32;
auto typed_data = ctx.logical_data(shape_of<slice<size_t, 2>>(nx, ny));
auto classic_data = ctx.logical_data(shape_of<slice<size_t, 2>>(nx, ny));
const auto part = make_partition(dim4(nx, ny), partition_spec{whole, blocked<0>}, grid.get_dims());
auto write = [] _CCCL_DEVICE(size_t x, size_t y, auto values) {
values(x, y) = x + 100 * y;
};
ctx.parallel_for(part, grid, typed_data.shape(), typed_data.write())->*decltype(write)(write);
ctx.parallel_for(blocked_partition(), grid, classic_data.shape(), classic_data.write())->*decltype(write)(write);
ctx.host_launch(typed_data.read(), classic_data.read())->*[=](auto typed, auto classic) {
for (size_t y = 0; y < ny; y++)
{
for (size_t x = 0; x < nx; x++)
{
EXPECT(typed(x, y) == x + 100 * y);
EXPECT(classic(x, y) == typed(x, y));
}
}
};
}
void test_cute_graph_backend(const exec_place& grid)
{
const size_t n = 1023;
graph_ctx ctx;
auto data = ctx.logical_data(shape_of<slice<size_t>>(n));
const auto part = make_partition(dim4(n), partition_spec{blocked<0>}, grid.get_dims());
ctx.parallel_for(part, grid, data.shape(), data.write())->*[] _CCCL_DEVICE(size_t i, auto values) {
values(i) = 5 * i + 3;
};
ctx.host_launch(data.read())->*[=](auto values) {
for (size_t i = 0; i < n; i++)
{
EXPECT(values(i) == 5 * i + 3);
}
};
ctx.finalize();
}
} // namespace
int main()
{
int ndevs;
cuda_safe_call(cudaGetDeviceCount(&ndevs));
stream_ctx ctx;
// A grid of two places (same device when only one GPU is present)
::std::vector<exec_place> places;
places.push_back(exec_place::device(0));
places.push_back(exec_place::device(ndevs > 1 ? 1 : 0));
auto grid = make_grid(mv(places));
test_cute_composite_cache(grid);
test_static_codegen_parity(ctx, grid);
// 1-D: dimension 0 blocked over the grid
{
const size_t n = 1024 * 1024;
auto lA = ctx.logical_data(shape_of<slice<size_t>>(n));
auto part = make_partition(dim4(n), partition_spec{blocked<0>}, grid.get_dims());
ctx.parallel_for(part, grid, lA.shape(), lA.write())->*[] _CCCL_DEVICE(size_t i, auto a) {
a(i) = 3 * i + 7;
};
ctx.host_launch(lA.read())->*[&](auto a) {
for (size_t i = 0; i < n; i++)
{
EXPECT(a(i) == 3 * i + 7);
}
};
}
// 3-D: dimension 1 blocked over the grid (the per-dimension expressiveness
// the classic blocked_partition cannot provide)
{
const size_t nx = 32, ny = 64, nz = 16;
auto lB = ctx.logical_data(shape_of<slice<size_t, 3>>(nx, ny, nz));
auto part = make_partition(dim4(nx, ny, nz), partition_spec{whole, blocked<0>, whole}, grid.get_dims());
ctx.parallel_for(part, grid, lB.shape(), lB.write())->*[] _CCCL_DEVICE(size_t x, size_t y, size_t z, auto b) {
b(x, y, z) = x + 100 * y + 10000 * z;
};
ctx.host_launch(lB.read())->*[&](auto b) {
for (size_t x = 0; x < nx; x++)
{
for (size_t y = 0; y < ny; y++)
{
for (size_t z = 0; z < nz; z++)
{
EXPECT(b(x, y, z) == x + 100 * y + 10000 * z);
}
}
}
};
}
// The classic stateless partitioners keep working through the same entry
{
const size_t n = 4096;
auto lC = ctx.logical_data(shape_of<slice<size_t>>(n));
ctx.parallel_for(blocked_partition(), grid, lC.shape(), lC.write())->*[] _CCCL_DEVICE(size_t i, auto c) {
c(i) = i;
};
ctx.host_launch(lC.read())->*[&](auto c) {
for (size_t i = 0; i < n; i++)
{
EXPECT(c(i) == i);
}
};
}
// Uneven extents: the padding phantoms are excluded by the sub-shape's
// predicate (CuTe predication), so odd sizes work end to end
{
const size_t n = 1023; // not divisible by 2 places
auto lD = ctx.logical_data(shape_of<slice<size_t>>(n));
auto part = make_partition(dim4(n), partition_spec{blocked<0>}, grid.get_dims());
ctx.parallel_for(part, grid, lD.shape(), lD.write())->*[] _CCCL_DEVICE(size_t i, auto d) {
d(i) = 2 * i + 1;
};
ctx.host_launch(lD.read())->*[&](auto d) {
for (size_t i = 0; i < n; i++)
{
EXPECT(d(i) == 2 * i + 1);
}
};
}
// Interior region: the box is a region within the tensor the partition was
// built for; each place computes its owned coordinates restricted to the
// box, and the boundary stays untouched
{
const size_t nx = 64, ny = 32;
auto lE = ctx.logical_data(shape_of<slice<size_t, 2>>(nx, ny));
auto part = make_partition(dim4(nx, ny), partition_spec{whole, blocked<0>}, grid.get_dims());
ctx.parallel_for(part, grid, lE.shape(), lE.write())->*[] _CCCL_DEVICE(size_t x, size_t y, auto e) {
e(x, y) = 7;
};
box interior({1ul, nx - 1}, {1ul, ny - 1});
ctx.parallel_for(part, grid, interior, lE.rw())->*[] _CCCL_DEVICE(size_t x, size_t y, auto e) {
e(x, y) = 100 + x + y;
};
ctx.host_launch(lE.read())->*[&](auto e) {
for (size_t x = 0; x < nx; x++)
{
for (size_t y = 0; y < ny; y++)
{
const bool inside = (x >= 1 && x < nx - 1 && y >= 1 && y < ny - 1);
EXPECT(e(x, y) == (inside ? 100 + x + y : 7));
}
}
};
}
// Boundary-style thin regions: iterate the face with a classic scale-free
// partitioner (tight, no discarded lanes) while keeping placement on the
// cute composite through explicit deps. This relies on separately
// constructed equal partitions producing the same composite identity -
// guarded here.
{
const size_t nx = 64, ny = 32;
auto lF = ctx.logical_data(shape_of<slice<size_t, 2>>(nx, ny));
auto part = make_partition(dim4(nx, ny), partition_spec{whole, blocked<0>}, grid.get_dims());
auto dist = cuda::experimental::places::make_composite_data_place(grid, part);
const auto equivalent_part = make_partition(dim4(nx, ny), partition_spec{whole, blocked<0>}, grid.get_dims());
EXPECT(dist == cuda::experimental::places::make_composite_data_place(grid, equivalent_part),
"cute composites from equal partitions must compare equal");
// Volumetric pass placed and decomposed by the partition
ctx.parallel_for(part, grid, lF.shape(), lF.write())->*[] _CCCL_DEVICE(size_t x, size_t y, auto f) {
f(x, y) = 1;
};
// Face update: classic iteration over the thin box, same placement
box face({0ul, nx}, {0ul, 1ul});
ctx.parallel_for(blocked_partition(), grid, face, lF.rw(dist))->*[] _CCCL_DEVICE(size_t x, size_t y, auto f) {
f(x, y) = 42;
};
ctx.host_launch(lF.read())->*[&](auto f) {
for (size_t x = 0; x < nx; x++)
{
for (size_t y = 0; y < ny; y++)
{
EXPECT(f(x, y) == (y == 0 ? 42 : 1));
}
}
};
}
ctx.finalize();
test_cute_graph_backend(grid);
printf("cute_parallel_for: all checks passed\n");
return 0;
}

56
cccl_upstream/cudax/test/stf/places/execution_policy_kernel_launch_test.cu Normal file
View File

@@ -0,0 +1,56 @@
//===----------------------------------------------------------------------===//
//
// 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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
//! file
//! !brief Check that multi-level launch specification are fulfilled
#include <cuda/experimental/stf.cuh>
#include <cassert>
#include <iostream>
using namespace cuda::experimental::stf;
int main()
{
stream_ctx ctx;
// Create a 3-level thread hierarchy specification that would expose the bug:
// Level 0: only 1 device to run on CI
// Level 1: 4 blocks per device (width 4)
// Level 2: 64 threads per block (width 64)
//
auto spec = par(hw_scope::device, 1, con<4>(hw_scope::block, con<64>(hw_scope::thread)));
int test_result = 0;
auto l_test_result = ctx.logical_data(make_slice(&test_result, 1));
ctx.launch(spec, exec_place::current_device(), l_test_result.rw())->*[] __device__(auto th, auto result) {
if (th.rank() == 0)
{
bool level0_correct = (th.size(0) == 1); // device level
bool level1_correct = (th.size(1) == 1 * 4) && (gridDim.x == 4); // blocks per device
bool level2_correct = (th.size(2) == 1 * 4 * 64) && (blockDim.x == 64); // threads per block
// Set test result based on whether all levels are correct
result[0] = level0_correct && level1_correct && level2_correct ? 1 : 0;
}
};
ctx.finalize();
if (test_result != 1)
{
fprintf(stderr, "FAIL: Hierarchy dimensions are incorrect!\n");
return 1;
}
return 0;
}

87
cccl_upstream/cudax/test/stf/places/managed.cu Normal file
View File

@@ -0,0 +1,87 @@
//===----------------------------------------------------------------------===//
//
// 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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
*
* @brief An AXPY kernel implemented with a task of the CUDA stream backend
* where the task accesses managed memory from the device
*
*/
#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
#include <iostream>
using namespace cuda::experimental::stf;
__global__ void axpy(double a, slice<const double> x, slice<double> y)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int nthreads = gridDim.x * blockDim.x;
for (int i = tid; i < x.size(); i += nthreads)
{
y(i) += a * x(i);
}
}
double X0(size_t i)
{
return sin((double) i);
}
double Y0(size_t i)
{
return cos((double) i);
}
int main()
{
// Verify whether this device can access memory concurrently from CPU and GPU.
int dev;
cuda_safe_call(cudaGetDevice(&dev));
assert(dev >= 0);
cudaDeviceProp prop;
cuda_safe_call(cudaGetDeviceProperties(&prop, dev));
if (!prop.concurrentManagedAccess)
{
fprintf(stderr, "Concurrent CPU/GPU access not supported, skipping test.\n");
return 0;
}
stream_ctx ctx;
const size_t N = 16;
double X[N], Y[N];
for (size_t i = 0; i < N; i++)
{
X[i] = X0(i);
Y[i] = Y0(i);
}
double alpha = 3.14;
auto lX = ctx.logical_data(X);
auto lY = ctx.logical_data(Y);
/* Compute Y = Y + alpha X, but leave X on the host and access it with mapped memory */
ctx.task(lX.read(data_place::managed()), lY.rw())->*[&](cudaStream_t s, auto dX, auto dY) {
axpy<<<16, 128, 0, s>>>(alpha, dX, dY);
};
ctx.finalize();
for (size_t i = 0; i < N; i++)
{
assert(fabs(Y[i] - (Y0(i) + alpha * X0(i))) < 0.0001);
assert(fabs(X[i] - X0(i)) < 0.0001);
}
}

92
cccl_upstream/cudax/test/stf/places/managed_from_shape.cu Normal file
View File

@@ -0,0 +1,92 @@
//===----------------------------------------------------------------------===//
//
// 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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
*
* @brief Make sure we can automatically allocate and use data in managed memory based on their shape
*
*/
#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
#include <iostream>
using namespace cuda::experimental::stf;
__global__ void axpy(double a, slice<const double> x, slice<double> y)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int nthreads = gridDim.x * blockDim.x;
for (int i = tid; i < x.size(); i += nthreads)
{
y(i) += a * x(i);
}
}
__host__ __device__ double X0(size_t i)
{
return sin((double) i);
}
double Y0(size_t i)
{
return cos((double) i);
}
int main()
{
// Verify whether this device can access memory concurrently from CPU and GPU.
int dev;
cuda_safe_call(cudaGetDevice(&dev));
assert(dev >= 0);
cudaDeviceProp prop;
cuda_safe_call(cudaGetDeviceProperties(&prop, dev));
if (!prop.concurrentManagedAccess)
{
fprintf(stderr, "Concurrent CPU/GPU access not supported, skipping test.\n");
return 0;
}
stream_ctx ctx;
const size_t N = 16;
double Y[N];
for (size_t i = 0; i < N; i++)
{
Y[i] = Y0(i);
}
double alpha = 3.14;
auto lX = ctx.logical_data(shape_of<slice<double>>(N));
auto lY = ctx.logical_data(Y);
// Make sure X is created automatically in managed memory
ctx.parallel_for(lX.shape(), lX.write(data_place::managed()))->*[] _CCCL_DEVICE(size_t i, auto X) {
X(i) = X0(i);
};
/* Compute Y = Y + alpha X, but leave X in managed memory */
ctx.task(lX.read(data_place::managed()), lY.rw())->*[&](cudaStream_t s, auto dX, auto dY) {
axpy<<<16, 128, 0, s>>>(alpha, dX, dY);
};
ctx.host_launch(lX.read(data_place::managed()), lY.read())->*[=](auto X, auto Y) {
for (size_t i = 0; i < N; i++)
{
EXPECT(fabs(Y(i) - (Y0(i) + alpha * X0(i))) < 0.0001);
EXPECT(fabs(X(i) - X0(i)) < 0.0001);
}
};
ctx.finalize();
}

90
cccl_upstream/cudax/test/stf/places/managed_from_user.cu Normal file
View File

@@ -0,0 +1,90 @@
//===----------------------------------------------------------------------===//
//
// 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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
*
* @brief An AXPY kernel implemented with a task of the CUDA stream backend
* where the task accesses managed memory from the device. This tests
* explicitly created managed memory, and passes it to a logical data.
*/
#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
#include <iostream>
using namespace cuda::experimental::stf;
__global__ void axpy(double a, slice<const double> x, slice<double> y)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int nthreads = gridDim.x * blockDim.x;
for (int i = tid; i < x.size(); i += nthreads)
{
y(i) += a * x(i);
}
}
double X0(size_t i)
{
return sin((double) i);
}
double Y0(size_t i)
{
return cos((double) i);
}
int main()
{
// Verify whether this device can access memory concurrently from CPU and GPU.
int dev;
cuda_safe_call(cudaGetDevice(&dev));
assert(dev >= 0);
cudaDeviceProp prop;
cuda_safe_call(cudaGetDeviceProperties(&prop, dev));
if (!prop.concurrentManagedAccess)
{
fprintf(stderr, "Concurrent CPU/GPU access not supported, skipping test.\n");
return 0;
}
stream_ctx ctx;
const size_t N = 16;
double Y[N];
double* X;
cuda_safe_call(cudaMallocManaged(&X, N * sizeof(double)));
for (size_t i = 0; i < N; i++)
{
X[i] = X0(i);
Y[i] = Y0(i);
}
double alpha = 3.14;
auto lX = ctx.logical_data(make_slice(X, N), data_place::managed());
auto lY = ctx.logical_data(Y);
/* Compute Y = Y + alpha X, but leave X in managed memory */
ctx.task(lX.read(data_place::managed()), lY.rw())->*[&](cudaStream_t s, auto dX, auto dY) {
axpy<<<16, 128, 0, s>>>(alpha, dX, dY);
};
ctx.finalize();
for (size_t i = 0; i < N; i++)
{
assert(fabs(Y[i] - (Y0(i) + alpha * X0(i))) < 0.0001);
assert(fabs(X[i] - X0(i)) < 0.0001);
}
}

77
cccl_upstream/cudax/test/stf/places/non_current_device.cu Normal file
View File

@@ -0,0 +1,77 @@
//===----------------------------------------------------------------------===//
//
// 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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
#include <cuda/experimental/__stf/stream/stream_ctx.cuh>
using namespace cuda::experimental::stf;
template <typename T>
__global__ void axpy(int n, T a, const T* x, T* y)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int nthreads = gridDim.x * blockDim.x;
for (int ind = tid; ind < n; ind += nthreads)
{
y[ind] += a * x[ind];
}
}
int main()
{
int ndevs;
cuda_safe_call(cudaGetDeviceCount(&ndevs));
cuda_safe_call(cudaSetDevice(0));
if (ndevs < 2)
{
fprintf(stderr, "Skipping test that needs at last 2 devices.\n");
return 0;
}
stream_ctx ctx;
const double alpha = 2.0;
const int n = 12;
double X[n], Y[n];
for (int ind = 0; ind < n; ind++)
{
X[ind] = 1.0 * ind;
Y[ind] = 2.0 * ind - 3.0;
}
auto handle_X = ctx.logical_data(X);
auto handle_Y = ctx.logical_data(Y);
/* Compute Y = Y + alpha X, but leave X on the host and access it with mapped memory */
ctx.task(exec_place::device(1), handle_X.read(), handle_Y.rw())->*[&](cudaStream_t stream, auto X, auto Y) {
axpy<<<16, 128, 0, stream>>>(n, alpha, X.data_handle(), Y.data_handle());
};
// Access Ask to use X, Y and Z on the host
ctx.task(exec_place::host(), handle_X.read(), handle_Y.read())->*[&](cudaStream_t stream, auto X, auto Y) {
cuda_safe_call(cudaStreamSynchronize(stream));
for (int ind = 0; ind < n; ind++)
{
// X unchanged
EXPECT(fabs(X(ind) - 1.0 * ind) < 0.00001);
// Y = Y + alpha X
EXPECT(fabs(Y(ind) - (-3.0 + ind * (2.0 + alpha))) < 0.00001);
}
};
ctx.finalize();
return 0;
}

53
cccl_upstream/cudax/test/stf/places/place_partition.cu Normal file
View File

@@ -0,0 +1,53 @@
//===----------------------------------------------------------------------===//
//
// 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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
#include <cuda/experimental/__places/place_partition.cuh>
#include <cuda/experimental/__stf/internal/stf_places_partition_into_stf.cuh>
#include <cuda/experimental/stf.cuh>
using namespace cuda::experimental::stf;
void print_partition(async_resources_handle& handle, exec_place place, place_partition_scope scope)
{
fprintf(stderr, "-----------\n");
fprintf(
stderr, "PARTITION %s (scope: %s):\n", place.to_string().c_str(), place_partition_scope_to_string(scope).c_str());
for (auto sub_place : place_partition(place, handle, scope))
{
fprintf(stderr, "[%s] subplace: %s\n", place.to_string().c_str(), sub_place.to_string().c_str());
}
fprintf(stderr, "-----------\n");
}
int main()
{
#if _CCCL_CTK_BELOW(12, 4)
fprintf(stderr, "Green contexts are not supported by this version of CUDA: skipping test.\n");
return 0;
#else // ^^^ _CCCL_CTK_BELOW(12, 4) ^^^ / vvv _CCCL_CTK_AT_LEAST(12, 4) vvv
async_resources_handle handle;
print_partition(handle, exec_place::all_devices(), place_partition_scope::cuda_device);
print_partition(handle, exec_place::all_devices(), place_partition_scope::cuda_stream);
print_partition(handle, exec_place::current_device(), place_partition_scope::cuda_stream);
print_partition(handle, exec_place::current_device(), place_partition_scope::green_context);
print_partition(handle, exec_place::current_device(), place_partition_scope::green_context);
print_partition(handle, exec_place::repeat(exec_place::current_device(), 4), place_partition_scope::green_context);
print_partition(handle, exec_place::current_device(), place_partition_scope::cuda_device);
print_partition(handle, exec_place::repeat(exec_place::current_device(), 4), place_partition_scope::cuda_stream);
#endif // ^^^ _CCCL_CTK_AT_LEAST(12, 4) ^^^
}

41
cccl_upstream/cudax/test/stf/places/recursion.cu Normal file
View File

@@ -0,0 +1,41 @@
//===----------------------------------------------------------------------===//
//
// 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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
#include <cuda/experimental/stf.cuh>
using namespace cuda::experimental::stf;
void rec_func(exec_place places)
{
if (places.size() == 1)
{
// places->print("SINGLE");
}
else
{
// places->print("REC");
for (int i = 0; i < 2; i++)
{
// Take every other places from the grid
auto half_places = partition_cyclic(places, dim4(2), pos4(i));
rec_func(half_places);
}
}
}
int main()
{
auto places = exec_place::all_devices();
// places->print("ALL");
rec_func(places);
return 0;
}

517
cccl_upstream/cudax/test/stf/places/replicated_data_place.cu Normal file
View File

@@ -0,0 +1,517 @@
//===----------------------------------------------------------------------===//
//
// 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 data_place::replicated — one copy of a logical data per grid
* member, read-only, fan-out on copy-in, per-place instance rebase in
* parallel_for. Covered on the stream backend, the graph backend, and
* inside a stackable conditional graph scope.
*/
#include <cuda/experimental/stf.cuh>
using namespace cuda::experimental::stf;
#if _CCCL_CTK_AT_LEAST(12, 4)
//! Replica instances are populated by copies issued INSIDE the nested
//! context of a stackable scope (the auto-push imports a single instance),
//! so the conditional body graph contains memcpy nodes. Some drivers reject
//! those at graph instantiation: probe for support so the stackable flavors
//! can waive instead of aborting.
bool conditional_body_memcpy_supported()
{
cudaGraph_t g;
cuda_safe_call(cudaGraphCreate(&g, 0));
cudaGraphConditionalHandle handle;
cuda_safe_call(cudaGraphConditionalHandleCreate(&handle, g, 1, cudaGraphCondAssignDefault));
cudaGraphNodeParams np{};
np.type = cudaGraphNodeTypeConditional;
np.conditional.handle = handle;
np.conditional.type = cudaGraphCondTypeWhile;
np.conditional.size = 1;
# if _CCCL_CTK_AT_LEAST(13, 0)
const cudaGraphNode_t cnode = cuda_try<cudaGraphAddNode>(g, nullptr, nullptr, 0, &np);
# else
const cudaGraphNode_t cnode = cuda_try<cudaGraphAddNode>(g, nullptr, 0, &np);
# endif
(void) cnode;
cudaGraph_t body = np.conditional.phGraph_out[0];
void* a = nullptr;
void* b = nullptr;
cuda_safe_call(cudaMalloc(&a, 8));
cuda_safe_call(cudaMalloc(&b, 8));
cudaGraphNode_t cp;
bool ok = (cudaGraphAddMemcpyNode1D(&cp, body, nullptr, 0, b, a, 8, cudaMemcpyDefault) == cudaSuccess);
if (ok)
{
cudaGraphExec_t e;
ok = (cudaGraphInstantiate(&e, g, 0) == cudaSuccess);
if (ok)
{
cuda_safe_call(cudaGraphExecDestroy(e));
}
}
cudaGetLastError(); // clear any probe failure
cuda_safe_call(cudaFree(a));
cuda_safe_call(cudaFree(b));
cuda_safe_call(cudaGraphDestroy(g));
return ok;
}
__global__ void probe_noop_kernel() {}
//! Conditional body graphs additionally require every kernel node to belong
//! to the SAME CUDA context ("all kernels ... must belong to the same CUDA
//! context"), and green contexts are distinct contexts. A stackable scope
//! running a green-grid task therefore mixes the primary context (the
//! condition-update kernel) with the grid's green contexts inside the body:
//! probe whether the driver accepts that at instantiation.
bool conditional_body_multi_context_supported(green_context_helper& gc)
{
cudaGraph_t g;
cuda_safe_call(cudaGraphCreate(&g, 0));
cudaGraphConditionalHandle handle;
cuda_safe_call(cudaGraphConditionalHandleCreate(&handle, g, 1, cudaGraphCondAssignDefault));
cudaGraphNodeParams np{};
np.type = cudaGraphNodeTypeConditional;
np.conditional.handle = handle;
np.conditional.type = cudaGraphCondTypeWhile;
np.conditional.size = 1;
# if _CCCL_CTK_AT_LEAST(13, 0)
const cudaGraphNode_t cnode = cuda_try<cudaGraphAddNode>(g, nullptr, nullptr, 0, &np);
# else
const cudaGraphNode_t cnode = cuda_try<cudaGraphAddNode>(g, nullptr, 0, &np);
# endif
(void) cnode;
cudaGraph_t body = np.conditional.phGraph_out[0];
cudaKernelNodeParams kp{};
kp.func = (void*) probe_noop_kernel;
kp.gridDim = dim3(1);
kp.blockDim = dim3(1);
kp.sharedMemBytes = 0;
kp.kernelParams = nullptr;
kp.extra = nullptr;
// one kernel in the current (primary) context, then one per green context
cudaGraphNode_t k;
cuda_safe_call(cudaGraphAddKernelNode(&k, body, nullptr, 0, &kp));
CUcontext prev;
cuda_safe_call(cuCtxGetCurrent(&prev));
for (size_t v = 0; v < 2; v++)
{
CUcontext gctx;
cuda_safe_call(cuCtxFromGreenCtx(&gctx, gc.get_view(v).g_ctx));
cuda_safe_call(cuCtxSetCurrent(gctx));
cuda_safe_call(cudaGraphAddKernelNode(&k, body, nullptr, 0, &kp));
}
cuda_safe_call(cuCtxSetCurrent(prev));
cudaGraphExec_t e;
bool ok = (cudaGraphInstantiate(&e, g, 0) == cudaSuccess);
if (ok)
{
cuda_safe_call(cudaGraphExecDestroy(e));
}
cudaGetLastError(); // clear any probe failure
cuda_safe_call(cudaGraphDestroy(g));
return ok;
}
#endif // _CCCL_CTK_AT_LEAST(12, 4)
int main()
{
#if _CCCL_CTK_BELOW(12, 4)
fprintf(stderr, "Waiving test: conditional nodes are only available since CUDA 12.4.\n");
return 0;
#else
const size_t n = 1 << 20;
const size_t nplaces = 2;
::std::vector<double> ref(n);
for (size_t i = 0; i < n; i++)
{
ref[i] = 0.25 * static_cast<double>(i % 1024) + 1.0;
}
const bool stackable_ok = conditional_body_memcpy_supported();
if (!stackable_ok)
{
printf("replicated data place: stackable flavors will be skipped (driver rejects memcpy nodes in conditional "
"body graphs)\n");
}
// ---- stream and graph backends
for (int use_graph = 0; use_graph < 2; use_graph++)
{
context ctx;
if (use_graph)
{
ctx = graph_ctx();
}
auto grid = exec_place::repeat(exec_place::current_device(), nplaces);
auto rep = data_place::replicated(grid);
auto lin = ctx.logical_data(&ref[0], {n});
auto lout = ctx.logical_data(shape_of<slice<double>>(n));
auto tok = ctx.token(); // a void_interface dep: the instances tuple is
// shorter than the deps tuple, which the
// replicated rebase must tolerate
// read at the replicated place: each shard must see the payload through
// its own replica, and results must match the reference everywhere
ctx.parallel_for(blocked_partition(), grid, lin.shape(), lin.read(rep), lout.write(), tok.write())
->*[] __device__(size_t i, auto in, auto out) {
out(i) = 2.0 * in(i);
};
ctx.host_launch(lout.read())->*[&](auto out) {
for (size_t i = 0; i < n; i++)
{
EXPECT(out(i) == 2.0 * ref[i]);
}
};
// read-only contract: any non-read access at a replicated place throws
bool thrown = false;
try
{
ctx.parallel_for(blocked_partition(), grid, lout.shape(), lout.rw(rep))->*[] __device__(size_t, auto) {};
}
catch (const ::std::invalid_argument&)
{
thrown = true;
}
EXPECT(thrown);
// deferred form on this backend: materialized against the launch's
// own execution place at acquire
auto lout_d = ctx.logical_data(shape_of<slice<double>>(n));
ctx.parallel_for(blocked_partition(), grid, lin.shape(), lin.read(data_place::replicated()), lout_d.write())
->*[] __device__(size_t i, auto in, auto out) {
out(i) = 6.0 * in(i);
};
ctx.host_launch(lout_d.read())->*[&](auto out) {
for (size_t i = 0; i < n; i++)
{
EXPECT(out(i) == 6.0 * ref[i]);
}
};
// merged-mode contract: declaring the SAME data as replicated-read and
// writable in one task merges the access modes past read; the combined
// dependency must be rejected at acquisition
bool thrown_merged = false;
try
{
ctx.parallel_for(blocked_partition(), grid, lout.shape(), lout.read(rep), lout.rw())
->*[] __device__(size_t, auto, auto) {};
}
catch (const ::std::invalid_argument&)
{
thrown_merged = true;
}
EXPECT(thrown_merged);
ctx.finalize();
printf("replicated data place: %s backend OK\n", use_graph ? "graph" : "stream");
}
// ---- green-context grid: member places are DISTINCT, so a replicated
// dep pins one real instance per member (the repeat grid above
// deduplicates to a single instance by place equality)
# if _CCCL_CTK_AT_LEAST(12, 4)
{
::std::optional<green_context_helper> gc_opt;
try
{
gc_opt.emplace(8, 0);
}
catch (...)
{
printf("replicated data place: green-context flavors skipped (unsupported hardware)\n");
}
if (gc_opt && gc_opt->get_count() >= 2)
{
auto& gc = *gc_opt;
::std::vector<exec_place> places;
places.push_back(exec_place::green_ctx(gc.get_view(0), true));
places.push_back(exec_place::green_ctx(gc.get_view(1), true));
auto ggrid = make_grid(mv(places));
auto grep = data_place::replicated(ggrid);
EXPECT(grep.instance_count() == 2);
EXPECT(grep.member(0) != grep.member(1));
context ctx;
for (int use_graph_g = 0; use_graph_g < 2; use_graph_g++)
{
if (use_graph_g)
{
ctx = graph_ctx();
}
auto lin = ctx.logical_data(&ref[0], {n});
auto lout = ctx.logical_data(shape_of<slice<double>>(n));
ctx.parallel_for(blocked_partition(), ggrid, lin.shape(), lin.read(grep), lout.write())
->*[] __device__(size_t i, auto in, auto out) {
out(i) = 3.0 * in(i);
};
ctx.host_launch(lout.read())->*[&](auto out) {
for (size_t i = 0; i < n; i++)
{
EXPECT(out(i) == 3.0 * ref[i]);
}
};
// deferred form: the grid is bound at acquire from the launch's
// execution place -- no grid repeated at the call site
auto lout2 = ctx.logical_data(shape_of<slice<double>>(n));
ctx.parallel_for(blocked_partition(), ggrid, lin.shape(), lin.read(data_place::replicated()), lout2.write())
->*[] __device__(size_t i, auto in, auto out) {
out(i) = 5.0 * in(i);
};
ctx.host_launch(lout2.read())->*[&](auto out) {
for (size_t i = 0; i < n; i++)
{
EXPECT(out(i) == 5.0 * ref[i]);
}
};
// scalar degenerate: on a single-place task the deferred form
// materializes to the place's affine data place
auto lout3 = ctx.logical_data(shape_of<slice<double>>(n));
ctx.parallel_for(lin.shape(), lin.read(data_place::replicated()), lout3.write())
->*[] __device__(size_t i, auto in, auto out) {
out(i) = 7.0 * in(i);
};
ctx.host_launch(lout3.read())->*[&](auto out) {
for (size_t i = 0; i < n; i++)
{
EXPECT(out(i) == 7.0 * ref[i]);
}
};
ctx.finalize();
printf("replicated data place: green-context grid (distinct member instances) %s OK\n",
use_graph_g ? "[graph]" : "[stream]");
printf("replicated data place: deferred form (grid-bound at acquire + scalar degenerate) %s OK\n",
use_graph_g ? "[graph]" : "[stream]");
}
// ---- axis-grouped replication on a (2, 2) grid: axis 0 = the two
// green domains (REPLICATED), axis 1 = two execution slots per domain
// (SHARED). One instance per domain, shared by its two slots.
{
::std::vector<exec_place> ps22;
ps22.push_back(exec_place::green_ctx(gc.get_view(0), true)); // (0,0)
ps22.push_back(exec_place::green_ctx(gc.get_view(1), true)); // (1,0)
ps22.push_back(exec_place::green_ctx(gc.get_view(0), true)); // (0,1)
ps22.push_back(exec_place::green_ctx(gc.get_view(1), true)); // (1,1)
auto grid22 = make_grid(mv(ps22), dim4(2, 2));
auto rep22 = data_place::replicated(grid22, replicate_over<0>);
// projection math: one instance per axis-0 coordinate
EXPECT(rep22.instance_count() == 2);
EXPECT(rep22.instance_of(0) == 0); // (0,0)
EXPECT(rep22.instance_of(1) == 1); // (1,0)
EXPECT(rep22.instance_of(2) == 0); // (0,1) shares domain 0's instance
EXPECT(rep22.instance_of(3) == 1); // (1,1) shares domain 1's instance
EXPECT(rep22.member(0) != rep22.member(1));
context gctx;
// the grouped resolution is backend-agnostic: same launch on the
// stream and graph backends
for (int use_graph22 = 0; use_graph22 < 2; use_graph22++)
{
if (use_graph22)
{
gctx = graph_ctx();
}
auto lgin = gctx.logical_data(&ref[0], {n});
auto lgout = gctx.logical_data(shape_of<slice<double>>(n));
gctx.parallel_for(blocked_partition(), grid22, lgin.shape(), lgin.read(rep22), lgout.write())
->*[] __device__(size_t i, auto in, auto out) {
out(i) = 4.0 * in(i);
};
gctx.host_launch(lgout.read())->*[&](auto out) {
for (size_t i = 0; i < n; i++)
{
EXPECT(out(i) == 4.0 * ref[i]);
}
};
gctx.finalize();
}
// and inside a stackable conditional scope: the auto-push imports
// the data once, the nested acquire pins one instance per axis-0
// group
const bool multi_ctx_ok = conditional_body_multi_context_supported(gc);
if (!multi_ctx_ok)
{
printf("replicated data place: green stackable flavors skipped (driver requires a single CUDA context in "
"conditional body graphs)\n");
}
if (stackable_ok && multi_ctx_ok)
{
stackable_ctx sctx;
auto lsin = sctx.logical_data(shape_of<slice<double>>(n)).set_symbol("g22in");
auto lsacc = sctx.logical_data(shape_of<slice<double>>(n)).set_symbol("g22acc");
sctx.parallel_for(lsin.shape(), lsin.write(), lsacc.write())->*[] __device__(size_t i, auto in, auto acc) {
in(i) = static_cast<double>(i % 32);
acc(i) = 0.0;
};
const size_t iters22 = 6;
{
auto rg = sctx.repeat_graph_scope(iters22);
sctx.parallel_for(blocked_partition(), grid22, lsin.shape(), lsin.read(rep22), lsacc.rw())
->*[] __device__(size_t i, auto in, auto acc) {
acc(i) += in(i);
};
}
sctx.host_launch(lsacc.read())->*[&](auto acc) {
for (size_t i = 0; i < n; i++)
{
EXPECT(acc(i) == static_cast<double>(iters22) * static_cast<double>(i % 32));
}
};
sctx.finalize();
}
// shared axes require co-located fibers: a (2, 2) arrangement whose
// axis-1 fiber crosses domains must be rejected at construction
::std::vector<exec_place> bad;
bad.push_back(exec_place::green_ctx(gc.get_view(0), true)); // (0,0)
bad.push_back(exec_place::green_ctx(gc.get_view(1), true)); // (1,0)
bad.push_back(exec_place::green_ctx(gc.get_view(1), true)); // (0,1) != (0,0)
bad.push_back(exec_place::green_ctx(gc.get_view(0), true)); // (1,1) != (1,0)
auto bad_grid = make_grid(mv(bad), dim4(2, 2));
bool thrown22 = false;
try
{
auto r = data_place::replicated(bad_grid, replicate_over<0>);
}
catch (const ::std::invalid_argument&)
{
thrown22 = true;
}
EXPECT(thrown22);
printf("replicated data place: axis-grouped replication on a (2,2) grid OK (stream/graph/stackable)\n");
}
}
else
{
printf("replicated data place: green-context flavor skipped (single group)\n");
}
}
# endif // _CCCL_CTK_AT_LEAST(12, 4)
// ---- stackable conditional graph scope: replicas are ordinary
// stream-ordered instances, so the auto-push (freeze + get within the
// nested context) and pop-time lifetime need no special handling
if (stackable_ok)
{
stackable_ctx ctx;
auto grid = exec_place::repeat(exec_place::current_device(), nplaces);
auto rep = data_place::replicated(grid);
auto lin = ctx.logical_data(shape_of<slice<double>>(n)).set_symbol("in");
auto lacc = ctx.logical_data(shape_of<slice<double>>(n)).set_symbol("acc");
ctx.parallel_for(lin.shape(), lin.write(), lacc.write())->*[] __device__(size_t i, auto in, auto acc) {
in(i) = static_cast<double>(i % 128);
acc(i) = 0.0;
};
const size_t iters = 10;
{
auto rg = ctx.repeat_graph_scope(iters);
ctx.parallel_for(blocked_partition(), grid, lin.shape(), lin.read(rep), lacc.rw())
->*[] __device__(size_t i, auto in, auto acc) {
acc(i) += in(i);
};
}
// the DEFERRED form through the stackable auto-push: the push imports
// the data at the context's default place; the nested task's acquire
// then materializes replicated() against its own execution place
{
auto rg = ctx.repeat_graph_scope(iters);
ctx.parallel_for(blocked_partition(), grid, lin.shape(), lin.read(data_place::replicated()), lacc.rw())
->*[] __device__(size_t i, auto in, auto acc) {
acc(i) += in(i);
};
}
ctx.host_launch(lacc.read())->*[&](auto acc) {
for (size_t i = 0; i < n; i++)
{
EXPECT(acc(i) == 2.0 * static_cast<double>(iters) * static_cast<double>(i % 128));
}
};
ctx.finalize();
printf("replicated data place: stackable conditional scope (concrete + deferred deps) OK\n");
}
// ---- stackable + DISTINCT member places: true multi-replica through the
// auto-push and a conditional scope (the repeat-grid section above only
// exercises the equal-places dedup degenerate)
# if _CCCL_CTK_AT_LEAST(12, 4)
{
::std::optional<green_context_helper> gc_opt;
try
{
gc_opt.emplace(8, 0);
}
catch (...)
{}
if (stackable_ok && gc_opt && gc_opt->get_count() >= 2 && conditional_body_multi_context_supported(*gc_opt))
{
auto& gc = *gc_opt;
stackable_ctx ctx;
::std::vector<exec_place> places;
places.push_back(exec_place::green_ctx(gc.get_view(0), true));
places.push_back(exec_place::green_ctx(gc.get_view(1), true));
auto ggrid = make_grid(mv(places));
auto lin = ctx.logical_data(shape_of<slice<double>>(n)).set_symbol("gin");
auto lacc = ctx.logical_data(shape_of<slice<double>>(n)).set_symbol("gacc");
ctx.parallel_for(lin.shape(), lin.write(), lacc.write())->*[] __device__(size_t i, auto in, auto acc) {
in(i) = static_cast<double>(i % 64);
acc(i) = 0.0;
};
const size_t iters = 8;
{
auto rg = ctx.repeat_graph_scope(iters);
ctx.parallel_for(blocked_partition(), ggrid, lin.shape(), lin.read(data_place::replicated()), lacc.rw())
->*[] __device__(size_t i, auto in, auto acc) {
acc(i) += in(i);
};
}
ctx.host_launch(lacc.read())->*[&](auto acc) {
for (size_t i = 0; i < n; i++)
{
EXPECT(acc(i) == static_cast<double>(iters) * static_cast<double>(i % 64));
}
};
ctx.finalize();
printf("replicated data place: stackable + distinct member replicas OK\n");
}
else
{
printf("replicated data place: stackable green flavor skipped\n");
}
}
# endif // _CCCL_CTK_AT_LEAST(12, 4)
printf("replicated_data_place: all checks passed\n");
return 0;
#endif
}

168
cccl_upstream/cudax/test/stf/places/replicated_frozen_import.cu Normal file
View File

@@ -0,0 +1,168 @@
//===----------------------------------------------------------------------===//
//
// 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 Frozen-import member walk: pushing a logical data at a concrete
* replicated place imports EVERY member instance into the nested context
* (population happens at the parent level), so a read at the replicated
* place resolves in the nested context without issuing any copy -- in
* particular no memcpy node lands in a conditional body graph.
*/
#include <cuda/experimental/stf.cuh>
using namespace cuda::experimental::stf;
int main()
{
#if _CCCL_CTK_BELOW(12, 4)
fprintf(stderr, "Waiving test: conditional nodes are only available since CUDA 12.4.\n");
return 0;
#else
const size_t n = 1 << 20;
const size_t nplaces = 2;
// ---- explicit push at the replicated place, then a conditional scope:
// the walk runs at push time (repeat grid: equal members dedup to one
// instance)
{
stackable_ctx ctx;
auto grid = exec_place::repeat(exec_place::current_device(), nplaces);
auto rep = data_place::replicated(grid);
auto lin = ctx.logical_data(shape_of<slice<double>>(n)).set_symbol("fin");
auto lacc = ctx.logical_data(shape_of<slice<double>>(n)).set_symbol("facc");
ctx.parallel_for(lin.shape(), lin.write(), lacc.write())->*[] __device__(size_t i, auto in, auto acc) {
in(i) = static_cast<double>(i % 128);
acc(i) = 0.0;
};
const size_t iters = 10;
{
auto rg = ctx.repeat_graph_scope(iters);
// push INSIDE the scope (at root level a push is a no-op): the member
// walk imports every replica; the accumulator is imported directly at
// the grid's composite place so no instance needs a transfer inside
// the scope
lin.push(access_mode::read, rep);
lacc.push(access_mode::rw, data_place::composite(blocked_partition(), grid));
ctx.parallel_for(blocked_partition(), grid, lin.shape(), lin.read(rep), lacc.rw())
->*[] __device__(size_t i, auto in, auto acc) {
acc(i) += in(i);
};
}
ctx.host_launch(lacc.read())->*[&](auto acc) {
for (size_t i = 0; i < n; i++)
{
EXPECT(acc(i) == static_cast<double>(iters) * static_cast<double>(i % 128));
}
};
ctx.finalize();
printf("frozen import: explicit push at replicated place OK\n");
}
// ---- read-only data: the auto-push sees the dependency's replicated
// place and walks the members without an explicit push
{
stackable_ctx ctx;
auto grid = exec_place::repeat(exec_place::current_device(), nplaces);
auto rep = data_place::replicated(grid);
auto lro = ctx.logical_data(shape_of<slice<double>>(n)).set_symbol("fro");
auto lacc = ctx.logical_data(shape_of<slice<double>>(n)).set_symbol("froacc");
ctx.parallel_for(lro.shape(), lro.write(), lacc.write())->*[] __device__(size_t i, auto in, auto acc) {
in(i) = static_cast<double>(i % 64);
acc(i) = 0.0;
};
lro.set_read_only();
const size_t iters = 6;
{
auto rg = ctx.repeat_graph_scope(iters);
lacc.push(access_mode::rw, data_place::composite(blocked_partition(), grid));
ctx.parallel_for(blocked_partition(), grid, lro.shape(), lro.read(rep), lacc.rw())
->*[] __device__(size_t i, auto in, auto acc) {
acc(i) += in(i);
};
}
ctx.host_launch(lacc.read())->*[&](auto acc) {
for (size_t i = 0; i < n; i++)
{
EXPECT(acc(i) == static_cast<double>(iters) * static_cast<double>(i % 64));
}
};
ctx.finalize();
printf("frozen import: auto-push walk for read-only data OK\n");
}
// ---- distinct member places (green contexts): the walk adopts one REAL
// instance per member. A plain (non-conditional) graph scope keeps this
// flavor clear of the driver rule that conditional bodies may not mix
// CUDA contexts.
# if _CCCL_CTK_AT_LEAST(12, 4)
{
::std::optional<green_context_helper> gc_opt;
try
{
gc_opt.emplace(8, 0);
}
catch (...)
{}
if (gc_opt && gc_opt->get_count() >= 2)
{
auto& gc = *gc_opt;
stackable_ctx ctx;
::std::vector<exec_place> places;
places.push_back(exec_place::green_ctx(gc.get_view(0), true));
places.push_back(exec_place::green_ctx(gc.get_view(1), true));
auto ggrid = make_grid(mv(places));
auto grep = data_place::replicated(ggrid);
EXPECT(grep.member(0) != grep.member(1));
auto lin = ctx.logical_data(shape_of<slice<double>>(n)).set_symbol("gfin");
auto lout = ctx.logical_data(shape_of<slice<double>>(n)).set_symbol("gfout");
ctx.parallel_for(lin.shape(), lin.write(), lout.write())->*[] __device__(size_t i, auto in, auto out) {
in(i) = static_cast<double>(i % 32);
out(i) = 0.0;
};
{
auto gs = ctx.graph_scope();
lin.push(access_mode::read, grep);
lout.push(access_mode::rw, data_place::composite(blocked_partition(), ggrid));
ctx.parallel_for(blocked_partition(), ggrid, lin.shape(), lin.read(grep), lout.rw())
->*[] __device__(size_t i, auto in, auto out) {
out(i) = 9.0 * in(i);
};
}
ctx.host_launch(lout.read())->*[&](auto out) {
for (size_t i = 0; i < n; i++)
{
EXPECT(out(i) == 9.0 * static_cast<double>(i % 32));
}
};
ctx.finalize();
printf("frozen import: distinct member instances adopted (green grid) OK\n");
}
else
{
printf("frozen import: green flavor skipped\n");
}
}
# endif // _CCCL_CTK_AT_LEAST(12, 4)
printf("replicated_frozen_import: all checks passed\n");
return 0;
#endif
}