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

1746
cccl_upstream/cudax/include/cuda/experimental/__places/cute_partition.cuh
View File

File diff suppressed because it is too large Load Diff

414
cccl_upstream/cudax/include/cuda/experimental/__places/data_place_impl.cuh
View File

@@ -1,414 +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 Concrete implementations of data_place_interface
*
* This file contains implementations for standard data place types:
* host, managed, device, invalid, affine, and device_auto.
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/experimental/__places/data_place_interface.cuh>
#include <cuda/experimental/__stf/utility/cuda_safe_call.cuh>
#include <cuda/experimental/__stf/utility/scope_guard.cuh>
namespace cuda::experimental::places
{
using ::cuda::experimental::stf::cuda_try;
/**
* @brief Implementation for the invalid data place
*/
class data_place_invalid final : public data_place_interface
{
public:
bool is_resolved() const override
{
return false;
}
int get_device_ordinal() const override
{
return data_place_interface::invalid;
}
::std::string to_string() const override
{
return "invalid";
}
size_t hash() const override
{
return ::std::hash<int>()(data_place_interface::invalid);
}
int cmp(const data_place_interface& other) const override
{
if (typeid(*this) != typeid(other))
{
return typeid(*this).before(typeid(other)) ? -1 : 1;
}
return 0;
}
void* allocate(::std::ptrdiff_t, cudaStream_t) const override
{
throw ::std::logic_error("Cannot allocate from invalid data_place");
}
void deallocate(void*, size_t, cudaStream_t) const override
{
throw ::std::logic_error("Cannot deallocate from invalid data_place");
}
bool allocation_is_stream_ordered() const override
{
return false;
}
};
/**
* @brief Implementation for the host (pinned memory) data place
*/
class data_place_host final : public data_place_interface
{
public:
bool is_resolved() const override
{
return true;
}
int get_device_ordinal() const override
{
return data_place_interface::host;
}
::std::string to_string() const override
{
return "host";
}
size_t hash() const override
{
return ::std::hash<int>()(data_place_interface::host);
}
int cmp(const data_place_interface& other) const override
{
if (typeid(*this) != typeid(other))
{
return typeid(*this).before(typeid(other)) ? -1 : 1;
}
return 0;
}
void* allocate(::std::ptrdiff_t size, cudaStream_t) const override
{
void* result = nullptr;
cuda_try(cudaMallocHost(&result, static_cast<size_t>(size)));
return result;
}
void deallocate(void* ptr, size_t, cudaStream_t) const override
{
cuda_try(cudaFreeHost(ptr));
}
bool allocation_is_stream_ordered() const override
{
return false;
}
CUresult mem_create(CUmemGenericAllocationHandle* handle, size_t size) const override
{
#if _CCCL_CTK_AT_LEAST(12, 2)
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_HOST;
prop.location.id = 0;
return cuMemCreate(handle, size, &prop, 0);
#else
(void) handle;
(void) size;
return CUDA_ERROR_NOT_SUPPORTED;
#endif
}
};
/**
* @brief Implementation for managed memory data place
*/
class data_place_managed final : public data_place_interface
{
public:
bool is_resolved() const override
{
return true;
}
int get_device_ordinal() const override
{
return data_place_interface::managed;
}
::std::string to_string() const override
{
return "managed";
}
size_t hash() const override
{
return ::std::hash<int>()(data_place_interface::managed);
}
int cmp(const data_place_interface& other) const override
{
if (typeid(*this) != typeid(other))
{
return typeid(*this).before(typeid(other)) ? -1 : 1;
}
return 0;
}
void* allocate(::std::ptrdiff_t size, cudaStream_t) const override
{
void* result = nullptr;
cuda_try(cudaMallocManaged(&result, static_cast<size_t>(size)));
return result;
}
void deallocate(void* ptr, size_t, cudaStream_t) const override
{
cuda_try(cudaFree(ptr));
}
bool allocation_is_stream_ordered() const override
{
return false;
}
};
/**
* @brief Implementation for a specific CUDA device data place
*/
class data_place_device final : public data_place_interface
{
public:
bool is_resolved() const override
{
return true;
}
explicit data_place_device(int device_id)
: device_id_(device_id)
{
_CCCL_ASSERT(device_id >= 0, "Device ID must be non-negative");
}
int get_device_ordinal() const override
{
return device_id_;
}
::std::string to_string() const override
{
return "dev" + ::std::to_string(device_id_);
}
size_t hash() const override
{
return ::std::hash<int>()(device_id_);
}
int cmp(const data_place_interface& other) const override
{
if (typeid(*this) != typeid(other))
{
return typeid(*this).before(typeid(other)) ? -1 : 1;
}
return (device_id_ > static_cast<const data_place_device&>(other).device_id_)
- (device_id_ < static_cast<const data_place_device&>(other).device_id_);
}
void* allocate(::std::ptrdiff_t size, cudaStream_t stream) const override
{
void* result = nullptr;
const int prev_dev = cuda_try<cudaGetDevice>();
if (prev_dev != device_id_)
{
cuda_try(cudaSetDevice(device_id_));
}
SCOPE(exit)
{
if (prev_dev != device_id_)
{
cuda_try(cudaSetDevice(prev_dev));
}
};
cuda_try(cudaMallocAsync(&result, static_cast<size_t>(size), stream));
return result;
}
void deallocate(void* ptr, size_t, cudaStream_t stream) const override
{
const int prev_dev = cuda_try<cudaGetDevice>();
if (prev_dev != device_id_)
{
cuda_try(cudaSetDevice(device_id_));
}
SCOPE(exit)
{
if (prev_dev != device_id_)
{
cuda_try(cudaSetDevice(prev_dev));
}
};
cuda_try(cudaFreeAsync(ptr, stream));
}
bool allocation_is_stream_ordered() const override
{
return true;
}
CUresult mem_create(CUmemGenericAllocationHandle* handle, size_t size) const override
{
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.location.id = device_id_;
return cuMemCreate(handle, size, &prop, 0);
}
private:
int device_id_;
};
/**
* @brief Implementation for the affine data place (uses exec_place's affine data place)
*/
class data_place_affine final : public data_place_interface
{
public:
bool is_resolved() const override
{
return false;
}
int get_device_ordinal() const override
{
return data_place_interface::affine;
}
::std::string to_string() const override
{
return "affine";
}
size_t hash() const override
{
return ::std::hash<int>()(data_place_interface::affine);
}
int cmp(const data_place_interface& other) const override
{
if (typeid(*this) != typeid(other))
{
return typeid(*this).before(typeid(other)) ? -1 : 1;
}
return 0;
}
void* allocate(::std::ptrdiff_t, cudaStream_t) const override
{
throw ::std::logic_error("Cannot allocate from affine data_place directly");
}
void deallocate(void*, size_t, cudaStream_t) const override
{
throw ::std::logic_error("Cannot deallocate from affine data_place directly");
}
bool allocation_is_stream_ordered() const override
{
return false;
}
};
/**
* @brief Implementation for device_auto data place (auto-select device)
*/
class data_place_device_auto final : public data_place_interface
{
public:
bool is_resolved() const override
{
return false;
}
int get_device_ordinal() const override
{
return data_place_interface::device_auto;
}
::std::string to_string() const override
{
return "auto";
}
size_t hash() const override
{
return ::std::hash<int>()(data_place_interface::device_auto);
}
int cmp(const data_place_interface& other) const override
{
if (typeid(*this) != typeid(other))
{
return typeid(*this).before(typeid(other)) ? -1 : 1;
}
return 0;
}
void* allocate(::std::ptrdiff_t, cudaStream_t) const override
{
throw ::std::logic_error("Cannot allocate from device_auto data_place directly");
}
void deallocate(void*, size_t, cudaStream_t) const override
{
throw ::std::logic_error("Cannot deallocate from device_auto data_place directly");
}
bool allocation_is_stream_ordered() const override
{
return true;
}
};
} // end namespace cuda::experimental::places

227
cccl_upstream/cudax/include/cuda/experimental/__places/data_place_interface.cuh
View File

@@ -1,227 +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 Abstract interface for data_place implementations
*
* This interface defines the contract that all data_place implementations must satisfy.
* It enables a clean polymorphic design where host, managed, device, composite, and
* custom places (e.g. green contexts) all implement a common interface.
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/std/__cccl/assert.h>
#include <cuda/experimental/__stf/utility/dimensions.cuh>
#include <cstddef>
#include <functional>
#include <limits>
#include <memory>
#include <stdexcept>
#include <string>
#include <typeinfo>
#include <cuda.h>
#include <cuda_runtime.h>
namespace cuda::experimental::places
{
using ::cuda::experimental::stf::dim4;
using ::cuda::experimental::stf::pos4;
// Forward declarations
class exec_place;
//! Function type for computing executor placement from data coordinates.
//! Uses an out-pointer convention so the signature is trivially representable
//! in FFI frameworks (ctypes, cffi, Rust) that cannot return C structs.
using partition_fn_t = void (*)(pos4* result, pos4 data_coords, dim4 data_dims, dim4 grid_dims);
/**
* @brief Abstract interface for data_place implementations
*
* All data_place types (host, managed, device, composite, future places) implement
* this interface. The data_place class holds a shared_ptr to this interface
* and delegates all operations to it.
*/
class data_place_interface
{
public:
virtual ~data_place_interface() = default;
/**
* @brief Special device ordinal values for non-device places
*
* Returned by get_device_ordinal() for places that don't correspond
* to a specific CUDA device.
*/
enum ord : int
{
invalid = ::std::numeric_limits<int>::min(),
composite = -5,
device_auto = -4,
affine = -3,
managed = -2,
host = -1,
};
// === Core properties ===
/**
* @brief Whether this place is fully resolved and ready for allocation
*
* Returns true for places that represent a concrete memory target:
* host, managed, device(N), composite, green_ctx, etc.
* Returns false for abstract/deferred places that need further
* resolution: invalid, affine, device_auto.
*/
virtual bool is_resolved() const = 0;
/**
* @brief Get the device ordinal for this place
*
* Returns:
* - >= 0 for specific CUDA devices
* - data_place_ordinals::host (-1) for host
* - data_place_ordinals::managed (-2) for managed
* - data_place_ordinals::affine (-3) for affine
* - data_place_ordinals::device_auto (-4) for device_auto
* - data_place_ordinals::composite (-5) for composite
* - data_place_ordinals::invalid for invalid
*/
virtual int get_device_ordinal() const = 0;
/**
* @brief Get a string representation of this place
*/
virtual ::std::string to_string() const = 0;
/**
* @brief Compute a hash value for this place
*/
virtual size_t hash() const = 0;
/**
* @brief Three-way comparison with another place
*
* @return -1 if *this < other, 0 if *this == other, 1 if *this > other
*/
virtual int cmp(const data_place_interface& other) const = 0;
// === Memory allocation ===
/**
* @brief Allocate memory at this place
*
* @param size Size of the allocation in bytes
* @param stream CUDA stream for stream-ordered allocations
* @return Pointer to allocated memory
* @throws std::runtime_error if allocation is not supported for this place type
*/
virtual void* allocate(::std::ptrdiff_t size, cudaStream_t stream) const = 0;
/**
* @brief Allocate memory at this place for a tensor with the given extents
*
* The default implementation ignores the tensor geometry and forwards to the
* byte-count allocate(); places whose physical placement depends on the
* geometry (composite places, whose partitioner maps element coordinates to
* places) override it with the real implementation.
*
* Extents follow the dimension-0-fastest linearization convention of
* dim4::get_index() (the STF slice convention). Row-major callers should
* present reversed extents (and a coordinate-reversing partitioner).
*
* @param data_dims Extents of the tensor
* @param elemsize Size of one element in bytes
* @param stream CUDA stream for stream-ordered allocations
* @return Pointer to allocated memory
*/
virtual void* allocate_nd(dim4 data_dims, size_t elemsize, cudaStream_t stream) const
{
return allocate(static_cast<::std::ptrdiff_t>(data_dims.size() * elemsize), stream);
}
/**
* @brief Deallocate memory at this place
*
* @param ptr Pointer to memory to deallocate
* @param size Size of the allocation
* @param stream CUDA stream for stream-ordered deallocations
*/
virtual void deallocate(void* ptr, size_t size, cudaStream_t stream) const = 0;
/**
* @brief Returns true if allocation/deallocation is stream-ordered
*/
virtual bool allocation_is_stream_ordered() const = 0;
/**
* @brief Create a physical memory allocation for this place (VMM API)
*
* Default implementation returns CUDA_ERROR_NOT_SUPPORTED.
* Subclasses that support VMM should override this.
*
* @param handle Output parameter for the allocation handle
* @param size Size of the allocation in bytes
* @return CUresult indicating success or failure
*/
virtual CUresult mem_create(CUmemGenericAllocationHandle* handle, size_t size) const
{
return CUDA_ERROR_NOT_SUPPORTED;
}
/**
* @brief Get the implementation for the affine exec_place (for custom place types)
*
* Custom data_place implementations (e.g. green contexts) override this to
* provide their own affine exec_place. Returns nullptr by default, which
* causes data_place::affine_exec_place() to fall through to the error path.
* The returned shared_ptr should be castable to shared_ptr<exec_place::impl>.
*/
virtual ::std::shared_ptr<void> get_affine_exec_impl() const
{
return nullptr;
}
// === Composite-specific (throw by default) ===
/**
* @brief Whether this place is a composite place (data distributed over a
* grid of places by a partitioner)
*/
virtual bool is_composite() const
{
return false;
}
/**
* @brief Get the partitioner function for composite places
* @throws std::logic_error if not a composite place
*/
virtual const partition_fn_t& get_partitioner() const
{
throw ::std::logic_error("get_partitioner() called on non-composite data_place");
}
};
} // end namespace cuda::experimental::places

295
cccl_upstream/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh
View File

@@ -1,295 +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 Execution place wrapping an externally-owned CUDA driver context
*
* This makes it possible to use a CUcontext created outside CUDASTF (e.g. a
* green context created through cuda.core in Python, or any other library)
* as an execution place. The place is non-owning: the caller must keep the
* context alive while the place (and the streams lazily created in its pool)
* is in use.
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/__runtime/ensure_current_context.h>
#include <cuda/experimental/__places/places.cuh>
#include <cuda/experimental/__stf/utility/hash.cuh>
#include <functional>
namespace cuda::experimental::places
{
/**
* @brief Implementation for execution places backed by an externally-owned CUcontext
*
* The context is used as-is: `activate()` saves the current context and makes
* this one current, `deactivate()` restores the saved context. Streams are
* created lazily in the place's own pool while the context is current, so they
* inherit whatever resources the context carries (e.g. the SM partition of a
* green context converted with `cuCtxFromGreenCtx`).
*
* Identity (`hash`/`cmp`) is keyed on the CUcontext handle, which uniquely
* identifies the underlying resources: `cuCtxFromGreenCtx` returns the same
* CUcontext for a given CUgreenCtx on every call.
*/
class exec_place_cuda_ctx_impl : public exec_place::impl
{
public:
/**
* @brief Construct an execution place from an externally-owned CUDA context
*
* @param ctx The CUDA driver context. Non-owning: the caller keeps it alive.
* @param devid The device ordinal the context belongs to, or -1 to derive it
* from the context (via cuCtxGetDevice).
* @param pool_size Number of streams in the place's stream pool.
*/
exec_place_cuda_ctx_impl(CUcontext ctx, int devid = -1, size_t pool_size = exec_place::impl::pool_size)
: exec_place_cuda_ctx_impl(ctx, resolve_devid(ctx, devid), stream_pool(pool_size))
{}
/**
* @brief Full-control constructor, also used by the green-context place
*
* @param ctx The CUDA driver context (already resolved, e.g. from cuCtxFromGreenCtx)
* @param devid The device ordinal (must be valid)
* @param pool A stream pool to use for this place (shared handle)
* @param affine The affine data place for this place
*/
exec_place_cuda_ctx_impl(CUcontext ctx, int devid, stream_pool pool, data_place affine)
: exec_place::impl(mv(affine))
, devid_(devid)
, driver_context_(ctx)
, pool_(mv(pool))
{
_CCCL_ASSERT(ctx != nullptr, "cuda_ctx exec_place requires a valid CUcontext");
_CCCL_ASSERT(devid_ >= 0, "cuda_ctx exec_place requires a valid device ordinal");
}
::std::shared_ptr<exec_place::impl> get_place(size_t idx) override
{
_CCCL_ASSERT(idx == 0, "Index out of bounds for cuda_ctx exec_place");
return shared_from_this();
}
exec_place activate(size_t idx) const override
{
_CCCL_ASSERT(idx == 0, "Index out of bounds for cuda_ctx exec_place");
// Save the current context wrapped as a place so deactivate can restore it
CUcontext current_ctx = cuda_try<cuCtxGetCurrent>();
exec_place result = exec_place(::std::make_shared<exec_place_cuda_ctx_impl>(saved_tag{}, current_ctx));
cuda_try<cuCtxSetCurrent>(driver_context_);
return result;
}
void deactivate(const exec_place& prev, size_t idx = 0) const override
{
_CCCL_ASSERT(idx == 0, "Index out of bounds for cuda_ctx exec_place");
auto prev_impl = ::std::static_pointer_cast<exec_place_cuda_ctx_impl>(prev.get_impl());
CUcontext saved_ctx = prev_impl->driver_context_;
cuda_try<cuCtxSetCurrent>(saved_ctx);
}
bool is_device() const override
{
return true;
}
::std::string to_string() const override
{
return "cuda_ctx(ctx=" + ::std::to_string(reinterpret_cast<::std::uintptr_t>(driver_context_))
+ " dev=" + ::std::to_string(devid_) + ")";
}
stream_pool& get_stream_pool(bool, exec_place_resources&, const exec_place&) const override
{
// This place carries its own pool and bypasses the registry. The user is
// responsible for keeping the underlying CUcontext alive while the pool
// is in use.
return pool_;
}
int cmp(const exec_place::impl& rhs) const override
{
if (typeid(*this) != typeid(rhs))
{
return typeid(*this).before(typeid(rhs)) ? -1 : 1;
}
const auto& other = static_cast<const exec_place_cuda_ctx_impl&>(rhs);
return ::std::less<CUcontext>{}(other.driver_context_, driver_context_)
- ::std::less<CUcontext>{}(driver_context_, other.driver_context_);
}
size_t hash() const override
{
return ::std::hash<CUcontext>()(driver_context_);
}
protected:
// Tag type for the internal "saved context" wrapper used by activate/deactivate.
// The type itself is protected so only this class (and derived classes) can
// name it; the constructor below must be public for make_shared.
struct saved_tag
{};
public:
// Wrap an existing context with no pool: only used to carry the saved
// context through the activate/deactivate round trip.
exec_place_cuda_ctx_impl(saved_tag, CUcontext saved_context)
: driver_context_(saved_context)
{}
protected:
exec_place_cuda_ctx_impl(CUcontext ctx, int devid, stream_pool pool)
: exec_place_cuda_ctx_impl(ctx, devid, mv(pool), data_place::device(devid))
{}
static int resolve_devid(CUcontext ctx, int devid)
{
::cuda::__ensure_current_context guard{ctx};
const int context_devid = static_cast<int>(cuda_try<cuCtxGetDevice>());
if (devid >= 0 && devid != context_devid)
{
throw ::std::invalid_argument("CUcontext device ordinal does not match devid");
}
return context_devid;
}
int devid_ = -1;
CUcontext driver_context_ = {};
mutable stream_pool pool_;
};
inline exec_place exec_place::cuda_context(CUcontext ctx, int devid, size_t pool_size)
{
if (ctx == nullptr)
{
// Reject eagerly: with an explicit devid a null context would construct a
// place that only fails (or silently unbinds the current context) at
// activation time.
throw ::std::invalid_argument("exec_place::cuda_context requires a valid CUcontext");
}
return exec_place(::std::make_shared<exec_place_cuda_ctx_impl>(ctx, devid, pool_size));
}
#ifdef UNITTESTED_FILE
namespace
{
//! RAII holder for the primary context of device 0, used by the unittests below.
struct primary_ctx_guard
{
primary_ctx_guard()
{
cuda_try<cuInit>(0);
dev = cuda_try<cuDeviceGet>(0);
ctx = cuda_try<cuDevicePrimaryCtxRetain>(dev);
}
~primary_ctx_guard()
{
cuda_try(cuDevicePrimaryCtxRelease(dev));
}
CUdevice dev = -1;
CUcontext ctx = nullptr;
};
} // namespace
UNITTEST("cuda_context exec_place equality")
{
primary_ctx_guard guard;
auto p0a = exec_place::cuda_context(guard.ctx, 0);
auto p0b = exec_place::cuda_context(guard.ctx, 0);
// Same context should be equal
EXPECT(p0a == p0b);
EXPECT(!(p0a != p0b));
// A cuda_context place is not a regular device place
auto dev0 = exec_place::device(0);
EXPECT(p0a != dev0);
EXPECT(!(p0a == dev0));
};
UNITTEST("cuda_context exec_place derives the device ordinal")
{
primary_ctx_guard guard;
// devid intentionally omitted: derived from the context via cuCtxGetDevice
auto p = exec_place::cuda_context(guard.ctx);
EXPECT(p.is_device());
EXPECT(p.affine_data_place() == data_place::device(0));
EXPECT(p == exec_place::cuda_context(guard.ctx, 0));
};
UNITTEST("cuda_context exec_place rejects a null context")
{
bool thrown = false;
try
{
auto p = exec_place::cuda_context(nullptr, 0);
}
catch (const ::std::invalid_argument&)
{
thrown = true;
}
EXPECT(thrown);
};
UNITTEST("cuda_context exec_place rejects a mismatched device ordinal")
{
primary_ctx_guard guard;
bool thrown = false;
try
{
auto p = exec_place::cuda_context(guard.ctx, 1);
}
catch (const ::std::invalid_argument&)
{
thrown = true;
}
EXPECT(thrown);
};
UNITTEST("cuda_context exec_place activate/deactivate round trip")
{
primary_ctx_guard guard;
auto p = exec_place::cuda_context(guard.ctx, 0);
CUcontext before = cuda_try<cuCtxGetCurrent>();
{
exec_place_scope scope(p);
EXPECT(cuda_try<cuCtxGetCurrent>() == guard.ctx);
}
EXPECT(cuda_try<cuCtxGetCurrent>() == before);
};
#endif // UNITTESTED_FILE
} // end namespace cuda::experimental::places

110
cccl_upstream/cudax/include/cuda/experimental/__places/exec/cuda_stream.cuh
View File

@@ -1,110 +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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief CUDA stream execution place implementation
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/experimental/__places/places.cuh>
namespace cuda::experimental::places
{
/**
* @brief Implementation for CUDA stream execution places
*/
class exec_place_cuda_stream_impl : public exec_place::impl
{
public:
exec_place_cuda_stream_impl(const augmented_stream& dstream)
: exec_place::impl(data_place::device(dstream.dev_id))
, dstream_(dstream)
, dummy_pool_(dstream)
{}
::std::shared_ptr<exec_place::impl> get_place(size_t idx) override
{
_CCCL_ASSERT(idx == 0, "Index out of bounds for cuda_stream exec_place");
return shared_from_this();
}
exec_place activate(size_t idx) const override
{
_CCCL_ASSERT(idx == 0, "Index out of bounds for cuda_stream exec_place");
return exec_place::device(dstream_.dev_id).get_impl()->activate(0);
}
void deactivate(const exec_place& prev, size_t idx = 0) const override
{
_CCCL_ASSERT(idx == 0, "Index out of bounds for cuda_stream exec_place");
exec_place::device(dstream_.dev_id).get_impl()->deactivate(prev, 0);
}
bool is_device() const override
{
return true;
}
stream_pool& get_stream_pool(bool, exec_place_resources&, const exec_place&) const override
{
// User-stream places carry their own single-stream pool and intentionally
// ignore the registry.
return dummy_pool_;
}
::std::string to_string() const override
{
return "cuda_stream(id=" + ::std::to_string(dstream_.id) + " dev=" + ::std::to_string(dstream_.dev_id) + ")";
}
int cmp(const exec_place::impl& rhs) const override
{
if (typeid(*this) != typeid(rhs))
{
return typeid(*this).before(typeid(rhs)) ? -1 : 1;
}
const auto& other = static_cast<const exec_place_cuda_stream_impl&>(rhs);
return (other.dstream_.stream < dstream_.stream) - (dstream_.stream < other.dstream_.stream);
}
size_t hash() const override
{
return ::std::hash<cudaStream_t>()(dstream_.stream);
}
private:
augmented_stream dstream_;
mutable stream_pool dummy_pool_;
};
inline exec_place exec_place::cuda_stream(cudaStream_t stream)
{
int devid = get_device_from_stream(stream);
return exec_place{
::std::make_shared<exec_place_cuda_stream_impl>(augmented_stream(stream, get_stream_id(stream), devid))};
}
inline exec_place exec_place::cuda_stream(const augmented_stream& dstream)
{
return exec_place{::std::make_shared<exec_place_cuda_stream_impl>(dstream)};
}
} // end namespace cuda::experimental::places

668
cccl_upstream/cudax/include/cuda/experimental/__places/exec/green_context.cuh
View File

@@ -1,668 +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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief Implementation of green context places
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/experimental/__places/data_place_interface.cuh>
#include <cuda/experimental/__places/exec/cuda_context.cuh>
#include <cuda/experimental/__places/exec/green_ctx_view.cuh>
#include <cuda/experimental/__places/places.cuh>
#include <cuda/experimental/__stf/utility/hash.cuh>
// Used only for unit tests, not in the actual implementation
#ifdef UNITTESTED_FILE
# include <map>
#endif
#if _CCCL_CTK_AT_LEAST(12, 4)
namespace cuda::experimental::places
{
/* Get the unique ID associated with a context (overloaded) */
inline unsigned long long get_cuda_context_id(CUcontext ctx)
{
return cuda_try<cuCtxGetId>(ctx);
}
/* Get the unique ID associated with a green context (overloaded) */
inline unsigned long long get_cuda_context_id(CUgreenCtx gctx)
{
return get_cuda_context_id(cuda_try<cuCtxFromGreenCtx>(gctx));
}
/**
* @brief data_place_interface implementation for green contexts
*
* Green contexts allow partitioning GPU resources (SMs, memory bandwidth)
* for fine-grained control over execution. This class provides the
* data_place_interface for green context-based data locations.
*/
class green_ctx_data_place_impl : public data_place_interface
{
public:
explicit green_ctx_data_place_impl(green_ctx_view view)
: view_(mv(view))
{}
bool is_resolved() const override
{
return true;
}
int get_device_ordinal() const override
{
return view_.devid;
}
::std::string to_string() const override
{
return "green_ctx(dev=" + ::std::to_string(view_.devid)
+ ", ctx=" + ::std::to_string(get_cuda_context_id(view_.g_ctx)) + ")";
}
size_t hash() const override
{
return hash_all(view_.g_ctx, view_.devid);
}
int cmp(const data_place_interface& other) const override
{
if (typeid(*this) != typeid(other))
{
return typeid(*this).before(typeid(other)) ? -1 : 1;
}
const auto& o = static_cast<const green_ctx_data_place_impl&>(other);
return (o.view_ < view_) - (view_ < o.view_);
}
const green_ctx_view& get_view() const
{
return view_;
}
void* allocate(::std::ptrdiff_t size, cudaStream_t stream) const override
{
void* result = nullptr;
cuda_try(cudaSetDevice(view_.devid));
cuda_try(cudaMallocAsync(&result, static_cast<size_t>(size), stream));
return result;
}
void deallocate(void* ptr, size_t /*size*/, cudaStream_t stream) const override
{
cuda_try(cudaFreeAsync(ptr, stream));
}
bool allocation_is_stream_ordered() const override
{
return true;
}
CUresult mem_create(CUmemGenericAllocationHandle* handle, size_t size) const override
{
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.location.id = view_.devid;
return cuMemCreate(handle, size, &prop, 0);
}
::std::shared_ptr<void> get_affine_exec_impl() const override;
private:
green_ctx_view view_;
};
/**
* @brief Create a green context data place
*
* @param gc_view The green context view
* @return data_place for the green context
*/
inline data_place make_green_ctx_data_place(const green_ctx_view& gc_view)
{
return data_place(::std::make_shared<green_ctx_data_place_impl>(gc_view));
}
/**
* @brief Helper class to create views of green contexts that can be used as execution places
*/
class green_context_helper
{
public:
/* Create green contexts with sm_count SMs per context on a specific device (current device by default) */
green_context_helper(int sm_count, int devid = cuda_try<cudaGetDevice>())
: devid(devid)
, numsm(sm_count)
{
assert(devid >= 0);
const int old_device = cuda_try<cudaGetDevice>();
// Change device only if necessary.
if (devid != old_device)
{
cuda_try(cudaSetDevice(devid));
}
/* Make sure we aren't requesting more SMs than the GPU has available */
int max_SMs = cuda_try<cudaDeviceGetAttribute>(cudaDevAttrMultiProcessorCount, devid);
assert(max_SMs >= int(numsm));
/* Determine the device's resources */
CUdevice device = cuda_try<cuDeviceGet>(devid);
/* Retain the primary ctx in order to get a set of SM resources for that device */
CUdevResource input;
CUcontext primaryCtx = cuda_try<cuDevicePrimaryCtxRetain>(device);
cuCtxGetDevResource(primaryCtx, &input, CU_DEV_RESOURCE_TYPE_SM);
cuDevicePrimaryCtxRelease(device);
// First we query how many groups should be created
unsigned int nbGroups;
cuda_try(cuDevSmResourceSplitByCount(nullptr, &nbGroups, &input, nullptr, 0, sm_count));
// Split the resources as requested
assert(nbGroups >= 1);
resources.resize(nbGroups);
cuda_try(cuDevSmResourceSplitByCount(resources.data(), &nbGroups, &input, &remainder, 0, sm_count));
/* Create a green context for each group */
ctxs.resize(nbGroups);
// Create pools of CUDA streams
pools.reserve(nbGroups);
for (int i = 0; i < static_cast<int>(nbGroups); i++)
{
if (resources[i].type != CU_DEV_RESOURCE_TYPE_INVALID)
{
// Create a descriptor and a green context with that descriptor:
/* The generated resource descriptor is necessary for the creation of green contexts via the
* cuGreenCtxCreate API. The API expects nbResources == 1, as there is only one type of resource and
* merging the same types of resource is currently not supported. */
CUdevResourceDesc localdesc = cuda_try<cuDevResourceGenerateDesc>(&resources[i], 1);
// Create a green context
ctxs[i] = cuda_try<cuGreenCtxCreate>(localdesc, device, CU_GREEN_CTX_DEFAULT_STREAM);
pools.emplace_back(exec_place::impl::pool_size);
}
}
}
green_context_helper() = default;
~green_context_helper() = default;
public:
size_t get_device_id() const
{
return devid;
}
CUgreenCtx partition(size_t partition = 0)
{
return ctxs[partition];
}
green_ctx_view get_view(size_t id)
{
return green_ctx_view(ctxs[id], pools[id], devid);
}
stream_pool& get_pool(size_t gc_id)
{
assert(gc_id < pools.size());
return pools[gc_id];
}
size_t get_count() const
{
return ctxs.size();
}
private:
friend class exec_place;
// resources to define how we split the device(s) into green contexts
::std::vector<CUdevResource> resources;
::std::vector<stream_pool> pools;
CUdevResource remainder = {};
int devid = -1;
// Number of SMs requested per green context
size_t numsm = 0;
::std::vector<CUgreenCtx> ctxs;
};
/*
* Green-context execution places are implemented with the generic
* externally-owned-context place (`exec_place_cuda_ctx_impl`): the only
* functional use of the CUgreenCtx is deriving its CUcontext, and
* cuCtxFromGreenCtx returns the same CUcontext on every call, so the
* conversion is done once here and the CUcontext is the canonical identity
* of the place. As a consequence, a place built from a green_ctx_view and a
* place built with exec_place::cuda_context() from the converted context
* compare equal, which is the desired semantics.
*
* The user is responsible for keeping the underlying CUgreenCtx alive while
* the place (and its stream pool) is in use.
*/
inline exec_place exec_place::green_ctx(const green_ctx_view& gc_view, bool use_green_ctx_data_place)
{
CUcontext ctx = cuda_try<cuCtxFromGreenCtx>(gc_view.g_ctx);
data_place affine = use_green_ctx_data_place ? make_green_ctx_data_place(gc_view) : data_place::device(gc_view.devid);
return exec_place(::std::make_shared<exec_place_cuda_ctx_impl>(ctx, gc_view.devid, gc_view.pool, mv(affine)));
}
inline ::std::shared_ptr<void> green_ctx_data_place_impl::get_affine_exec_impl() const
{
return exec_place::green_ctx(view_).get_impl();
}
inline data_place data_place::green_ctx(const green_ctx_view& gc_view)
{
return make_green_ctx_data_place(gc_view);
}
# ifdef UNITTESTED_FILE
UNITTEST("green context exec_place equality")
{
green_context_helper gc_helper(8, 0); // 8 SMs per green context
// Need at least 2 green contexts for the test
if (gc_helper.get_count() < 2)
{
return;
}
auto gc0_view = gc_helper.get_view(0);
auto gc1_view = gc_helper.get_view(1);
// Create exec_places from different green contexts (default: use_green_ctx_data_place=false)
auto p0a = exec_place::green_ctx(gc0_view);
auto p0b = exec_place::green_ctx(gc0_view); // same green context as p0a
auto p1 = exec_place::green_ctx(gc1_view); // different green context
// Same green context should be equal
EXPECT(p0a == p0b);
EXPECT(!(p0a != p0b));
// Different green contexts should NOT be equal
EXPECT(p0a != p1);
EXPECT(!(p0a == p1));
// Green context exec_place should not be equal to regular device exec_place
auto dev0 = exec_place::device(0);
EXPECT(p0a != dev0);
EXPECT(!(p0a == dev0));
};
UNITTEST("green_ctx place equals the cuda_context place wrapping the same partition")
{
green_context_helper gc_helper(8, 0); // 8 SMs per green context
if (gc_helper.get_count() < 1)
{
return;
}
auto view = gc_helper.get_view(0);
// Identity is keyed on the converted CUcontext, which is a stable accessor
// (cuCtxFromGreenCtx returns the same handle on every call), so both
// construction routes must yield equal places.
auto p_green = exec_place::green_ctx(view);
auto p_ctx = exec_place::cuda_context(cuda_try<cuCtxFromGreenCtx>(view.g_ctx), view.devid);
EXPECT(p_green == p_ctx);
EXPECT(!(p_green != p_ctx));
};
UNITTEST("green context data_place equality")
{
green_context_helper gc_helper(8, 0);
if (gc_helper.get_count() < 2)
{
return;
}
auto gc0_view = gc_helper.get_view(0);
auto gc1_view = gc_helper.get_view(1);
// Create green context data places
auto dp0a = data_place::green_ctx(gc0_view);
auto dp0b = data_place::green_ctx(gc0_view);
auto dp1 = data_place::green_ctx(gc1_view);
// Same green context data place should be equal
EXPECT(dp0a == dp0b);
EXPECT(!(dp0a != dp0b));
// Different green context data places should NOT be equal
EXPECT(dp0a != dp1);
EXPECT(!(dp0a == dp1));
// Green context data place should not be equal to regular device data place
auto dev0 = data_place::device(0);
EXPECT(dp0a != dev0);
EXPECT(!(dp0a == dev0));
// Green context data place should be resolved but not a plain device
EXPECT(dp0a.is_resolved());
EXPECT(!dp0a.is_device());
EXPECT(dev0.is_resolved());
EXPECT(dev0.is_device());
};
UNITTEST("green context data_place to_string distinguishes contexts on the same device")
{
green_context_helper gc_helper(8, 0);
// Two green contexts on the same device are needed to exercise the fact that
// to_string() embeds the green context handle in addition to the device id.
if (gc_helper.get_count() < 2)
{
return;
}
auto gc0_view = gc_helper.get_view(0);
auto gc1_view = gc_helper.get_view(1);
auto dp0 = data_place::green_ctx(gc0_view);
auto dp1 = data_place::green_ctx(gc1_view);
const ::std::string s0 = dp0.to_string();
const ::std::string s1 = dp1.to_string();
// Both places live on device 0 and expose the green context in their name.
EXPECT(s0.find("dev=0") != ::std::string::npos);
EXPECT(s1.find("dev=0") != ::std::string::npos);
EXPECT(s0.find("ctx=") != ::std::string::npos);
EXPECT(s1.find("ctx=") != ::std::string::npos);
// Different green contexts on the same device must produce different names.
EXPECT(s0 != s1);
// A fresh place for the same green context yields the same name.
auto dp0_copy = data_place::green_ctx(gc0_view);
EXPECT(dp0.to_string() == dp0_copy.to_string());
};
UNITTEST("green context exec_place equality with green_ctx_data_place flag")
{
green_context_helper gc_helper(8, 0);
if (gc_helper.get_count() < 2)
{
return;
}
auto gc0_view = gc_helper.get_view(0);
auto gc1_view = gc_helper.get_view(1);
// Create exec_places with use_green_ctx_data_place=true
auto p0a = exec_place::green_ctx(gc0_view, true);
auto p0b = exec_place::green_ctx(gc0_view, true);
auto p1 = exec_place::green_ctx(gc1_view, true);
// Same green context should be equal
EXPECT(p0a == p0b);
// Different green contexts should NOT be equal
EXPECT(p0a != p1);
// Affine data place should be resolved but not a plain device when use_green_ctx_data_place=true
EXPECT(p0a.affine_data_place().is_resolved());
EXPECT(!p0a.affine_data_place().is_device());
};
UNITTEST("green context exec_place and data_place with different data place modes")
{
green_context_helper gc_helper(8, 0);
if (gc_helper.get_count() < 2)
{
return;
}
auto gc0_view = gc_helper.get_view(0);
auto gc1_view = gc_helper.get_view(1);
// Create exec_places for same green context with different use_green_ctx_data_place settings
auto ep0_device_affine = exec_place::green_ctx(gc0_view, false); // affine = device data place
auto ep0_green_affine = exec_place::green_ctx(gc0_view, true); // affine = green ctx data place
// Same green context exec_places should be equal regardless of data place mode
// (exec_place identity is about the green context, not the affine data place)
EXPECT(ep0_device_affine == ep0_green_affine);
// But their affine data places should be different
EXPECT(ep0_device_affine.affine_data_place() != ep0_green_affine.affine_data_place());
EXPECT(ep0_device_affine.affine_data_place().is_device());
EXPECT(!ep0_green_affine.affine_data_place().is_device());
// Different green contexts should NOT be equal, regardless of data place mode
auto ep1_device_affine = exec_place::green_ctx(gc1_view, false);
auto ep1_green_affine = exec_place::green_ctx(gc1_view, true);
EXPECT(ep0_device_affine != ep1_device_affine);
EXPECT(ep0_device_affine != ep1_green_affine);
EXPECT(ep0_green_affine != ep1_device_affine);
EXPECT(ep0_green_affine != ep1_green_affine);
// Test green context data places directly
auto dp0 = data_place::green_ctx(gc0_view);
auto dp1 = data_place::green_ctx(gc1_view);
// Different green context data places should NOT be equal
EXPECT(dp0 != dp1);
// Green context data place should NOT equal regular device data place
EXPECT(dp0 != data_place::device(0));
// Green context data place should equal affine of exec_place with use_green_ctx_data_place=true
EXPECT(dp0 == ep0_green_affine.affine_data_place());
};
UNITTEST("green context data_place as unordered_map key")
{
green_context_helper gc_helper(8, 0);
if (gc_helper.get_count() < 2)
{
return;
}
auto gc0_view = gc_helper.get_view(0);
auto gc1_view = gc_helper.get_view(1);
// Create green context-specific data places (the kind used when
// use_green_ctx_data_place = true). These are distinct from data_place::device(0).
auto dp0 = data_place::green_ctx(gc0_view);
auto dp1 = data_place::green_ctx(gc1_view);
// Different green contexts on the same device must be distinguished as different keys.
EXPECT(dp0 != dp1);
// Both are different from the regular device data place
EXPECT(dp0 != data_place::device(0));
EXPECT(dp1 != data_place::device(0));
::std::unordered_map<data_place, int, hash<data_place>> map;
// Insert green context data places - different green contexts should be different keys
map[dp0] = 100;
map[dp1] = 200;
// Verify lookups work correctly
EXPECT(map[dp0] == 100);
EXPECT(map[dp1] == 200);
EXPECT(map.size() == 2);
// Verify that a new data_place for the same green context finds the same entry
auto dp0_copy = data_place::green_ctx(gc0_view);
EXPECT(map[dp0_copy] == 100);
// Mix with regular device data place
map[data_place::device(0)] = 300;
EXPECT(map.size() == 3);
EXPECT(map[data_place::device(0)] == 300);
// Green context data place and device data place should be different keys
EXPECT(map[dp0] == 100); // Still 100, not overwritten
};
UNITTEST("green context exec_place as unordered_map key")
{
green_context_helper gc_helper(8, 0);
if (gc_helper.get_count() < 2)
{
return;
}
auto gc0_view = gc_helper.get_view(0);
auto gc1_view = gc_helper.get_view(1);
// Create exec_places without use_green_ctx_data_place flag (default).
// Their affine data_place is data_place::device(0), not a green context-specific one.
auto ep0 = exec_place::green_ctx(gc0_view);
auto ep1 = exec_place::green_ctx(gc1_view);
// Both share the same affine data_place (the device), but they must still be
// distinguished as different exec_place keys in the map.
EXPECT(ep0.affine_data_place() == ep1.affine_data_place());
EXPECT(ep0.affine_data_place() == data_place::device(0));
EXPECT(ep0 != ep1);
::std::unordered_map<exec_place, int, hash<exec_place>> map;
// Insert green context exec places - different green contexts should be different keys
// even though they share the same affine data_place
map[ep0] = 100;
map[ep1] = 200;
// Verify lookups work correctly
EXPECT(map[ep0] == 100);
EXPECT(map[ep1] == 200);
EXPECT(map.size() == 2);
// Verify that a new exec_place for the same green context finds the same entry
auto ep0_copy = exec_place::green_ctx(gc0_view);
EXPECT(map[ep0_copy] == 100);
// Mix with regular device exec place - should be a different key
map[exec_place::device(0)] = 300;
EXPECT(map.size() == 3);
EXPECT(map[exec_place::device(0)] == 300);
// Green context exec place should still have its value
EXPECT(map[ep0] == 100);
};
UNITTEST("green context data_place as std::map key")
{
green_context_helper gc_helper(8, 0);
if (gc_helper.get_count() < 2)
{
return;
}
auto gc0_view = gc_helper.get_view(0);
auto gc1_view = gc_helper.get_view(1);
auto dp0 = data_place::green_ctx(gc0_view);
auto dp1 = data_place::green_ctx(gc1_view);
// Different green contexts must be distinguished
EXPECT(dp0 != dp1);
::std::map<data_place, int> map;
// Insert green context data places
map[dp0] = 100;
map[dp1] = 200;
// Verify lookups work correctly
EXPECT(map[dp0] == 100);
EXPECT(map[dp1] == 200);
EXPECT(map.size() == 2);
// Verify that a new data_place for the same green context finds the same entry
auto dp0_copy = data_place::green_ctx(gc0_view);
EXPECT(map[dp0_copy] == 100);
// Mix with regular device data place
map[data_place::device(0)] = 300;
EXPECT(map.size() == 3);
EXPECT(map[data_place::device(0)] == 300);
EXPECT(map[dp0] == 100); // Still 100, not overwritten
};
UNITTEST("green context exec_place as std::map key")
{
green_context_helper gc_helper(8, 0);
if (gc_helper.get_count() < 2)
{
return;
}
auto gc0_view = gc_helper.get_view(0);
auto gc1_view = gc_helper.get_view(1);
auto ep0 = exec_place::green_ctx(gc0_view);
auto ep1 = exec_place::green_ctx(gc1_view);
// Both share the same affine data_place but must be distinguished
EXPECT(ep0.affine_data_place() == ep1.affine_data_place());
EXPECT(ep0 != ep1);
::std::map<exec_place, int> map;
// Insert green context exec places
map[ep0] = 100;
map[ep1] = 200;
// Verify lookups work correctly
EXPECT(map[ep0] == 100);
EXPECT(map[ep1] == 200);
EXPECT(map.size() == 2);
// Verify that a new exec_place for the same green context finds the same entry
auto ep0_copy = exec_place::green_ctx(gc0_view);
EXPECT(map[ep0_copy] == 100);
// Mix with regular device exec place
map[exec_place::device(0)] = 300;
EXPECT(map.size() == 3);
EXPECT(map[exec_place::device(0)] == 300);
EXPECT(map[ep0] == 100); // Still 100
};
# endif // UNITTESTED_FILE
} // end namespace cuda::experimental::places
#endif // _CCCL_CTK_AT_LEAST(12, 4)

83
cccl_upstream/cudax/include/cuda/experimental/__places/exec/green_ctx_view.cuh
View File

@@ -1,83 +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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
/**
* @file
* @brief Implementation of green context views
*/
#include <cuda/experimental/__places/stream_pool.cuh>
#include <cuda/experimental/__stf/utility/hash.cuh>
#if _CCCL_CTK_AT_LEAST(12, 4)
namespace cuda::experimental::places
{
template <typename T>
struct hash;
// Green contexts are only supported since CUDA 12.4
/**
* @brief View of a green context and a pool of CUDA streams
*/
class green_ctx_view
{
public:
green_ctx_view(CUgreenCtx g_ctx, stream_pool pool, int devid)
: g_ctx(g_ctx)
, pool(::cuda::experimental::stf::mv(pool))
, devid(devid)
{}
CUgreenCtx g_ctx;
stream_pool pool;
int devid;
bool operator==(const green_ctx_view& other) const
{
return (g_ctx == other.g_ctx) && (devid == other.devid);
}
bool operator<(const green_ctx_view& other) const
{
if (g_ctx != other.g_ctx)
{
return g_ctx < other.g_ctx;
}
return devid < other.devid;
}
};
/**
* @brief Specialization of `places::hash` for `green_ctx_view`
*/
template <>
struct hash<green_ctx_view>
{
::std::size_t operator()(const green_ctx_view& k) const
{
return ::cuda::experimental::stf::hash_all(k.g_ctx, k.devid);
}
};
} // end namespace cuda::experimental::places
#endif // _CCCL_CTK_AT_LEAST(12, 4)

134
cccl_upstream/cudax/include/cuda/experimental/__places/exec_place_resources.cuh
View File

@@ -1,134 +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 Standalone per-place stream-pool registry.
*
* `exec_place_resources` owns a `{compute, data}` `stream_pool` slot for every
* pooled place it is queried with. Slots are created lazily on first use and
* destroyed with the registry. The registry depends only on `stream_pool.cuh`
* and a forward declaration of `exec_place`; it can be embedded in any
* resource container (e.g. `async_resources_handle`) without pulling in STF.
*
* Keys are `exec_place::impl*` pointers. Pooled implementations (`device(N)`,
* `host()`) live as process-wide singleton impls, so pointer identity matches
* place identity for them. Self-contained implementations (`cuda_stream`,
* green-context, grid) override `get_stream_pool` and never reach the
* registry.
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/experimental/__places/stream_pool.cuh>
#include <mutex>
#include <unordered_map>
namespace cuda::experimental::places
{
/**
* @brief Default size of each per-place stream pool created by the registry.
*
* `exec_place::impl::pool_size` and `data_pool_size` are aliases to these
* values so `places.cuh` can keep its public surface unchanged.
*/
inline constexpr ::std::size_t exec_place_default_pool_size = 4;
inline constexpr ::std::size_t exec_place_default_data_pool_size = 4;
/**
* @brief A registry of per-place stream pools keyed by `exec_place::impl*`.
*
* For every distinct pooled impl pointer the registry is queried with, it
* owns one `{compute, data}` pair of `stream_pool`s, created lazily on first
* lookup with sizes `exec_place_default_pool_size` /
* `exec_place_default_data_pool_size`.
*
* The map itself is mutex-guarded. The mutex is only held across the
* find/insert into the map; subsequent stream creation (which happens lazily
* inside `stream_pool::next`) runs outside the lock, so contention is limited
* to slow-path task submission.
*
* Lifetime: each entry's pool is owned by the registry. Destroying the
* registry destroys every pool it has created (and their cached
* `cudaStream_t` handles). Consequently, a registry must not outlive the
* CUDA primary context(s) of the devices it has cached streams for; with
* this design, registries are typically embedded in an
* `async_resources_handle` and share the lifetime of the owning STF context.
*
* Caveats for externally-owned places:
* - User-stream places (`exec_place::cuda_stream(s)`) carry their own
* single-stream pool and never participate in the registry.
* - Green-context places carry their own pool (constructed from the
* `green_ctx_view`) and also bypass the registry. The user must keep the
* underlying `CUgreenCtx` alive as long as the place is used.
*/
class exec_place_resources
{
public:
struct per_place_pools
{
per_place_pools()
: compute(exec_place_default_pool_size)
, data(exec_place_default_data_pool_size)
{}
stream_pool compute;
stream_pool data;
};
exec_place_resources() = default;
exec_place_resources(const exec_place_resources&) = delete;
exec_place_resources& operator=(const exec_place_resources&) = delete;
exec_place_resources(exec_place_resources&&) = delete;
exec_place_resources& operator=(exec_place_resources&&) = delete;
/**
* @brief Look up (or lazily create) the `{compute, data}` pool slot for the
* supplied impl pointer.
*
* Thread-safe: the mutex is held only across the find/insert. The returned
* reference is stable for the lifetime of the registry (`std::unordered_map`
* preserves node addresses across rehashes).
*/
[[nodiscard]] per_place_pools& get(const void* impl_key)
{
::std::scoped_lock lock(mtx_);
auto it = map_.find(impl_key);
if (it == map_.end())
{
it = map_.emplace(impl_key, per_place_pools{}).first;
}
return it->second;
}
/// @brief Number of per-place entries currently cached. Mainly for tests.
[[nodiscard]] ::std::size_t size() const
{
::std::scoped_lock lock(mtx_);
return map_.size();
}
private:
mutable ::std::mutex mtx_;
::std::unordered_map<const void*, per_place_pools> map_;
};
} // namespace cuda::experimental::places

630
cccl_upstream/cudax/include/cuda/experimental/__places/localized_array.cuh
View File

@@ -1,630 +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.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief Implementation of the localized_array class which dispatches a VMM
* allocation over multiple data places using a partitioner function.
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/std/__tuple_dir/get.h>
#include <cuda/std/__tuple_dir/tuple.h>
#include <cuda/experimental/__places/places.cuh>
#include <array>
#include <functional>
#include <random>
#include <unordered_map>
#include <vector>
namespace cuda::experimental::places
{
/**
* @brief Check if localized allocation statistics should be printed
*/
inline bool localized_alloc_stats_enabled()
{
static bool enabled = [] {
const char* env = ::std::getenv("CUDASTF_LOCALIZED_ALLOC_STATS");
return env != nullptr && ::std::string(env) != "0";
}();
return enabled;
}
//! Default number of elements sampled per block to decide the block's owner
inline constexpr size_t localized_placement_default_probes = 10;
/**
* @brief Statistics describing how a localized allocation - or a dry-run
* evaluation of one - distributes a tensor over data places.
*
* Produced by evaluate_localized_placement() and by localized_array (see
* localized_array::get_stats()). This is the returnable form of the report
* previously only printed to stderr under CUDASTF_LOCALIZED_ALLOC_STATS.
*/
struct localized_stats
{
size_t total_bytes = 0; //!< requested payload size in bytes
size_t vm_bytes = 0; //!< block-rounded virtual reservation size in bytes
size_t block_size = 0; //!< placement granularity in bytes
size_t nblocks = 0; //!< number of placement blocks
size_t nallocs = 0; //!< physical allocations after merging same-owner runs
size_t total_samples = 0; //!< probes drawn by the block-owner sampler
size_t matching_samples = 0; //!< probes agreeing with the chosen block owner
//! Bytes backed by each place, keyed by data_place::to_string()
::std::unordered_map<::std::string, size_t> bytes_per_place;
//! Bytes owned by each grid position, keyed by the position's linear index
//! (dim4::get_index of the pos4; friendlier than strings across FFI)
::std::unordered_map<size_t, size_t> bytes_per_grid_index;
//! Fraction of sampled elements whose owner matches the block-majority
//! owner: an estimate of the fraction of bytes that end up local to their
//! owner once ownership is quantized to blocks.
double accuracy() const
{
return total_samples == 0 ? 1.0 : static_cast<double>(matching_samples) / static_cast<double>(total_samples);
}
};
/**
* @brief Placement granularity used when the caller does not specify one:
* the device allocation granularity when a device is present, or the
* customary 2 MiB VMM granularity for GPU-free (offline) evaluation. The
* granularity query is the only driver interaction.
*/
inline size_t default_placement_block_size()
{
int ndevs = 0;
if (cudaGetDeviceCount(&ndevs) == cudaSuccess && ndevs > 0)
{
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.location.id = 0;
return cuda_try<cuMemGetAllocationGranularity>(&prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM);
}
cudaGetLastError();
return 2 * 1024 * 1024;
}
/**
* @brief Decide the owner of each placement block by sampled majority vote.
*
* For each block, `probes` elements are sampled (reproducibly: one seeded
* generator for the whole computation) and the most frequent owner wins. The
* majority vote is what tolerates partitions whose boundaries do not align
* with the block granularity: a block straddling two owners goes to the one
* owning most of it.
*
* @param owner_of Callable mapping a linear element index to the pos4 of its
* owner in the grid
* @param nblocks Number of placement blocks
* @param block_size_bytes Size of a placement block in bytes (must be at
* least ``elemsize``)
* @param elemsize Size of one element in bytes (must be at least 1)
* @param total_elems Total number of elements (probes are clipped to it)
* @param probes Number of samples per block
* @param stats Accumulates total/matching sample counts
*/
template <typename OwnerFn>
::std::vector<pos4> compute_block_owners(
OwnerFn&& owner_of,
size_t nblocks,
size_t block_size_bytes,
size_t elemsize,
size_t total_elems,
size_t probes,
localized_stats& stats)
{
if (elemsize == 0 || block_size_bytes < elemsize)
{
throw ::std::invalid_argument("placement blocks must hold at least one element (elemsize in [1, block size])");
}
const size_t block_elems = block_size_bytes / elemsize;
// Fixed seed: placement must be reproducible from one run to the next
::std::mt19937 gen(0x5EED);
::std::uniform_int_distribution<size_t> dis(0, block_elems - 1);
probes = ::std::max<size_t>(1, ::std::min(probes, block_elems));
::std::vector<pos4> owners;
owners.reserve(nblocks);
::std::vector<pos4> sampled_pos(probes);
for (size_t i = 0; i < nblocks; i++)
{
// First element of the block (exact, so non-dividing element sizes do
// not accumulate drift across blocks)
const size_t block_start = i * block_size_bytes / elemsize;
for (size_t sample = 0; sample < probes; sample++)
{
// Clip: the last block may extend past the payload
const size_t index = ::std::min(block_start + dis(gen), total_elems - 1);
sampled_pos[sample] = owner_of(index);
}
::std::unordered_map<pos4, size_t, ::cuda::experimental::stf::hash<pos4>> sample_cnt;
for (const auto& s : sampled_pos)
{
++sample_cnt[s];
}
size_t max_cnt = 0;
pos4 max_pos;
for (const auto& s : sample_cnt)
{
if (s.second > max_cnt)
{
max_pos = s.first;
max_cnt = s.second;
}
}
stats.total_samples += probes;
stats.matching_samples += max_cnt;
owners.push_back(max_pos);
}
return owners;
}
/**
* @brief Call `fn(owner, first_block, num_blocks)` for each maximal run of
* consecutive blocks with the same owner.
*/
template <typename F>
void for_each_owner_run(const ::std::vector<pos4>& owners, F&& fn)
{
for (size_t i = 0; i < owners.size();)
{
const pos4 p = owners[i];
size_t j = 0;
while ((i + j < owners.size()) && (owners[i + j] == p))
{
j++;
}
fn(p, i, j);
i += j;
}
}
/**
* @brief An allocator that takes a mapping function to dispatch an allocation over multiple data places.
*
* This is the mechanism used to implement the data_place of a grid of execution places.
* Uses CUDA Virtual Memory Management (VMM) to create a contiguous virtual address range
* backed by physical allocations on different devices according to the partitioner.
*/
class localized_array
{
struct metadata
{
metadata(data_place place_, size_t size_, size_t offset_)
: alloc_handle{}
, place(mv(place_))
, size(size_)
, offset(offset_)
{}
CUmemGenericAllocationHandle alloc_handle;
const data_place place;
size_t size;
size_t offset;
};
public:
template <typename F>
localized_array(
exec_place grid,
partition_fn_t mapper,
F&& delinearize,
size_t total_size,
size_t elemsize,
dim4 data_dims,
size_t probes = localized_placement_default_probes)
: grid(mv(grid))
, mapper(mv(mapper))
, total_size_bytes(total_size * elemsize)
, data_dims(data_dims)
, elemsize(elemsize)
{
const dim4 grid_dims = this->grid.get_dims();
init(
[&](size_t index) {
const pos4 coords = delinearize(index);
pos4 eplace_coords(0);
this->mapper(&eplace_coords, coords, this->data_dims, grid_dims);
return eplace_coords;
},
total_size,
probes);
}
/**
* @brief Construct from a generic owner function instead of a raw
* partition_fn_t mapper (e.g. a stateful partition object). The owner
* function maps a linear element index to the grid position owning it and
* is only used during construction.
*/
localized_array(exec_place grid,
const ::std::function<pos4(size_t)>& owner_of,
size_t total_size,
size_t elemsize,
dim4 data_dims,
size_t probes = localized_placement_default_probes)
: grid(mv(grid))
, total_size_bytes(total_size * elemsize)
, data_dims(data_dims)
, elemsize(elemsize)
{
init(owner_of, total_size, probes);
}
localized_array() = delete;
localized_array(const localized_array&) = delete;
localized_array(localized_array&&) = delete;
localized_array& operator=(const localized_array&) = delete;
localized_array& operator=(localized_array&&) = delete;
~localized_array()
{
for (auto& item : meta)
{
size_t offset = item.offset;
size_t sz = item.size;
cuda_try(cuMemUnmap(base_ptr + offset, sz));
cuda_try(cuMemRelease(item.alloc_handle));
}
cuda_try(cuMemAddressFree(base_ptr, vm_total_size_bytes));
}
void* get_base_ptr() const
{
return reinterpret_cast<void*>(base_ptr);
}
/**
* @brief Placement statistics of this allocation (see localized_stats)
*/
const localized_stats& get_stats() const
{
return stats;
}
/*
* This equality operator is used to find entries in an allocation cache which match a specific request
*/
template <typename... P>
bool operator==(::cuda::std::tuple<P&...> t) const
{
// tuple arguments :
// 0 : grid, 1 : mapper, 2 : delinearize function, 3 : total size, 4 elem_size, 5 : data_dims
bool result = grid == ::cuda::std::get<0>(t) && mapper == ::cuda::std::get<1>(t)
&& this->total_size_bytes == ::cuda::std::get<3>(t) * ::cuda::std::get<4>(t)
&& elemsize == ::cuda::std::get<4>(t) && data_dims == ::cuda::std::get<5>(t);
if (result)
{
assert(this->total_size_bytes == ::cuda::std::get<3>(t) * ::cuda::std::get<4>(t));
assert(data_dims == ::cuda::std::get<5>(t));
}
return result;
}
private:
void init(const ::std::function<pos4(size_t)>& owner_of, size_t total_size, size_t probes)
{
if (elemsize == 0)
{
throw ::std::invalid_argument("localized_array requires an element size of at least 1 byte");
}
cuda_try(cudaFree(nullptr));
const int ndevs = cuda_try<cudaGetDeviceCount>();
CUdevice dev = cuda_try<cuCtxGetDevice>();
int supportsVMM = cuda_try<cuDeviceGetAttribute>(CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED, dev);
EXPECT(supportsVMM == 1, "Cannot create a localized_array object on this machine because it does not support VMM.");
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.location.id = dev;
size_t alloc_granularity_bytes = cuda_try<cuMemGetAllocationGranularity>(&prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM);
block_size_bytes = alloc_granularity_bytes;
vm_total_size_bytes =
((total_size_bytes + alloc_granularity_bytes - 1) / alloc_granularity_bytes) * alloc_granularity_bytes;
size_t nblocks = vm_total_size_bytes / alloc_granularity_bytes;
base_ptr = cuda_try<cuMemAddressReserve>(vm_total_size_bytes, 0ULL, 0ULL, 0ULL);
::std::vector<CUmemAccessDesc> accessDesc(ndevs);
for (int d = 0; d < ndevs; d++)
{
accessDesc[d].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
accessDesc[d].location.id = d;
accessDesc[d].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
}
stats.total_bytes = total_size_bytes;
stats.vm_bytes = vm_total_size_bytes;
stats.block_size = block_size_bytes;
stats.nblocks = nblocks;
const ::std::vector<pos4> owners =
compute_block_owners(owner_of, nblocks, block_size_bytes, elemsize, total_size, probes, stats);
meta.reserve(nblocks);
for_each_owner_run(owners, [&](pos4 p, size_t first_block, size_t num_blocks) {
data_place place = grid_pos_to_place(p);
size_t alloc_size = num_blocks * block_size_bytes;
stats.bytes_per_place[place.to_string()] += alloc_size;
stats.bytes_per_grid_index[this->grid.get_dims().get_index(p)] += alloc_size;
meta.emplace_back(mv(place), alloc_size, first_block * block_size_bytes);
});
stats.nallocs = meta.size();
if (localized_alloc_stats_enabled())
{
print_stats(owners);
}
for (auto& item : meta)
{
int item_dev = device_ordinal(item.place);
cuda_try(item.place.mem_create(&item.alloc_handle, item.size));
_CCCL_ASSERT(item.offset + item.size <= vm_total_size_bytes, "Allocation offset out of bounds");
cuda_try(cuMemMap(base_ptr + item.offset, item.size, 0ULL, item.alloc_handle, 0ULL));
for (int d = 0; d < ndevs; d++)
{
int set_access = 1;
if (item_dev != d)
{
set_access = cuda_try<cudaDeviceCanAccessPeer>(d, item_dev);
if (!set_access)
{
fprintf(stderr, "Warning : Cannot enable peer access between devices %d and %d\n", d, item_dev);
}
}
if (set_access == 1)
{
cuda_try(cuMemSetAccess(base_ptr + item.offset, item.size, &accessDesc[d], 1ULL));
}
}
}
}
void print_stats(const ::std::vector<pos4>& owners)
{
fprintf(stderr, "\n=== Localized Array Allocation Statistics ===\n");
fprintf(stderr, "Total size: %zu bytes (%.2f MB)\n", stats.total_bytes, stats.total_bytes / (1024.0 * 1024.0));
fprintf(stderr, "VM reservation: %zu bytes (%.2f MB)\n", stats.vm_bytes, stats.vm_bytes / (1024.0 * 1024.0));
fprintf(stderr, "Block size: %zu bytes (%.2f KB)\n", stats.block_size, stats.block_size / 1024.0);
fprintf(stderr, "Number of blocks: %zu (merged into %zu allocations)\n", stats.nblocks, stats.nallocs);
fprintf(stderr, "Number of places: %zu\n", stats.bytes_per_place.size());
fprintf(stderr, "\nAllocation distribution by place:\n");
for (const auto& entry : stats.bytes_per_place)
{
double pct = 100.0 * entry.second / stats.vm_bytes;
fprintf(stderr,
" %s: %zu bytes (%.2f MB, %.1f%%)\n",
entry.first.c_str(),
entry.second,
entry.second / (1024.0 * 1024.0),
pct);
}
if (stats.total_samples > 0)
{
fprintf(stderr,
"\nPlacement accuracy: %.1f%% (%zu/%zu samples matched chosen position)\n",
100.0 * stats.accuracy(),
stats.matching_samples,
stats.total_samples);
}
fprintf(stderr, "\nAllocation map (%zu allocations):\n", meta.size());
fprintf(stderr, " %-6s %-12s %-12s %-10s %s\n", "Index", "Offset", "Size", "Blocks", "Place");
fprintf(stderr, " %-6s %-12s %-12s %-10s %s\n", "-----", "------", "----", "------", "-----");
for (size_t idx = 0; idx < meta.size(); idx++)
{
const auto& item = meta[idx];
size_t num_blocks = item.size / stats.block_size;
size_t start_block = item.offset / stats.block_size;
(void) start_block;
fprintf(stderr,
" %-6zu %-12zu %-12zu %-10zu %s\n",
idx,
item.offset,
item.size,
num_blocks,
item.place.to_string().c_str());
}
fprintf(stderr, "\nBlock ownership map (each char = 1 block, 0-9/a-z = place index):\n ");
::std::unordered_map<::std::string, char> place_to_char;
char next_char = '0';
for (size_t i = 0; i < owners.size(); i++)
{
::std::string place_str = grid_pos_to_place(owners[i]).to_string();
if (place_to_char.find(place_str) == place_to_char.end())
{
place_to_char[place_str] = next_char;
if (next_char == '9')
{
next_char = 'a';
}
else
{
next_char++;
}
}
fprintf(stderr, "%c", place_to_char[place_str]);
if ((i + 1) % 80 == 0)
{
fprintf(stderr, "\n ");
}
}
fprintf(stderr, "\n");
fprintf(stderr, "\n Legend:\n");
for (const auto& entry : place_to_char)
{
fprintf(stderr, " %c = %s\n", entry.second, entry.first.c_str());
}
fprintf(stderr, "==============================================\n\n");
}
data_place grid_pos_to_place(pos4 grid_pos)
{
return grid.get_place(grid_pos).affine_data_place();
}
exec_place grid;
partition_fn_t mapper = nullptr;
::std::vector<metadata> meta;
size_t block_size_bytes = 0;
size_t total_size_bytes = 0;
size_t vm_total_size_bytes = 0;
CUdeviceptr base_ptr = 0;
dim4 data_dims;
size_t elemsize = 0;
localized_stats stats;
};
/**
* @brief Evaluate - without allocating anything - how a localized allocation
* would distribute a tensor over the places of a grid.
*
* Runs the exact same block-owner decision procedure as localized_array and
* returns the resulting statistics, so callers can score a candidate mapping
* (and tune its parameters) before committing memory.
*
* Extents follow the dimension-0-fastest convention of dim4::get_index().
*
* @param grid Grid of execution places the mapper distributes over
* @param mapper Partition function mapping element coordinates to a place
* @param data_dims Extents of the tensor
* @param elemsize Size of one element in bytes
* @param probes Number of samples per block for the majority vote
* @param block_size Placement granularity in bytes; 0 selects the device
* allocation granularity when a device is present, or a 2 MiB default
* otherwise (this granularity query is the only driver interaction)
*/
[[nodiscard]] inline localized_stats evaluate_localized_placement(
const exec_place& grid,
partition_fn_t mapper,
dim4 data_dims,
size_t elemsize,
size_t probes = localized_placement_default_probes,
size_t block_size = 0)
{
if (block_size == 0)
{
block_size = default_placement_block_size();
}
localized_stats stats;
const size_t total_elems = data_dims.size();
stats.total_bytes = total_elems * elemsize;
stats.vm_bytes = ((stats.total_bytes + block_size - 1) / block_size) * block_size;
stats.block_size = block_size;
stats.nblocks = stats.vm_bytes / block_size;
const dim4 grid_dims = grid.get_dims();
const ::std::vector<pos4> owners = compute_block_owners(
[&](size_t index) {
pos4 eplace_coords(0);
mapper(&eplace_coords, data_dims.index_to_pos(index), data_dims, grid_dims);
return eplace_coords;
},
stats.nblocks,
block_size,
elemsize,
total_elems,
probes,
stats);
for_each_owner_run(owners, [&](pos4 p, size_t /*first_block*/, size_t num_blocks) {
const data_place place = grid.get_place(p).affine_data_place();
stats.bytes_per_place[place.to_string()] += num_blocks * block_size;
stats.bytes_per_grid_index[grid.get_dims().get_index(p)] += num_blocks * block_size;
stats.nallocs++;
});
return stats;
}
inline ::std::unordered_map<void*, ::std::unique_ptr<localized_array>>& get_composite_alloc_registry()
{
static ::std::unordered_map<void*, ::std::unique_ptr<localized_array>> reg;
return reg;
}
inline void* allocate_composite_data_place(const data_place_composite& p, dim4 data_dims, size_t elemsize)
{
const exec_place& grid = p.get_grid();
const partition_fn_t& mapper = p.get_partitioner();
// Linear memory follows the dimension-0-fastest convention of
// dim4::get_index(), like STF slices; the partitioner receives true element
// coordinates within data_dims.
auto delinearize = [data_dims](size_t i) {
return data_dims.index_to_pos(i);
};
auto arr = ::std::make_unique<localized_array>(grid, mapper, delinearize, data_dims.size(), elemsize, data_dims);
void* ptr = arr->get_base_ptr();
get_composite_alloc_registry()[ptr] = ::std::move(arr);
return ptr;
}
inline void deallocate_composite_data_place(void* ptr)
{
get_composite_alloc_registry().erase(ptr);
}
} // namespace cuda::experimental::places

129
cccl_upstream/cudax/include/cuda/experimental/__places/machine.cuh
View File

@@ -1,129 +0,0 @@
//===----------------------------------------------------------------------===//
//
// Part of CUDA Experimental 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.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief The machine class abstract low-level CUDA mechanisms such as enabling P2P accesses
*
* This class should also provide information about the topology of the machine
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/experimental/__stf/utility/cuda_safe_call.cuh>
#include <cuda/experimental/__utility/meyers_singleton.cuh>
#include <cstdio>
#include <cuda_runtime_api.h>
namespace cuda::experimental::places::reserved
{
using ::cuda::experimental::stf::cuda_try;
/**
* @brief Singleton object abstracting a machine able to set up CUDA peer accesses.
*
*/
class machine : public ::cuda::experimental::meyers_singleton<machine>
{
protected:
machine()
: ndevices(cuda_try<cudaGetDeviceCount>())
{
// TODO: remove this call? Currently anyone getting a machine gets peer access
// and subsequent explicit calls to enable_peer_accesses() are no-ops.
enable_peer_accesses();
}
// Nobody can copy or assign.
machine& operator=(const machine&) = delete;
machine(const machine&) = delete;
// Clients can't destroy an object (because they can't create one in the first place).
~machine() = default;
public:
void enable_peer_accesses()
{
// Only once
if (initialized_peer_accesses)
{
return;
}
int current_dev = cuda_try<cudaGetDevice>();
for (int d = 0; d < ndevices; d++)
{
cuda_try(cudaSetDevice(d));
cudaMemPool_t mempool = cuda_try<cudaDeviceGetDefaultMemPool>(d);
for (int peer_d = 0; peer_d < ndevices; peer_d++)
{
if (peer_d == d)
{
continue;
}
int can_access_peer = cuda_try<cudaDeviceCanAccessPeer>(d, peer_d);
uint64_t threshold = UINT64_MAX;
cuda_try(cudaMemPoolSetAttribute(mempool, cudaMemPoolAttrReleaseThreshold, &threshold));
if (can_access_peer)
{
cudaError_t res = cudaDeviceEnablePeerAccess(peer_d, 0);
assert(res == cudaErrorPeerAccessAlreadyEnabled || res == cudaSuccess);
if (res == cudaErrorPeerAccessAlreadyEnabled)
{
fprintf(stderr, "[DEV %d] peer access already enabled with device %d\n", d, peer_d);
}
// Enable access to remote memory pool
cudaMemAccessDesc desc = {.location = {.type = cudaMemLocationTypeDevice, .id = peer_d},
.flags = cudaMemAccessFlagsProtReadWrite};
cuda_try(cudaMemPoolSetAccess(mempool, &desc, 1 /* numDescs */));
// ::std::cout << "[DEV " << d << "] cudaMemPoolSetAccess to peer "<< peer_d << ::'\n';
}
else
{
fprintf(stderr, "[DEV %d] cannot enable peer access with device %d\n", d, peer_d);
}
}
}
cuda_try(cudaSetDevice(current_dev));
initialized_peer_accesses = true;
}
// Naive solution
int get_ith_closest_node(int node, int ith)
{
int nnodes = ndevices + 1;
return (node + ith) % nnodes;
}
private:
bool initialized_peer_accesses = false;
int ndevices;
};
} // namespace cuda::experimental::places::reserved

178
cccl_upstream/cudax/include/cuda/experimental/__places/partitions/blocked_partition.cuh
View File

@@ -1,178 +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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief Definition of the `blocked_partition` strategy
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/experimental/__places/partitions/cyclic_shape.cuh>
#include <cuda/experimental/__places/places.cuh>
#include <cuda/experimental/__stf/utility/dimensions.cuh>
namespace cuda::experimental::places
{
template <::std::ptrdiff_t which_dim = -1>
class blocked_partition_custom
{
public:
blocked_partition_custom() = default;
template <size_t dimensions>
_CCCL_HOST_DEVICE static auto apply(const box<dimensions>& in, pos4 place_position, dim4 grid_dims)
{
::std::array<::std::pair<::std::ptrdiff_t, ::std::ptrdiff_t>, dimensions> bounds;
size_t target_dim = (which_dim == -1) ? dimensions - 1 : size_t(which_dim);
if (target_dim > dimensions - 1)
{
target_dim = dimensions - 1;
}
// if constexpr (dimensions > 1) {
for (size_t d = 0; d < dimensions; d++)
{
// First position in this dimension (included)
bounds[d].first = in.get_begin(d);
// Last position in this dimension (excluded)
bounds[d].second = in.get_end(d);
}
// }
size_t nplaces = grid_dims.x;
::std::ptrdiff_t dim_beg = bounds[target_dim].first;
::std::ptrdiff_t dim_end = bounds[target_dim].second;
size_t cnt = dim_end - dim_beg;
::std::ptrdiff_t part_size = (cnt + nplaces - 1) / nplaces;
// If first = second, this means it's an empty shape. This may happen
// when there are more entries in grid_dims than in the shape for
// example
bounds[target_dim].first = ::std::min(dim_beg + part_size * place_position.x, dim_end);
bounds[target_dim].second = ::std::min(dim_beg + part_size * (place_position.x + 1), dim_end);
return box(bounds);
}
template <typename mdspan_shape_t>
_CCCL_HOST_DEVICE static auto apply(const mdspan_shape_t& in, pos4 place_position, dim4 grid_dims)
{
constexpr size_t dimensions = mdspan_shape_t::rank();
::std::array<::std::pair<::std::ptrdiff_t, ::std::ptrdiff_t>, dimensions> bounds;
for (size_t d = 0; d < dimensions; d++)
{
bounds[d].first = 0;
bounds[d].second = in.extent(d);
}
// Last position in this dimension (excluded)
size_t target_dim = (which_dim == -1) ? dimensions - 1 : size_t(which_dim);
if (target_dim > dimensions - 1)
{
target_dim = dimensions - 1;
}
::std::ptrdiff_t dim_end = in.extent(target_dim);
// The last dimension is split across the different places
size_t nplaces = grid_dims.x;
size_t part_size = (in.extent(target_dim) + nplaces - 1) / nplaces;
bounds[target_dim].first = ::std::min((::std::ptrdiff_t) part_size * place_position.x, dim_end);
bounds[target_dim].second = ::std::min((::std::ptrdiff_t) part_size * (place_position.x + 1), dim_end);
return box<dimensions>(bounds);
}
_CCCL_HOST_DEVICE static void get_executor(pos4* result, pos4 data_coords, dim4 data_dims, dim4 grid_dims)
{
// Find the largest dimension
size_t rank = data_dims.get_rank();
size_t target_dim = (which_dim == -1) ? rank : size_t(which_dim);
if (target_dim > rank)
{
target_dim = rank;
}
size_t extent = data_dims.get(target_dim);
size_t nplaces = grid_dims.x;
_CCCL_ASSERT(nplaces > 0, "blocked partition requires a non-empty grid");
size_t part_size = (extent + nplaces - 1) / nplaces;
// A zero part_size (empty extent, or extent + nplaces - 1 wrapping) would
// make the division below SIGFPE; allocate_nd() rejects such geometries
// before the mapper runs
_CCCL_ASSERT(part_size > 0, "blocked partition applied to an empty or wrapping extent");
// Get the coordinate in the selected dimension
size_t c = data_coords.get(target_dim);
*result = pos4(c / part_size);
}
};
//! Partitions a multidimensional box or shape into contiguous blocks along a selected dimension.
//!
//! This partitioning strategy divides the data space into contiguous blocks, distributing them
//! across execution places. By default, partitioning occurs along the last dimension, but a
//! specific dimension can be selected using the template parameter. This approach ensures
//! good spatial locality and is particularly effective for regular data access patterns.
//!
//! When mapping element coordinates (get_executor), the selected dimension is
//! clamped to the highest axis whose extent is greater than one: -1 always
//! selects that axis, and a larger explicit dimension is clamped down to it
//! (e.g. blocked_partition_custom<2> on extents {n, 1, 1, 1} partitions along
//! axis 0).
using blocked_partition = blocked_partition_custom<>;
#ifdef UNITTESTED_FILE
UNITTEST("blocked partition with very large data arrays")
{
// blocked_partition splits along a single dimension (the highest-rank one
// by default) and maps to grid_dims.x places.
// 400 x 300 x 200 x 1000 = 24,000,000,000 elements (~24 billion)
dim4 massive_4d_dims(400, 300, 200, 1000);
EXPECT(massive_4d_dims.size() == 24000000000ULL);
EXPECT(massive_4d_dims.size() > (1ULL << 34));
// Partition the t-dimension (rank 3, extent 1000) into 4 blocks
dim4 grid_dims(4);
pos4 first_coord(0, 0, 0, 0);
pos4 middle_coord(200, 150, 100, 500);
pos4 last_coord(399, 299, 199, 999);
pos4 first_pos, middle_pos, last_pos;
blocked_partition::get_executor(&first_pos, first_coord, massive_4d_dims, grid_dims);
blocked_partition::get_executor(&middle_pos, middle_coord, massive_4d_dims, grid_dims);
blocked_partition::get_executor(&last_pos, last_coord, massive_4d_dims, grid_dims);
// part_size = ceil(1000/4) = 250
// t=0 -> block 0, t=500 -> block 2, t=999 -> block 3
EXPECT(first_pos.x == 0);
EXPECT(middle_pos.x == 2);
EXPECT(last_pos.x == 3);
};
#endif // UNITTESTED_FILE
} // namespace cuda::experimental::places

391
cccl_upstream/cudax/include/cuda/experimental/__places/partitions/cyclic_shape.cuh
View File

@@ -1,391 +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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief Define explicit shapes
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/experimental/__stf/utility/dimensions.cuh>
namespace cuda::experimental::places
{
using ::cuda::experimental::stf::box;
using ::cuda::experimental::stf::dim4;
using ::cuda::experimental::stf::each;
using ::cuda::experimental::stf::pos4;
/**
* @brief An cyclic shape is a shape or rank 'dimensions' where the bounds are
* explicit in each dimension, and where we jump between elements with a
* specific stride.
*
* @tparam dimensions the rank of the shape
*/
template <size_t dimensions = 1>
class cyclic_shape
{
public:
///@{ @name Constructors
/// Construct and explicit shape from a list of lower and upper bounds
_CCCL_HOST_DEVICE explicit cyclic_shape(const ::std::array<::std::tuple<size_t, size_t, size_t>, dimensions>& list)
{
size_t i = 0;
for (auto& e : list)
{
begins[i] = ::std::get<0>(e);
ends[i] = ::std::get<1>(e);
strides[i] = ::std::get<2>(e);
// printf("CYCLIC SHAPE : dim %ld : %ld %ld %ld\n", i, begin, end, stride);
i++;
}
}
/// Construct and explicit shape from a list of lower and upper bounds
_CCCL_HOST_DEVICE cyclic_shape(const ::std::array<size_t, dimensions>& begins_,
const ::std::array<size_t, dimensions>& ends_,
const ::std::array<size_t, dimensions>& strides_)
: begins(begins_)
, ends(ends_)
, strides(strides_)
{}
_CCCL_HOST_DEVICE void print() const
{
printf("CYCLIC SHAPE\n");
for (size_t i = 0; i < dimensions; i++)
{
printf("\t%ld:%ld:%ld\n", begins[i], ends[i], strides[i]);
}
}
///@}
/// Get the total number of elements in this explicit shape
_CCCL_HOST_DEVICE ::std::ptrdiff_t size() const
{
size_t res = 1;
for (size_t d = 0; d < dimensions; d++)
{
res *= (ends[d] - begins[d] + strides[d] - 1) / strides[d];
}
return res;
}
// Overload the equality operator to check if two shapes are equal
_CCCL_HOST_DEVICE bool operator==(const cyclic_shape& other) const
{
for (size_t i = 0; i < dimensions; ++i)
{
if (begins[i] != other.begins[i] || ends[i] != other.ends[i] || strides[i] != other.strides[i])
{
// printf("BEGIN[%ld] %ld != OTHER BEGIN[%ld] %ld\n", i, begins[i] , i , other.begins[i]);
return false;
}
}
return true;
}
/// get the dimensionnality of the explicit shape
_CCCL_HOST_DEVICE size_t get_rank() const
{
return dimensions;
}
// Iterator class for cyclic_shape
class iterator
{
private:
cyclic_shape* shape;
size_t index;
::std::array<size_t, dimensions> current; // Array to store the current position in each dimension
public:
_CCCL_HOST_DEVICE iterator(cyclic_shape& s, size_t idx = 0)
: shape(&s)
, index(idx)
, current(s.begins)
{}
// Overload the dereference operator to get the current position
_CCCL_HOST_DEVICE auto& operator*()
{
if constexpr (dimensions == 1)
{
return current[0];
}
else
{
return current;
}
}
// Overload the pre-increment operator to move to the next position
_CCCL_HOST_DEVICE iterator& operator++()
{
index++;
for (auto dim : each(0, dimensions))
{
current[dim] += shape->strides[dim]; // Increment the current position by the stride
// fprintf(stderr, "current[%ld] += %ld\n", dim, shape.get_stride(dim));
// printf("TEST current[%ld] (%ld) > shape.ends[%ld] (%ld)\n", dim, current[dim], dim, shape.ends[dim]);
if (current[dim] < shape->ends[dim])
{
break;
}
// Wrap around to the begin value if the current position exceeds the end value
current[dim] = shape->begins[dim];
}
return *this;
}
// Overload the equality operator to check if two iterators are equal
_CCCL_HOST_DEVICE bool operator==(const iterator& other) const
{ /*printf("EQUALITY TEST index %d %d shape equal ? %s\n", index,
other.index, (&shape == &other.shape)?"yes":"no"); */
// printf("check == : index %d %d => %s shape equal ? %s\n", index, other.index, (index ==
// other.index)?"yes":"no", (shape == other.shape)?"yes":"no");
return shape == other.shape && index == other.index;
}
// Overload the inequality operator to check if two iterators are not equal
_CCCL_HOST_DEVICE bool operator!=(const iterator& other) const
{
return !(*this == other);
}
};
// Functions to create the begin and end iterators
_CCCL_HOST_DEVICE iterator begin()
{
return iterator(*this);
}
_CCCL_HOST_DEVICE iterator end()
{
size_t total_positions = 1;
for (auto i : each(0, dimensions))
{
auto begin = begins[i]; // included
auto end = ends[i]; // excluded
auto stride = strides[i];
total_positions *= (end - begin + stride - 1) / stride;
}
// The first invalid position is index with "total_positions", even if there is no entry in the range
return iterator(*this, total_positions);
}
private:
::std::array<size_t, dimensions> begins;
::std::array<size_t, dimensions> ends;
::std::array<size_t, dimensions> strides;
};
/**
* @brief Apply a round-robin distribution of elements
*/
class cyclic_partition
{
public:
cyclic_partition() = default;
template <size_t dimensions>
_CCCL_HOST_DEVICE static auto apply(const box<dimensions>& in, pos4 place_position, dim4 grid_dims)
{
::std::array<size_t, dimensions> begins;
::std::array<size_t, dimensions> ends;
::std::array<size_t, dimensions> strides;
for (size_t d = 0; d < dimensions; d++)
{
begins[d] = in.get_begin(d) + place_position.get(d);
ends[d] = in.get_end(d);
strides[d] = grid_dims.get(d);
}
return cyclic_shape<dimensions>(begins, ends, strides);
}
template <typename mdspan_shape_t>
_CCCL_HOST_DEVICE static auto apply(const mdspan_shape_t& in, pos4 place_position, dim4 grid_dims)
{
constexpr size_t dimensions = mdspan_shape_t::rank();
::std::array<::std::tuple<size_t, size_t, size_t>, dimensions> bounds;
for (size_t d = 0; d < dimensions; d++)
{
// Can't assign the whole tuple because the assignment needs to run on device.
::std::get<0>(bounds[d]) = place_position.get(d);
::std::get<1>(bounds[d]) = in.extent(d);
::std::get<2>(bounds[d]) = grid_dims.get(d);
}
return cyclic_shape<dimensions>(bounds);
}
/**
* @brief Inverse of `apply` for zero-based data coordinates: in a round-robin
* distribution the owner of an element is its coordinate modulo the grid
* extent, independently in each dimension.
*/
_CCCL_HOST_DEVICE static void get_executor(pos4* result, pos4 data_coords, dim4 /*data_dims*/, dim4 grid_dims)
{
_CCCL_ASSERT(data_coords.x >= 0 && data_coords.y >= 0 && data_coords.z >= 0 && data_coords.t >= 0,
"get_executor requires zero-based (non-negative) coordinates");
_CCCL_ASSERT(grid_dims.x >= 1 && grid_dims.y >= 1 && grid_dims.z >= 1 && grid_dims.t >= 1,
"get_executor requires nonzero grid extents");
*result = pos4(data_coords.x % static_cast<::std::ptrdiff_t>(grid_dims.x),
data_coords.y % static_cast<::std::ptrdiff_t>(grid_dims.y),
data_coords.z % static_cast<::std::ptrdiff_t>(grid_dims.z),
data_coords.t % static_cast<::std::ptrdiff_t>(grid_dims.t));
}
};
#ifdef UNITTESTED_FILE
UNITTEST("cyclic_shape<3>")
{
// Expect to iterate over Card({0, 2, 4, 6}x{1, 3}x{10, 15}) = 4*2*2 = 16 items
const size_t expected_cnt = 16;
size_t cnt = 0;
cyclic_shape<3> shape{{::std::make_tuple(0, 7, 2), ::std::make_tuple(1, 5, 2), ::std::make_tuple(10, 20, 5)}};
for ([[maybe_unused]] const auto& pos : shape)
{
//// Use the position in each dimension
// ::std::cout << "(";
// for (const auto& p: pos) {
// ::std::cout << p << ", ";
//}
// ::std::cout << ")" << ::'\n';
EXPECT(cnt < expected_cnt);
cnt++;
}
EXPECT(cnt == expected_cnt);
};
UNITTEST("cyclic_shape<1> size rounds up stragglers")
{
cyclic_shape<1> shape{{::std::make_tuple(0, 7, 2)}};
size_t cnt = 0;
for ([[maybe_unused]] const auto& pos : shape)
{
EXPECT(cnt < 4);
cnt++;
}
EXPECT(cnt == 4);
EXPECT(shape.size() == 4);
};
UNITTEST("empty cyclic_shape<1>")
{
cyclic_shape<1> shape{{::std::make_tuple(7, 7, 1)}};
auto it_end = shape.end();
auto it_begin = shape.begin();
if (it_end != it_begin)
{
fprintf(stderr, "Error: begin() != end()\n");
abort();
}
// There should be no entry in this range
for ([[maybe_unused]] const auto& pos : shape)
{
abort();
}
};
UNITTEST("apply cyclic ")
{
box<3> e({{0, 7}, {1, 5}, {10, 20}});
size_t dim0 = 2;
size_t dim1 = 2;
size_t cnt = 0;
size_t expected_cnt = 7 * 4 * 10;
for (size_t i0 = 0; i0 < dim0; i0++)
{
for (size_t i1 = 0; i1 < dim1; i1++)
{
auto c = cyclic_partition::apply(e, pos4(i0, i1), dim4(dim0, dim1));
for ([[maybe_unused]] const auto& pos : c)
{
//// Use the position in each dimension
// ::std::cout << " (";
// for (const auto& p: pos) {
// ::std::cout << p << ", ";
//}
// ::std::cout << ")" << ::'\n';
// avoid infinite loops
EXPECT(cnt < expected_cnt);
cnt++;
}
}
}
// We must have gone over all elements exactly once
EXPECT(cnt == expected_cnt);
};
UNITTEST("cyclic get_executor is the inverse of apply")
{
// Zero-based box so that apply and get_executor use the same coordinate
// origin: apply assigns coordinate c to place (c % grid extent) per dimension.
box<3> e({{0, 7}, {0, 5}, {0, 20}});
const size_t dim0 = 2;
const size_t dim1 = 3;
const dim4 grid_dims(dim0, dim1);
const dim4 data_dims(7, 5, 20);
size_t cnt = 0;
for (size_t i0 = 0; i0 < dim0; i0++)
{
for (size_t i1 = 0; i1 < dim1; i1++)
{
auto c = cyclic_partition::apply(e, pos4(i0, i1), grid_dims);
for (const auto& pos : c)
{
pos4 owner;
cyclic_partition::get_executor(&owner, pos4(pos[0], pos[1], pos[2]), data_dims, grid_dims);
EXPECT(owner == pos4(i0, i1));
cnt++;
}
}
}
// Every element must be visited exactly once and map back to its place
EXPECT(cnt == 7 * 5 * 20);
};
#endif // UNITTESTED_FILE
} // namespace cuda::experimental::places

188
cccl_upstream/cudax/include/cuda/experimental/__places/partitions/tiled_partition.cuh
View File

@@ -1,188 +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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief Definition of the `tiled_partition` strategy
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/experimental/__places/places.cuh>
namespace cuda::experimental::places
{
namespace reserved
{
/*
* Define a tiled transformation to a shape of mdspan
*/
template <size_t tile_size, typename mdspan_shape_t>
class tiled_mdspan_shape
{
public:
// parts id, nparts ?
tiled_mdspan_shape(const mdspan_shape_t& s, size_t part_id, size_t nparts)
: original_shape(s)
, part_id(part_id)
, nparts(nparts)
{}
/* The number of elements in this part */
size_t size() const
{
_CCCL_ASSERT(mdspan_shape_t::rank() == 1, "Tiled mdspan shape only implemented in 1D yet");
const size_t n = original_shape.size();
// 0000 1111 2222 0000 11xx xxxx xxxx
// S=4, n=18
//
// nparts=3
// ntiles = (18+3)/4 = 5
// tile_per_part = (5+2)/3 = 2 (1 or 2 tiles per part)
// cnt = 2*4 = 8
// last_elem = {0 => (0+((2-1)*3)+1)*8=(1+3)*8=16} (n >= last_elem) => cnt = 8
// {1 => (1+((2-1)*3)+1)*8=(2+3)*8=20} extra = min(4, 20-18)=2 => cnt = 6
// {2 => (2+((2-1)*3)+1)*8=(3+3)*8=24} extra = min(4, 24-18)=4 => cnt = 4
// How many tiles if we round up to a multiple of tile_size ?
const size_t ntiles = (n + tile_size - 1) / tile_size;
const size_t tile_per_part = (ntiles + nparts - 1) / nparts;
// If all parts are the same
size_t cnt = tile_per_part * tile_size;
// Assuming all parts have tile_per_part tiles, the last tile of the
// part starts the last tile has index (part_id + ((tile_per_part-1)*
// nparts)), and ends at the beginning of the next tile (hence +1).
const size_t last_elem = (part_id + (tile_per_part - 1) * nparts + 1) * tile_size;
// Remove extra elements (if any) by computing what would be the last
// element for this part
if (last_elem > n)
{
size_t extra_elems = ::std::min(tile_size, last_elem - n);
cnt -= extra_elems;
}
return cnt;
}
using coords_t = typename mdspan_shape_t::coords_t;
_CCCL_HOST_DEVICE coords_t index_to_coords(size_t index) const
{
// First transform from nparts and part_id to one coordinate
const size_t remain = index % tile_size;
const size_t tile_id = index / tile_size;
// Stage 2: apply original shape's transformation to stage1
return original_shape.index_to_coords((part_id + tile_id * nparts) * tile_size + remain);
}
private:
mdspan_shape_t original_shape;
size_t part_id;
size_t nparts;
};
} // end namespace reserved
/**
* @brief Tiled partition strategy applied on a shape of mdspan
*
* @tparam `tile_size` size of the tiles
* @tparam `mdspan_shape_t` shape of a mdspan
*
* Since there is no partial template deduction on classes, we provide a
* function to implement tiled<tile_size> and have the other type deduced.
*/
template <size_t tile_size, typename mdspan_shape_t>
auto tiled(const mdspan_shape_t s, size_t part_id, size_t nparts)
{
return reserved::tiled_mdspan_shape<tile_size, mdspan_shape_t>(s, part_id, nparts);
}
template <size_t tile_size>
class tiled_partition
{
static_assert(tile_size > 0);
public:
tiled_partition() = default;
template <typename mdspan_shape_t>
static const reserved::tiled_mdspan_shape<tile_size, mdspan_shape_t>
apply(const mdspan_shape_t& in, pos4 place_position, dim4 grid_dims)
{
// TODO assert 1D !
assert(grid_dims.x > 0);
return reserved::tiled_mdspan_shape<tile_size, mdspan_shape_t>(in, place_position.x, grid_dims.x);
}
_CCCL_HOST_DEVICE static void get_executor(pos4* result, pos4 data_coords, dim4 /*unused*/, dim4 grid_dims)
{
assert(grid_dims.x > 0);
*result = pos4((data_coords.x / tile_size) % grid_dims.x);
}
};
#ifdef UNITTESTED_FILE
UNITTEST("Composite data place equality")
{
auto all = exec_place::all_devices();
auto repeated_dev0 = exec_place::repeat(exec_place::device(0), 3);
using P = tiled_partition<128>;
using P2 = tiled_partition<64>;
/* Same partitioning operator, same execution place */
EXPECT(data_place::composite(P(), all) == data_place::composite(P(), all));
/* Make sure we do not have a false positive in the test below */
EXPECT(all != repeated_dev0);
/* Same partitioning operator, different execution place */
EXPECT(data_place::composite(P(), repeated_dev0) != data_place::composite(P(), all));
/* Different partitioning operator, same execution place */
EXPECT(data_place::composite(P(), all) != data_place::composite(P2(), all));
};
UNITTEST("tiled partition with large 1D data")
{
const size_t large_1d_size = (1ULL << 36);
// Test coordinate that makes tiling calculation obvious
const size_t test_coord = (1ULL << 34);
pos4 large_coords(test_coord); // 1D coordinate
dim4 data_dims(large_1d_size); // 1D data space
dim4 grid_dims(32); // 32 places in grid
constexpr size_t tile_size = 1000;
pos4 tile_pos;
tiled_partition<tile_size>::get_executor(&tile_pos, large_coords, data_dims, grid_dims);
EXPECT(tile_pos.x == (test_coord / tile_size) % grid_dims.x);
};
#endif // UNITTESTED_FILE
} // namespace cuda::experimental::places

329
cccl_upstream/cudax/include/cuda/experimental/__places/place_partition.cuh
View File

@@ -1,329 +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-2024 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief Facilities to manipulate subset of places
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cuda/experimental/__places/exec/cuda_stream.cuh>
#include <cuda/experimental/__places/exec/green_context.cuh>
#include <cuda/experimental/__places/places.cuh>
#include <cuda/experimental/__stf/internal/async_resources_handle.cuh>
namespace cuda::experimental::places
{
/**
* @brief Defines a partitioning granularity
*
* This should be used in combination with `place_partition`
*/
enum class place_partition_scope
{
cuda_device,
green_context,
cuda_stream,
};
/**
* @brief Convert a place_partition_scope value to a string (for debugging purpose)
* @param scope The partitioning granularity to convert
* @return A string representation of `scope` (e.g. "cuda_device", "green_context", "cuda_stream")
*/
inline ::std::string place_partition_scope_to_string(place_partition_scope scope)
{
switch (scope)
{
case place_partition_scope::cuda_device:
return "cuda_device";
case place_partition_scope::green_context:
return "green_context";
case place_partition_scope::cuda_stream:
return "cuda_stream";
}
abort();
return "unknown";
}
// TODO method to get the scope of an exec place
/**
* @brief Get subsets of an execution place.
*
* Computes a vector of execution places that partition the input place at a
* given granularity (see `place_partition_scope`). For example, a grid place
* can be partitioned into devices, or into green contexts, or into CUDA streams.
*
* Use the constructors that take `::cuda::experimental::stf::async_resources_handle&` when partitioning at
* `cuda_stream` or `green_context` scope (stream and green-context resources
* are obtained from the handle). The constructors without a handle support only
* `cuda_device` scope. Green context scope requires CUDA 12.4 or later.
*
* Iteration over subplaces is provided via `begin()` / `end()`; `to_exec_place()` builds
* an `exec_place` grid from the subplaces.
*/
class place_partition
{
public:
/** @brief Partition an execution place into a vector of subplaces (with async resource handle).
* @param place The execution place to partition (e.g. grid or device)
* @param handle Handle used to obtain stream or green-context resources when scope is cuda_stream or green_context
* @param scope Partitioning granularity (cuda_device, green_context, or cuda_stream)
*/
place_partition(
exec_place place, ::cuda::experimental::stf::async_resources_handle& handle, place_partition_scope scope)
{
#if _CCCL_CTK_BELOW(12, 4)
_CCCL_ASSERT(scope != place_partition_scope::green_context, "Green contexts unsupported.");
#endif // _CCCL_CTK_BELOW(12, 4)
compute_subplaces(handle, mv(place), scope);
}
/** @brief Partition an execution place into a vector of subplaces (no async handle).
* Only `cuda_device` scope is supported; green_context and cuda_stream require a handle.
* @param place The execution place to partition
* @param scope Partitioning granularity (must be cuda_device when no handle is provided)
*/
place_partition(const exec_place& place, place_partition_scope scope)
{
#if _CCCL_CTK_BELOW(12, 4)
_CCCL_ASSERT(scope != place_partition_scope::green_context, "Green contexts need an async resource handle.");
#endif // _CCCL_CTK_BELOW(12, 4)
compute_subplaces_no_handle(place, scope);
}
/** @brief Partition a vector of execution places into a single vector of subplaces (with async handle).
* @param handle Handle for stream or green-context resources when scope is cuda_stream or green_context
* @param places Input execution places to partition
* @param scope Partitioning granularity
*/
place_partition(::cuda::experimental::stf::async_resources_handle& handle,
const ::std::vector<::std::shared_ptr<exec_place>>& places,
place_partition_scope scope)
{
for (const auto& place : places)
{
compute_subplaces(handle, *place, scope);
}
}
/** @brief Partition a grid of execution places into a single vector of subplaces (with async handle).
* @param handle Handle for stream or green-context resources when scope is cuda_stream or green_context
* @param grid Input execution place grid to partition
* @param scope Partitioning granularity
*/
place_partition(
::cuda::experimental::stf::async_resources_handle& handle, const exec_place& grid, place_partition_scope scope)
{
::std::vector<::std::shared_ptr<exec_place>> places;
places.reserve(grid.size());
for (size_t i = 0; i < grid.size(); ++i)
{
places.push_back(::std::make_shared<exec_place>(grid.get_place(i)));
}
for (const auto& place : places)
{
compute_subplaces(handle, *place, scope);
}
}
/** @brief Partition a vector of execution places into a single vector of subplaces (no async handle).
* Only cuda_device scope is supported.
* @param places Input execution places to partition
* @param scope Partitioning granularity (must be cuda_device)
*/
place_partition(const ::std::vector<::std::shared_ptr<exec_place>>& places, place_partition_scope scope)
{
for (const auto& place : places)
{
compute_subplaces_no_handle(*place, scope);
}
}
~place_partition() = default;
/** Iteration over subplaces. */
using iterator = ::std::vector<exec_place>::iterator;
using const_iterator = ::std::vector<exec_place>::const_iterator;
/** @brief Iterator to the first subplace. @return Begin iterator. */
iterator begin()
{
return sub_places.begin();
}
/** @brief Past-the-end iterator for subplaces. @return End iterator. */
iterator end()
{
return sub_places.end();
}
/** @brief Const iterator to the first subplace. @return Begin const iterator. */
const_iterator begin() const
{
return sub_places.begin();
}
/** @brief Past-the-end const iterator. @return End const iterator. */
const_iterator end() const
{
return sub_places.end();
}
/** @brief Number of subplaces in the partition. @return Size of the partition. */
size_t size() const
{
return sub_places.size();
}
/** @brief Get the i-th subplace (mutable).
* @param i Index in [0, size()).
* @return Reference to the i-th exec_place.
*/
exec_place& get(size_t i)
{
return sub_places[i];
}
/** @brief Get the i-th subplace (const).
* @param i Index in [0, size()).
* @return Const reference to the i-th exec_place.
*/
const exec_place& get(size_t i) const
{
return sub_places[i];
}
/** @brief Build an exec_place from the subplaces.
* @return A grid view of the partitioned execution places, or single place if size == 1.
*/
exec_place to_exec_place() const
{
return make_grid(sub_places);
}
private:
/** @brief Compute the subplaces of a place at the specified granularity (scope) into the sub_places vector */
void compute_subplaces(
::cuda::experimental::stf::async_resources_handle& handle, exec_place place, place_partition_scope scope)
{
// Handle multi-element grids by recursively partitioning
if (place.size() > 1 && scope == place_partition_scope::cuda_stream)
{
for (auto& device_p : place_partition(mv(place), handle, place_partition_scope::cuda_device))
{
auto device_p_places = place_partition(device_p, handle, place_partition_scope::cuda_stream).sub_places;
sub_places.insert(sub_places.end(), device_p_places.begin(), device_p_places.end());
}
return;
}
// Handle scalar places (including 1-element grids) for cuda_stream scope
if (place.size() == 1 && scope == place_partition_scope::cuda_stream)
{
// Get the underlying scalar place (for 1-element grids, get the single element)
exec_place scalar_place = place.is_device() ? place : place.get_place(0);
if (!scalar_place.is_device())
{
// Host or other non-device place - no streams to partition into
sub_places.push_back(mv(place));
return;
}
auto& pool = scalar_place.get_stream_pool(true, handle.get_place_resources());
for (size_t i = 0; i < pool.size(); i++)
{
sub_places.push_back(exec_place::cuda_stream(pool.next(scalar_place)));
}
return;
}
// Green contexts are only supported since CUDA 12.4
#if _CCCL_CTK_AT_LEAST(12, 4)
if (place.size() > 1 && scope == place_partition_scope::green_context)
{
// Recursively partition grid into devices, then into green contexts
for (auto& device_p : place_partition(mv(place), handle, place_partition_scope::cuda_device))
{
auto device_p_places = place_partition(device_p, handle, place_partition_scope::green_context).sub_places;
sub_places.insert(sub_places.end(), device_p_places.begin(), device_p_places.end());
}
return;
}
// Handle scalar places (including 1-element grids) for green_context scope
if (place.size() == 1 && scope == place_partition_scope::green_context)
{
exec_place scalar_place = place.is_device() ? place : place.get_place(0);
if (!scalar_place.is_device())
{
sub_places.push_back(mv(place));
return;
}
int dev_id = device_ordinal(scalar_place.affine_data_place());
const char* env = getenv("CUDASTF_GREEN_CONTEXT_SIZE");
int sm_cnt = env ? atoi(env) : 8;
auto h = handle.get_gc_helper(dev_id, sm_cnt);
size_t cnt = h->get_count();
for (size_t i = 0; i < cnt; i++)
{
sub_places.push_back(exec_place::green_ctx(h->get_view(i)));
}
return;
}
#endif
// If the scope requires no handle
compute_subplaces_no_handle(place, scope);
}
/** @brief Compute the subplaces of a place at the specified granularity (scope) into the sub_places vector */
void compute_subplaces_no_handle(const exec_place& place, place_partition_scope scope)
{
#if _CCCL_CTK_BELOW(12, 4)
_CCCL_ASSERT(scope != place_partition_scope::green_context, "Green contexts scope need an async resource handle.");
#endif // _CCCL_CTK_BELOW(12, 4)
_CCCL_ASSERT(scope != place_partition_scope::cuda_stream, "CUDA stream scope needs an async resource handle.");
if (scope == place_partition_scope::cuda_device)
{
for (size_t i = 0; i < place.size(); ++i)
{
sub_places.push_back(place.get_place(i));
}
return;
}
assert(!"Internal error: unreachable code.");
}
/** A vector with all subplaces (computed once in compute_subplaces) */
::std::vector<exec_place> sub_places;
};
// Deferred implementation because we need place_partition
template <typename... Args>
auto exec_place::partition_by_scope(Args&&... args)
{
return place_partition(*this, ::std::forward<Args>(args)...).to_exec_place();
}
} // end namespace cuda::experimental::places

2302
cccl_upstream/cudax/include/cuda/experimental/__places/places.cuh
View File

File diff suppressed because it is too large Load Diff

290
cccl_upstream/cudax/include/cuda/experimental/__places/stream_pool.cuh
View File

@@ -1,290 +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.
//
//===----------------------------------------------------------------------===//
/**
* @file
* @brief Stream pool and augmented stream types used by places.
*
* Definitions of stream_pool::next() live in places.cuh (because we need to
* activate the place to create a stream in the appropriate CUDA context).
*/
#pragma once
#include <cuda/__cccl_config>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif
#include <cuda/experimental/__stf/utility/cuda_safe_call.cuh>
#include <memory>
#include <mutex>
#include <vector>
#include <cuda.h>
#include <cuda_runtime.h>
namespace cuda::experimental::places
{
using ::cuda::experimental::stf::cuda_try;
using ::cuda::experimental::stf::mv;
class exec_place;
/**
* @brief Is @p stream currently participating in a CUDA graph capture?
*
* Returns `false` for `nullptr` (the legacy default stream is never
* capturing). `cudaStreamIsCapturing` is itself capture-safe, so this can be
* called from contexts where most other driver queries (`cuStreamGetId`,
* `cudaStreamGetDevice`, ...) would be rejected with
* `cudaErrorStreamCaptureUnsupported` and would *invalidate* the in-flight
* capture.
*/
[[nodiscard]] inline bool is_stream_capturing(cudaStream_t stream)
{
return cuda_try<cudaStreamIsCapturing>(stream) != cudaStreamCaptureStatusNone;
}
/**
* @brief Computes the CUDA device in which the stream was created
*/
inline int get_device_from_stream(cudaStream_t stream)
{
if (stream == nullptr)
{
return cuda_try<cudaGetDevice>();
}
// cudaStreamGetDevice/cuStreamGetCtx are not permitted while the stream is
// participating in capture. Use the active device, which is the device on
// which the capture is being constructed.
if (is_stream_capturing(stream))
{
return cuda_try<cudaGetDevice>();
}
#if _CCCL_CTK_AT_LEAST(12, 8)
return cuda_try<cudaStreamGetDevice>(stream);
#else
auto stream_driver = CUstream(stream);
CUcontext ctx = cuda_try<cuStreamGetCtx>(stream_driver);
cuda_try(cuCtxPushCurrent(ctx));
CUdevice stream_dev = cuda_try<cuCtxGetDevice>();
static_cast<void>(cuda_try<cuCtxPopCurrent>());
return static_cast<int>(stream_dev);
#endif
}
/** Sentinel for "no stream" / empty slot. Distinct from any value returned by cuStreamGetId. */
inline constexpr unsigned long long k_no_stream_id = static_cast<unsigned long long>(-1);
/**
* @brief Returns the unique stream ID from the CUDA driver (cuStreamGetId).
* @param stream A valid CUDA stream, or nullptr.
* @return The stream's unique ID, or k_no_stream_id if stream is nullptr.
*
* When @p stream is participating in CUDA graph capture, querying the driver
* stream ID is not permitted (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). In that case
* this returns k_no_stream_id; callers that cache syncs treat unknown ids as
* ``never skip cudaStreamWaitEvent`` (see async_resources_handle).
*/
inline unsigned long long get_stream_id(cudaStream_t stream)
{
if (stream == nullptr)
{
return k_no_stream_id;
}
// ``cuStreamGetId`` is not capture-safe: during
// ``cudaStreamCaptureModeThreadLocal`` / ``Global`` it rejects the query
// *and* invalidates the capture itself. Gate on ``cudaStreamIsCapturing``
// (which is safe) and conservatively report an unknown stream ID while
// capture is in flight.
if (is_stream_capturing(stream))
{
return k_no_stream_id;
}
unsigned long long id = cuda_try<cuStreamGetId>(reinterpret_cast<CUstream>(stream));
_CCCL_ASSERT(id != k_no_stream_id, "Internal error: cuStreamGetId returned k_no_stream_id");
return id;
}
/**
* @brief A CUDA stream augmented with its pre-resolved driver identity and home device.
*
* Carries:
* - the stream itself,
* - the stream's unique ID from the CUDA driver (cuStreamGetId), or k_no_stream_id when unknown,
* - the device index on which the stream resides.
*
* The id and device are looked up at construction so downstream consumers
* (the stream-reuse logic in stream_task, the (src_id, dst_id)-keyed sync-skip
* cache in async_resources_handle, etc.) never have to re-pay the driver
* round-trip and can use the augmented stream as a cache key.
*/
struct augmented_stream
{
augmented_stream() = default;
augmented_stream(cudaStream_t stream, unsigned long long id, int dev_id = -1)
: stream(stream)
, id(id)
, dev_id(dev_id)
{}
/** Construct from stream only; id is from cuStreamGetId, dev_id is -1 (filled lazily when needed). */
explicit augmented_stream(cudaStream_t stream)
: stream(stream)
, id(get_stream_id(stream))
, dev_id(-1)
{}
cudaStream_t stream = nullptr;
unsigned long long id = k_no_stream_id;
int dev_id = -1;
};
/**
* @brief A stream_pool object stores a set of streams associated to a specific
* CUDA context (device, green context, ...)
*
* This class uses a PIMPL idiom so that it is copyable and movable with shared
* semantics: copies refer to the same underlying pool of streams.
*
* When a slot is empty, next(place) activates the place (RAII guard) and calls
* place.create_stream(). Defined in places.cuh.
*/
class stream_pool
{
struct impl
{
explicit impl(size_t n)
: payload(n, augmented_stream(nullptr, k_no_stream_id, -1))
{}
// Construct from an augmented stream, this is used to create a stream pool with a single stream.
explicit impl(augmented_stream ds)
: payload(1, mv(ds))
, externally_owned(true)
{}
// Release every stream the pool has lazily created. We intentionally
// skip entries that came from an externally-owned `decorated_stream`
// (single-stream pool built from a user-supplied stream, used for
// `exec_place::cuda_stream(s)`); those are not ours to destroy.
//
// `cudaStreamDestroy` is documented to be asynchronous when work is
// still pending on the stream: the call returns immediately and CUDA
// releases the stream's resources once the device has completed its
// pending work. That contract is what makes it safe to tear the pool
// down at the end of an STF context without blocking on a caller stream
// synchronize, as long as the outbound event chain has already been
// recorded back onto the user stream.
~impl() noexcept
{
if (externally_owned)
{
return;
}
for (auto& ds : payload)
{
if (ds.stream != nullptr)
{
// Stream destruction can fail during CUDA runtime teardown; the
// destructor has no useful way to report or recover from that.
(void) cudaStreamDestroy(ds.stream);
ds.stream = nullptr;
}
}
}
impl(const impl&) = delete;
impl& operator=(const impl&) = delete;
mutable ::std::mutex mtx;
::std::vector<augmented_stream> payload;
size_t index = 0;
bool externally_owned = false;
};
::std::shared_ptr<impl> pimpl;
public:
stream_pool() = default;
explicit stream_pool(size_t n)
: pimpl(::std::make_shared<impl>(n))
{}
// Construct from an augmented stream, this is used to create a stream pool with a single stream.
explicit stream_pool(augmented_stream ds)
: pimpl(::std::make_shared<impl>(mv(ds)))
{}
stream_pool(const stream_pool&) = default;
stream_pool(stream_pool&&) = default;
stream_pool& operator=(const stream_pool&) = default;
stream_pool& operator=(stream_pool&&) = default;
/**
* @brief Get the next stream in the pool; when a slot is empty, activate the place (RAII guard) and call
* place.create_stream(). Defined in places.cuh so the pool can use exec_place_scope and exec_place::create_stream().
*/
augmented_stream next(const exec_place& place);
using iterator = ::std::vector<augmented_stream>::iterator;
iterator begin()
{
return pimpl->payload.begin();
}
iterator end()
{
return pimpl->payload.end();
}
/**
* @brief Number of streams in the pool
*
* CUDA streams are initialized lazily, so this gives the number of slots
* available in the pool, not the number of streams initialized.
*/
size_t size() const
{
::std::scoped_lock locker(pimpl->mtx);
return pimpl->payload.size();
}
explicit operator bool() const
{
return pimpl != nullptr;
}
bool operator==(const stream_pool& other) const
{
return pimpl == other.pimpl;
}
bool operator<(const stream_pool& other) const
{
return pimpl < other.pimpl;
}
};
} // namespace cuda::experimental::places