NCCL 2.27.3-1
Symmetric memory API and symmetric kernels
* Redesign from the ground up, enabling major latency and bandwidth
improvements.
* Add new API calls to register user-allocated memory among communicator
ranks into a NCCL window: ncclCommWindowRegister() and
ncclCommWindowDeregister(). The calls currently support symmetric
registration for P2P and NVLS, and require VMM memory buffers (i.e.,
CUMEM must be operational).
* Implement specialized kernels taking advantage of symmetrically
registered memory, with performance gains expected particularly for
small to medium message sizes.
* The kernels support 32 bit floating point types and smaller, and sum as
the reduction operator, with no more than one collective operation per
group.
* Floating point summation is always done in fp32 accumulators (with the
exception of fp8 on NVLS, where it uses fp16 inside the switch). Thus,
the accuracy with fp8 and fp16 data types should be much improved.
* This initial implementation supports non-network communicators only (P2P
and NVLS transports).
* To explore this functionality users need to use the new memory
registration API calls with the NCCL_WIN_COLL_SYMMETRIC flag and all
ranks of a communicator must pass buffers at the same offset in the same
registration when invoking a collective NCCL operation.
Add support for DGX Spark.
Add support for DirectNIC (CX8) to the internal IB plugin.
Add a new ncclCommShrink() API call
* It is a non-collective call similar to ncclCommSplit(), which makes it
possible to exclude some (possibly unresponsive) ranks from the parent
communicator.
Add support for loading multiple network plugins
* This enables the creation of generic containers that can work across a
range of providers.
* Allow NCCL_NET_PLUGIN to accept a comma-separated list of plugins to
load.
NVLink SHARP (NVLS) improvements
* Implement NVLS+IB SHARP support for AllGather and ReduceScatter with
user buffer registration. This improves performance and reduces the
number of CTAs needed to achieve peak bandwidth.
* Gracefully fall back by default to other transports if NVLS
initialization fails (the old behavior of returning an error code from a
NCCL call can be preserved by setting NCCL_NVLS_ENABLE=1).
* Decrease the NVLS channel count to 24 on Blackwell systems with multiple
NVLink domains per communicator.
* Enable fine-tuning of NCCL behavior per communicator using new
"ncclConfig_t" members "collnetEnable", "CTAPolicy", and "nvlsCTAs".
Profiler improvements
* Extend the init function by adding communicator name, comm id (hash),
rank, number of ranks, number of nodes, and the NCCL log function to the
argument list. This makes the name and the comm id available to all
events in the communicator without explicitly passing them to each
individual event. Add the communicator id and rank to the profiler trace
filename. Now, the communicator name can be set via a new "ncclConfig_t"
member "commName".
* Improve the accuracy of the GPU kernel events by providing GPU-generated
timestamps for the start and stop of every NCCL operation.
* Harmonize proxy events, removing overlaps between ProxyOp and ProxyStep
states.
* Add support for network-defined event updates (through
"recordEventState").
* Report the correct number of channels used by every collective/p2p
operation (used to be set to nMaxChannels for collectives and absent for
p2ps).
* Fix the logic on proxyCtrl Idle/Active events (Issue #1162).
* Fix an issue where the network proxy profiler could lose track of an
event identifier (Issue #1682).
* Improve the backward compatibility with plugins older than v4.
* Ensure that the work counters are 0-initialized.
* Fix a potential race condition in the network profiler that could result
in an event being linked to a wrong parent.
MNNVL improvements
* Increase to 16 the number of NICs used to communicate between MNNVL
domains on GB200 systems, to optimize the performance of collective
operations.
* Add support for more complex MNNVL topologies with up to 32 NICs per
node.
* If the MNNVL fabric initialization was unsuccessful, NCCL will now fail
by default, so as to avoid inadvertently falling back to a potentially
much slower network transport. Such failures are typically due to a
misconfigured IMEX support on the system. To continue without MNNVL,
restart the job with NCCL_MNNVL_ENABLE=0.
* Fix a potential hang in alltoall-like communication patterns at a scale
of over 80 ranks.
* Make NCCL_P2P_DISABLE=1 imply NCCL_MNNVL_ENABLE=0 (so the latter no
longer needs to be specified on MNNVL systems).
* Fix an initialization failure when NCCL_TOPO_FILE is used on MNNVL
systems.
* Fix the graph search to exclude non-local NICs.
* Fix the SHM transport to use fabric handles on MNNVL systems.
NIC Fusion improvements
* Disable the creation of fused NICs for physical devices that haven't
been merged.
* Flatten multiple ports to a single PCI device within the internal IB
plugin and reparent dual-port NICs under the first PCI parent. If the
parent is not a PCI switch, PCI devices for fused NICs won't be
duplicated.
* Route traffic on GB200-CX8 systems through DirectNIC, not the host
interface.
Improve support for platforms with C2C connectivity (e.g., GB200)
* Enable GPUDirect RDMA for the NICs by default.
* Add support for P2C (PXN over C2C) and the LL128 protocol.
Extend NCCL fault tolerance in multithreaded scenarios
* Support the creation of multiple nonblocking communicators within a
single group and polling in parallel for the completion using multiple
threads (one per communicator).
Enable ncclImplicitOrderLaunch for CUDA 12.9+
* This can potentially speed up NCCL_IMPLICIT_LAUNCH_ORDER.
Improve the netSocket transport latency and control
* Provide finer control over the size of the socket send/receive buffers,
the task size, and the number of sockets that a single peer can open.
* Add support for the inlining of small messages behind the header when
using multiple sockets per connection.
Improve the readability of the CPU affinity in the debug output
* Print it as a range string rather than a bitmask.
Fix a potential race condition in graph execution
* A contention could arise when mixing graph and non-graph execution.
Improve PXN connection code
* Avoid duplicate and unused connections.
RAS fixes
* Fix a memory corruption at job termination time in case of a previously
failed initialization of a RAS socket connection.
* Fix a race condition leading to a crash when generating a RAS report
during communicator initialization (Issues #1669, #1718).
* Fix a potential race condition when gathering data for a RAS status
report.
Fix a potential memory corruption in ncclCommSplit()
* Memory could get corrupted when resource sharing was in use and the size
of the NVLink domain in the new communicator was smaller than in the old
one.
Fix asynchronous graph upload
* Fix a small memory leak.
* Fix oversychronization.
Add a check for out-of-memory conditions in ncclMemAlloc()
Clean up the NCCL socket code
* accept() will retry also if just reading the magic failed (Issue #1613).
* connect() will retry also if poll() did not return a POLLOUT event
(Issue #1618).
* Add error checking in a few instances (Issue #1539).
* Fix the loop condition in ncclFindInterfaceMatchSubnet() (Issue #1574).
* Clean up the debug output, downgrading WARN messages to INFO in
non-critical cases, and printing the peer's address where relevant.
Switch NCCL_DEBUG_FILE to line buffering
* This should help avoid mixed-up partial output lines in multithreaded
cases.
Other minor fixes
* Improve the checks for buffer overflows in the graph code (Issue #1585).
* Extend logging and state clearing to all four events in the internal IB
plugin (Issue #1650).
* Fix the error path in case IB communication is not ready (Issue #1489).
* Add ECE logging for IB fabric.
* Fix various minor issues in the graph module (Issue #1635).
* Clean up the debug output in the graph code, downgrading WARN messages
to INFO in non-critical cases.
* Add a missing argument to a directSend() call (Issue #1628).
* Remove duplicate code in sendProxySetup() (Issue #1420).
* Fix the order of arguments of cudaDeviceCanAccessPeer() (Issue #1507).
* Fix compiler warnings with GCC 14.
* Fix a typo in a comment (Issue #1236).
[ROCm/rccl commit: 72d2432094]
This commit is contained in:
Kamil Iskra
@@ -38,18 +38,18 @@ struct IsFloatingPoint<double>: std::true_type {};
|
||||
// 3. Have constructor taking `uint64_t opArg`.
|
||||
|
||||
template<typename T>
|
||||
struct FuncCopy { using EltType = T; __device__ FuncCopy(uint64_t opArg=0) {}; };
|
||||
struct FuncCopy { using EltType = T; __device__ __forceinline__ FuncCopy(uint64_t opArg=0) {}; };
|
||||
template<typename T>
|
||||
struct FuncSum { using EltType = T; __device__ FuncSum(uint64_t opArg=0) {}; };
|
||||
struct FuncSum { using EltType = T; __device__ __forceinline__ FuncSum(uint64_t opArg=0) {}; };
|
||||
template<typename T>
|
||||
struct FuncProd { using EltType = T; __device__ FuncProd(uint64_t opArg=0) {}; };
|
||||
struct FuncProd { using EltType = T; __device__ __forceinline__ FuncProd(uint64_t opArg=0) {}; };
|
||||
|
||||
template<typename T>
|
||||
struct FuncMinMax {
|
||||
using EltType = T;
|
||||
BytePack<sizeof(T)> xormask; // only used by integers
|
||||
bool isMinNotMax; // only used by floats
|
||||
__device__ FuncMinMax(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncMinMax(uint64_t opArg=0) {
|
||||
xormask.native = opArg;
|
||||
isMinNotMax = (opArg&1)==0;
|
||||
}
|
||||
@@ -64,13 +64,13 @@ template<typename T> struct FuncSumPostDiv;
|
||||
template<typename Fn>
|
||||
struct RedOpArg { // default case: no argument
|
||||
static constexpr bool ArgUsed = false;
|
||||
__device__ static uint64_t loadArg(void *ptr) { return 0; }
|
||||
__device__ __forceinline__ static uint64_t loadArg(void *ptr) { return 0; }
|
||||
};
|
||||
|
||||
template<typename T>
|
||||
struct RedOpArg<FuncMinMax<T>> {
|
||||
static constexpr bool ArgUsed = true;
|
||||
__device__ static uint64_t loadArg(void *ptr) {
|
||||
__device__ __forceinline__ static uint64_t loadArg(void *ptr) {
|
||||
union { uint64_t u64; T val; };
|
||||
u64 = 0;
|
||||
val = *(T*)ptr;
|
||||
@@ -84,6 +84,11 @@ struct RedOpArg<FuncMinMax<T>> {
|
||||
// of elements. These classes are intended to be specialized for specific
|
||||
// combinations of reduction function and pack size.
|
||||
|
||||
template<typename A, typename B, int EltPerPackA>
|
||||
struct Apply_Cast/*{
|
||||
static BytePack<EltPerPackA*sizeof(B)/sizeof(A)> cast(BytePack<EltPerPackA*sizeof(A)> a);
|
||||
}*/;
|
||||
|
||||
template<typename Fn, int EltPerPack>
|
||||
struct Apply_Reduce /*{
|
||||
static BytePack<EltPerPack*sizeof(T)> reduce(
|
||||
@@ -111,16 +116,60 @@ struct Apply_LoadMultimem/*{
|
||||
static BytePack<BytePerPack> load(Fn fn, uintptr_t addr);
|
||||
}*/;
|
||||
|
||||
|
||||
// Helpers for dealing with BytePack<0>'s
|
||||
template<typename A, typename B, int EltPerPack>
|
||||
struct Apply_Cast_MaybeEmpty: Apply_Cast<A, B, EltPerPack> {};
|
||||
template<typename A, typename B>
|
||||
struct Apply_Cast_MaybeEmpty<A, B, /*EltPerPack=*/0> {
|
||||
__device__ constexpr static BytePack<0> cast(BytePack<0> a) { return {}; }
|
||||
};
|
||||
|
||||
template<typename Fn, int EltPerPack>
|
||||
struct Apply_Reduce_MaybeEmpty: Apply_Reduce<Fn, EltPerPack> {};
|
||||
template<typename Fn>
|
||||
struct Apply_Reduce_MaybeEmpty<Fn, 0> {
|
||||
__device__ constexpr static BytePack<0> reduce(Fn fn, BytePack<0> a, BytePack<0> b) { return {}; }
|
||||
};
|
||||
|
||||
template<typename Fn, int EltPerPack>
|
||||
struct Apply_PreOp_MaybeEmpty: Apply_PreOp<Fn, EltPerPack> {};
|
||||
template<typename Fn>
|
||||
struct Apply_PreOp_MaybeEmpty<Fn, 0> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
__device__ constexpr static BytePack<0> preOp(Fn fn, BytePack<0> a) { return {}; }
|
||||
};
|
||||
|
||||
template<typename Fn, int EltPerPack>
|
||||
struct Apply_PostOp_MaybeEmpty: Apply_PostOp<Fn, EltPerPack> {};
|
||||
template<typename Fn>
|
||||
struct Apply_PostOp_MaybeEmpty<Fn, 0> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
__device__ constexpr static BytePack<0> postOp(Fn fn, BytePack<0> a) { return {}; }
|
||||
};
|
||||
|
||||
template<typename Fn, int BytePerPack>
|
||||
struct Apply_LoadMultimem_MaybeEmpty: Apply_LoadMultimem<Fn, BytePerPack> {};
|
||||
template<typename Fn>
|
||||
struct Apply_LoadMultimem_MaybeEmpty<Fn, 0> {
|
||||
__device__ constexpr static BytePack<0> load(Fn fn, uintptr_t addr) { return {}; }
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Public API for calling the trait classes. These take the data elements as a
|
||||
// pack of any type, which could be a BytePack<?> or any integral type (uint64_t,
|
||||
// uint32_t, etc.), and will return a new pack where each element has been
|
||||
// transformed appropriately.
|
||||
|
||||
template<typename A, typename B, typename PackA>
|
||||
__device__ __forceinline__ BytePack<BytePackOf<PackA>::Size*sizeof(B)/sizeof(A)> applyCast(PackA a) {
|
||||
return Apply_Cast_MaybeEmpty<A, B, BytePackOf<PackA>::Size/sizeof(A)>::cast(toPack(a));
|
||||
}
|
||||
|
||||
template<typename Fn, typename Pack>
|
||||
__device__ __forceinline__ Pack applyReduce(Fn fn, Pack a, Pack b) {
|
||||
return fromPack<Pack>(
|
||||
Apply_Reduce<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
Apply_Reduce_MaybeEmpty<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
::reduce(fn, toPack(a), toPack(b))
|
||||
);
|
||||
}
|
||||
@@ -128,7 +177,7 @@ __device__ __forceinline__ Pack applyReduce(Fn fn, Pack a, Pack b) {
|
||||
template<typename Fn, typename Pack>
|
||||
__device__ __forceinline__ Pack applyPreOp(Fn fn, Pack a) {
|
||||
return fromPack<Pack>(
|
||||
Apply_PreOp<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
Apply_PreOp_MaybeEmpty<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
::preOp(fn, toPack(a))
|
||||
);
|
||||
}
|
||||
@@ -136,23 +185,107 @@ __device__ __forceinline__ Pack applyPreOp(Fn fn, Pack a) {
|
||||
template<typename Fn, typename Pack>
|
||||
__device__ __forceinline__ Pack applyPostOp(Fn fn, Pack a) {
|
||||
return fromPack<Pack>(
|
||||
Apply_PostOp<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
Apply_PostOp_MaybeEmpty<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
::postOp(fn, toPack(a))
|
||||
);
|
||||
}
|
||||
|
||||
template<typename Fn, int BytePerPack>
|
||||
__device__ __forceinline__ BytePack<BytePerPack> applyLoadMultimem(Fn fn, uintptr_t addr) {
|
||||
return Apply_LoadMultimem<Fn, BytePerPack>::load(fn, addr);
|
||||
return Apply_LoadMultimem_MaybeEmpty<Fn, BytePerPack>::load(fn, addr);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Apply_Cast
|
||||
|
||||
template<typename A, typename B, int EltPerPack>
|
||||
struct Apply_Cast {
|
||||
__device__ __forceinline__ static BytePack<EltPerPack*sizeof(B)> cast(BytePack<EltPerPack*sizeof(A)> a) {
|
||||
BytePack<EltPerPack*sizeof(B)> b;
|
||||
b.half[0] = Apply_Cast<A, B, EltPerPack/2>::cast(a.half[0]);
|
||||
b.half[1] = Apply_Cast<A, B, EltPerPack/2>::cast(a.half[1]);
|
||||
return b;
|
||||
}
|
||||
};
|
||||
|
||||
template<typename A, typename B>
|
||||
struct Apply_Cast<A, B, /*EltPerPack=*/1> {
|
||||
__device__ __forceinline__ static BytePack<sizeof(B)> cast(BytePack<sizeof(A)> a) {
|
||||
return toPack(B(fromPack<A>(a)));
|
||||
}
|
||||
};
|
||||
|
||||
template<>
|
||||
struct Apply_Cast<__half, float, /*EltPerPack=*/1> {
|
||||
__device__ __forceinline__ static BytePack<sizeof(float)> cast(BytePack<sizeof(__half)> a) {
|
||||
return toPack(__half2float(fromPack<__half>(a)));
|
||||
}
|
||||
};
|
||||
template<>
|
||||
struct Apply_Cast<float, __half, /*EltPerPack=*/1> {
|
||||
__device__ __forceinline__ static BytePack<sizeof(__half)> cast(BytePack<sizeof(float)> a) {
|
||||
return toPack(__float2half_rn(fromPack<float>(a)));
|
||||
}
|
||||
};
|
||||
|
||||
template<>
|
||||
struct Apply_Cast<__half, float, /*EltPerPack=*/2> {
|
||||
__device__ __forceinline__ static BytePack<4*2> cast(BytePack<2*2> a) {
|
||||
return toPack(__half22float2(fromPack<__half2>(a)));
|
||||
}
|
||||
};
|
||||
template<>
|
||||
struct Apply_Cast<float, __half, /*EltPerPack=*/2> {
|
||||
__device__ __forceinline__ static BytePack<2*2> cast(BytePack<4*2> a) {
|
||||
return toPack(__float22half2_rn(fromPack<float2>(a)));
|
||||
}
|
||||
};
|
||||
|
||||
#if defined(__CUDA_BF16_TYPES_EXIST__) && (CUDART_RUNTIME >= 12000 || __CUDA_ARCH__ >= 800)
|
||||
template<>
|
||||
struct Apply_Cast<__nv_bfloat16, float, /*EltPerPack=*/2> {
|
||||
__device__ __forceinline__ static BytePack<4*2> cast(BytePack<2*2> a) {
|
||||
return toPack(__bfloat1622float2(fromPack<__nv_bfloat162>(a)));
|
||||
}
|
||||
};
|
||||
template<>
|
||||
struct Apply_Cast<float ,__nv_bfloat16, /*EltPerPack=*/2> {
|
||||
__device__ __forceinline__ static BytePack<2*2> cast(BytePack<4*2> a) {
|
||||
return toPack(__float22bfloat162_rn(fromPack<float2>(a)));
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
#define EASY_CAST(A, B, EltPerPack, VecA, VecB) \
|
||||
template<> \
|
||||
struct Apply_Cast<A, B, EltPerPack> { \
|
||||
__device__ __forceinline__ static BytePack<sizeof(B)*EltPerPack> cast(BytePack<sizeof(A)*EltPerPack> a) { \
|
||||
return toPack(VecB(fromPack<VecA>(a))); \
|
||||
} \
|
||||
}; \
|
||||
template<> \
|
||||
struct Apply_Cast<B, A, EltPerPack> { \
|
||||
__device__ __forceinline__ static BytePack<sizeof(A)*EltPerPack> cast(BytePack<sizeof(B)*EltPerPack> b) { \
|
||||
return toPack(VecA(fromPack<VecB>(b))); \
|
||||
} \
|
||||
};
|
||||
|
||||
#if defined(__CUDA_FP8_TYPES_EXIST__)
|
||||
EASY_CAST(__nv_fp8_e5m2, float, 2, __nv_fp8x2_e5m2, float2)
|
||||
EASY_CAST(__nv_fp8_e5m2, float, 4, __nv_fp8x4_e5m2, float4)
|
||||
|
||||
EASY_CAST(__nv_fp8_e4m3, float, 2, __nv_fp8x2_e4m3, float2)
|
||||
EASY_CAST(__nv_fp8_e4m3, float, 4, __nv_fp8x4_e4m3, float4)
|
||||
#endif
|
||||
#undef EASY_CAST
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Apply_Reduce
|
||||
|
||||
// Nonsensical base case
|
||||
template<typename Fn>
|
||||
struct Apply_Reduce<Fn, /*EltPerPack=*/0> {
|
||||
__device__ static BytePack<0> reduce(Fn fn, BytePack<0> a, BytePack<0> b) {
|
||||
__device__ __forceinline__ static BytePack<0> reduce(Fn fn, BytePack<0> a, BytePack<0> b) {
|
||||
return {};
|
||||
}
|
||||
};
|
||||
@@ -164,7 +297,7 @@ struct Apply_Reduce<Fn, /*EltPerPack=*/0> {
|
||||
template<typename Fn, int EltPerPack>
|
||||
struct Apply_Reduce {
|
||||
template<int Size>
|
||||
__device__ static BytePack<Size> reduce(Fn fn, BytePack<Size> a, BytePack<Size> b) {
|
||||
__device__ __forceinline__ static BytePack<Size> reduce(Fn fn, BytePack<Size> a, BytePack<Size> b) {
|
||||
a.half[0] = Apply_Reduce<Fn, EltPerPack/2>::reduce(fn, a.half[0], b.half[0]);
|
||||
a.half[1] = Apply_Reduce<Fn, EltPerPack/2>::reduce(fn, a.half[1], b.half[1]);
|
||||
return a;
|
||||
@@ -174,25 +307,25 @@ struct Apply_Reduce {
|
||||
// Base case definitions (EltPerPack == 1)
|
||||
template<typename T>
|
||||
struct Apply_Reduce<FuncCopy<T>, /*EltPerPack=*/1> {
|
||||
__device__ static BytePack<sizeof(T)> reduce(FuncCopy<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncCopy<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
return a;
|
||||
}
|
||||
};
|
||||
template<typename T>
|
||||
struct Apply_Reduce<FuncSum<T>, /*EltPerPack=*/1> {
|
||||
__device__ static BytePack<sizeof(T)> reduce(FuncSum<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncSum<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
return toPack<T>(fromPack<T>(a) + fromPack<T>(b));
|
||||
}
|
||||
};
|
||||
template<typename T>
|
||||
struct Apply_Reduce<FuncProd<T>, /*EltPerPack=*/1> {
|
||||
__device__ static BytePack<sizeof(T)> reduce(FuncProd<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncProd<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
return toPack<T>(fromPack<T>(a) * fromPack<T>(b));
|
||||
}
|
||||
};
|
||||
template<typename T>
|
||||
struct Apply_Reduce<FuncMinMax<T>, /*EltPerPack=*/1> {
|
||||
__device__ static BytePack<sizeof(T)> reduce(FuncMinMax<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncMinMax<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
return (a.native ^ fn.xormask.native) < (b.native ^ fn.xormask.native) ? a : b;
|
||||
}
|
||||
};
|
||||
@@ -200,7 +333,7 @@ struct Apply_Reduce<FuncMinMax<T>, /*EltPerPack=*/1> {
|
||||
// Optimizations for specfic types and element count combinations:
|
||||
template<>
|
||||
struct Apply_Reduce<FuncSum<uint8_t>, /*EltPerPack=*/4> {
|
||||
__device__ static BytePack<4> reduce(FuncSum<uint8_t> fn, BytePack<4> a, BytePack<4> b) {
|
||||
__device__ __forceinline__ static BytePack<4> reduce(FuncSum<uint8_t> fn, BytePack<4> a, BytePack<4> b) {
|
||||
constexpr uint32_t even = 0x00ff00ffu;
|
||||
uint32_t x = (a.native & even) + (b.native & even);
|
||||
uint32_t y = (a.native & ~even) + (b.native & ~even);
|
||||
@@ -236,7 +369,7 @@ struct Apply_Reduce<FuncMinMax<uint8_t>, /*EltPerPack=*/4> {
|
||||
|
||||
template<>
|
||||
struct Apply_Reduce<FuncProd<uint8_t>, /*EltPerPack=*/4> {
|
||||
__device__ static BytePack<4> reduce(FuncProd<uint8_t> fn, BytePack<4> apack, BytePack<4> bpack) {
|
||||
__device__ __forceinline__ static BytePack<4> reduce(FuncProd<uint8_t> fn, BytePack<4> apack, BytePack<4> bpack) {
|
||||
uint32_t a = apack.native;
|
||||
uint32_t b = bpack.native;
|
||||
uint32_t ab0 = (a*b) & 0xffu;
|
||||
@@ -332,7 +465,7 @@ template<typename Fn, int EltPerPack>
|
||||
struct Apply_PreOp {
|
||||
static constexpr bool IsIdentity = Apply_PreOp<Fn, EltPerPack/2>::IsIdentity;
|
||||
template<int Size>
|
||||
__device__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
|
||||
__device__ __forceinline__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
|
||||
#if __cpp_if_constexpr
|
||||
if constexpr(!IsIdentity) {
|
||||
#else
|
||||
@@ -352,7 +485,7 @@ template<typename Fn>
|
||||
struct Apply_PreOp<Fn, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
template<int Size>
|
||||
__device__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
|
||||
__device__ __forceinline__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
|
||||
return a;
|
||||
}
|
||||
};
|
||||
@@ -360,7 +493,7 @@ struct Apply_PreOp<Fn, /*EltPerPack=*/1> {
|
||||
template<typename Fn>
|
||||
struct Apply_PreOp<Fn, /*EltPerPack=*/0> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
__device__ static BytePack<0> preOp(Fn fn, BytePack<0> a) {
|
||||
__device__ __forceinline__ static BytePack<0> preOp(Fn fn, BytePack<0> a) {
|
||||
return {};
|
||||
}
|
||||
};
|
||||
@@ -373,7 +506,7 @@ template<typename Fn, int EltPerPack>
|
||||
struct Apply_PostOp {
|
||||
static constexpr bool IsIdentity = Apply_PostOp<Fn, EltPerPack/2>::IsIdentity;
|
||||
template<int Size>
|
||||
__device__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
|
||||
__device__ __forceinline__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
|
||||
#if __cpp_if_constexpr
|
||||
if constexpr(!IsIdentity) {
|
||||
#else
|
||||
@@ -393,7 +526,7 @@ template<typename Fn>
|
||||
struct Apply_PostOp<Fn, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
template<int Size>
|
||||
__device__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
|
||||
__device__ __forceinline__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
|
||||
return a;
|
||||
}
|
||||
};
|
||||
@@ -401,7 +534,7 @@ struct Apply_PostOp<Fn, /*EltPerPack=*/1> {
|
||||
template<typename Fn>
|
||||
struct Apply_PostOp<Fn, /*EltPerPack=*/0> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
__device__ static BytePack<0> postOp(Fn fn, BytePack<0> a) {
|
||||
__device__ __forceinline__ static BytePack<0> postOp(Fn fn, BytePack<0> a) {
|
||||
return {};
|
||||
}
|
||||
};
|
||||
@@ -413,7 +546,7 @@ struct Apply_PostOp<Fn, /*EltPerPack=*/0> {
|
||||
template<typename T>
|
||||
struct RedOpArg<FuncPreMulSum<T>> {
|
||||
static constexpr bool ArgUsed = true;
|
||||
__device__ static uint64_t loadArg(void *ptr) {
|
||||
__device__ __forceinline__ static uint64_t loadArg(void *ptr) {
|
||||
union { uint64_t u64; T val; };
|
||||
u64 = 0;
|
||||
val = *(T*)ptr;
|
||||
@@ -426,7 +559,7 @@ template<typename T>
|
||||
struct FuncPreMulSum {
|
||||
using EltType = T;
|
||||
T scalar;
|
||||
__device__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
union { uint64_t u64; T val; };
|
||||
u64 = opArg;
|
||||
scalar = val;
|
||||
@@ -441,7 +574,7 @@ struct FuncPreMulSum<half> {
|
||||
using EltType = half;
|
||||
#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
|
||||
__half2 scalar;
|
||||
__device__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
union { uint64_t u64; __half val; };
|
||||
u64 = opArg;
|
||||
scalar.x = val;
|
||||
@@ -449,7 +582,7 @@ struct FuncPreMulSum<half> {
|
||||
}
|
||||
#else
|
||||
float scalar;
|
||||
__device__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
union { uint64_t u64; __half val; };
|
||||
u64 = opArg;
|
||||
scalar = (float)val;
|
||||
@@ -466,7 +599,7 @@ struct FuncPreMulSum<half> {
|
||||
using EltType = __nv_bfloat16;
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
__nv_bfloat162 scalar;
|
||||
__device__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
union { uint64_t u64; __nv_bfloat16 val; };
|
||||
u64 = opArg;
|
||||
scalar.x = val;
|
||||
@@ -474,7 +607,7 @@ struct FuncPreMulSum<half> {
|
||||
}
|
||||
#else
|
||||
float scalar;
|
||||
__device__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
union { uint64_t u64; __nv_bfloat16 val; };
|
||||
u64 = opArg;
|
||||
scalar = __bfloat162float(val);
|
||||
@@ -489,7 +622,7 @@ struct FuncPreMulSum<half> {
|
||||
struct FuncPreMulSum<__nv_fp8_e4m3> {
|
||||
using EltType = __nv_fp8_e4m3;
|
||||
__half2 scalar2;
|
||||
__device__ FuncPreMulSum(uint64_t opArg) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg) {
|
||||
union { uint64_t u64; __nv_fp8_storage_t val; };
|
||||
u64 = opArg;
|
||||
scalar2.x = __half(__nv_cvt_fp8_to_halfraw(val, __NV_E4M3));
|
||||
@@ -501,7 +634,7 @@ struct FuncPreMulSum<half> {
|
||||
struct FuncPreMulSum<__nv_fp8_e5m2> {
|
||||
using EltType = __nv_fp8_e5m2;
|
||||
__half2 scalar2;
|
||||
__device__ FuncPreMulSum(uint64_t opArg) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg) {
|
||||
union { uint64_t u64; __nv_fp8_storage_t val; };
|
||||
u64 = opArg;
|
||||
scalar2.x = __half(__nv_cvt_fp8_to_halfraw(val, __NV_E5M2));
|
||||
@@ -513,7 +646,7 @@ struct FuncPreMulSum<half> {
|
||||
|
||||
template<typename T, int EltPerPack>
|
||||
struct Apply_Reduce<FuncPreMulSum<T>, EltPerPack> {
|
||||
__device__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncPreMulSum<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncPreMulSum<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
|
||||
// FuncPreMulSum reduce dispatches to FuncSum.
|
||||
return Apply_Reduce<FuncSum<T>, EltPerPack>::reduce(FuncSum<T>(), a, b);
|
||||
}
|
||||
@@ -523,7 +656,7 @@ struct Apply_Reduce<FuncPreMulSum<T>, EltPerPack> {
|
||||
template<typename T>
|
||||
struct Apply_PreOp<FuncPreMulSum<T>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(T)> preOp(FuncPreMulSum<T> fn, BytePack<sizeof(T)> a) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> preOp(FuncPreMulSum<T> fn, BytePack<sizeof(T)> a) {
|
||||
return toPack<T>(fromPack<T>(a) * fn.scalar);
|
||||
}
|
||||
};
|
||||
@@ -534,7 +667,7 @@ struct Apply_PreOp<FuncPreMulSum<T>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(half)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half)> a) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(half)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half)> a) {
|
||||
#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
|
||||
return toPack<half>(__hmul(fromPack<half>(a), fn.scalar.x));
|
||||
#else
|
||||
@@ -546,7 +679,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/2> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(half2)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half2)> a) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(half2)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half2)> a) {
|
||||
return toPack<half2>(__hmul2(fromPack<half2>(a), fn.scalar));
|
||||
}
|
||||
};
|
||||
@@ -559,7 +692,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<__nv_bfloat16>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_bfloat16)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_bfloat16)> preOp(
|
||||
FuncPreMulSum<__nv_bfloat16> fn, BytePack<sizeof(__nv_bfloat16)> a
|
||||
) {
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
@@ -573,7 +706,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<__nv_bfloat16>, /*EltPerPack=*/2> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_bfloat162)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_bfloat162)> preOp(
|
||||
FuncPreMulSum<__nv_bfloat16> fn, BytePack<sizeof(__nv_bfloat162)> a
|
||||
) {
|
||||
return toPack<__nv_bfloat162>(__hmul2(fromPack<__nv_bfloat162>(a), fn.scalar));
|
||||
@@ -590,7 +723,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e4m3>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_fp8_e4m3)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8_e4m3)> preOp(
|
||||
FuncPreMulSum<__nv_fp8_e4m3> fn, BytePack<sizeof(__nv_fp8_e4m3)> a
|
||||
) {
|
||||
return toPack<__nv_fp8_e4m3>(__nv_fp8_e4m3(__hmul(__half(fromPack<__nv_fp8_e4m3>(a)), fn.scalar2.x)));
|
||||
@@ -599,7 +732,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e4m3>, /*EltPerPack=*/2> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_fp8x2_e4m3)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8x2_e4m3)> preOp(
|
||||
FuncPreMulSum<__nv_fp8_e4m3> fn, BytePack<sizeof(__nv_fp8x2_e4m3)> a
|
||||
) {
|
||||
return toPack<__nv_fp8x2_e4m3>(__nv_fp8x2_e4m3(__hmul2(__half2(fromPack<__nv_fp8x2_e4m3>(a)), fn.scalar2)));
|
||||
@@ -609,7 +742,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e5m2>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_fp8_e5m2)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8_e5m2)> preOp(
|
||||
FuncPreMulSum<__nv_fp8_e5m2> fn, BytePack<sizeof(__nv_fp8_e5m2)> a
|
||||
) {
|
||||
return toPack<__nv_fp8_e5m2>(__nv_fp8_e5m2(__hmul(__half(fromPack<__nv_fp8_e5m2>(a)), fn.scalar2.x)));
|
||||
@@ -618,7 +751,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e5m2>, /*EltPerPack=*/2> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_fp8x2_e5m2)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8x2_e5m2)> preOp(
|
||||
FuncPreMulSum<__nv_fp8_e5m2> fn, BytePack<sizeof(__nv_fp8x2_e5m2)> a
|
||||
) {
|
||||
return toPack<__nv_fp8x2_e5m2>(__nv_fp8x2_e5m2(__hmul2(__half2(fromPack<__nv_fp8x2_e5m2>(a)), fn.scalar2)));
|
||||
@@ -633,7 +766,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<typename T>
|
||||
struct RedOpArg<FuncSumPostDiv<T>> {
|
||||
static constexpr bool ArgUsed = true;
|
||||
__device__ static uint64_t loadArg(void *ptr) {
|
||||
__device__ __forceinline__ static uint64_t loadArg(void *ptr) {
|
||||
return *(uint64_t*)ptr;
|
||||
}
|
||||
};
|
||||
@@ -646,12 +779,12 @@ struct FuncSumPostDiv {
|
||||
uint32_t divisor:31, isSigned:1;
|
||||
UintType recip;
|
||||
|
||||
__device__ FuncSumPostDiv(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncSumPostDiv(uint64_t opArg=0) {
|
||||
isSigned = opArg & 1;
|
||||
divisor = opArg >> 1;
|
||||
recip = UintType(-1)/divisor;
|
||||
}
|
||||
__device__ T divide(T x) {
|
||||
__device__ __forceinline__ T divide(T x) {
|
||||
// x is negative iff we are in signed mode and the top bit is set
|
||||
bool xneg = isSigned && (x & ~(T(-1)>>1));
|
||||
// Compute abs(x):
|
||||
@@ -673,7 +806,7 @@ struct FuncSumPostDiv {
|
||||
template<typename T, int EltPerPack>
|
||||
struct Apply_Reduce<FuncSumPostDiv<T>, EltPerPack>:
|
||||
Apply_Reduce<FuncSum<T>, EltPerPack> {
|
||||
__device__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncSumPostDiv<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncSumPostDiv<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
|
||||
// FuncSumPostDiv reduce dispatches to FuncSum.
|
||||
return Apply_Reduce<FuncSum<T>, EltPerPack>::reduce(FuncSum<T>(), a, b);
|
||||
}
|
||||
@@ -682,7 +815,7 @@ struct Apply_Reduce<FuncSumPostDiv<T>, EltPerPack>:
|
||||
template<typename T>
|
||||
struct Apply_PostOp<FuncSumPostDiv<T>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(T)> postOp(FuncSumPostDiv<T> fn, BytePack<sizeof(T)> a) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> postOp(FuncSumPostDiv<T> fn, BytePack<sizeof(T)> a) {
|
||||
return toPack<T>(fn.divide(fromPack<T>(a)));
|
||||
}
|
||||
};
|
||||
@@ -690,120 +823,145 @@ struct Apply_PostOp<FuncSumPostDiv<T>, /*EltPerPack=*/1> {
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Apply_LoadMultimem
|
||||
|
||||
#define SIZEOF_BytePack_field_u16 2
|
||||
#define PTX_REG_BytePack_field_u16 "h"
|
||||
#define RegCode_for_size_1 "r"
|
||||
#define RegCode_for_size_2 "h"
|
||||
#define RegCode_for_size_4 "r"
|
||||
#define RegCode_for_size_8 "l"
|
||||
|
||||
#define SIZEOF_BytePack_field_u32 4
|
||||
#define PTX_REG_BytePack_field_u32 "r"
|
||||
#define RegSize_for_size_1 4
|
||||
#define RegSize_for_size_2 2
|
||||
#define RegSize_for_size_4 4
|
||||
#define RegSize_for_size_8 8
|
||||
|
||||
#define SIZEOF_BytePack_field_u64 8
|
||||
#define PTX_REG_BytePack_field_u64 "l"
|
||||
#define PtxAcc_for_u32
|
||||
#define PtxAcc_for_s32
|
||||
#define PtxAcc_for_s64
|
||||
#define PtxAcc_for_u64
|
||||
#define PtxAcc_for_f32
|
||||
#define PtxAcc_for_f64
|
||||
#if CUDART_VERSION >= 12020
|
||||
#define PtxAcc_for_f16 ".acc::f32"
|
||||
#define PtxAcc_for_bf16 ".acc::f32"
|
||||
#define PtxAcc_for_f16x2 ".acc::f32"
|
||||
#define PtxAcc_for_bf16x2 ".acc::f32"
|
||||
#else
|
||||
#define PtxAcc_for_f16
|
||||
#define PtxAcc_for_bf16
|
||||
#define PtxAcc_for_f16x2
|
||||
#define PtxAcc_for_bf16x2
|
||||
#endif
|
||||
#define PtxAcc_for_e4m3 ".acc::f16"
|
||||
#define PtxAcc_for_e5m2 ".acc::f16"
|
||||
#define PtxAcc_for_e4m3x4 ".acc::f16"
|
||||
#define PtxAcc_for_e5m2x4 ".acc::f16"
|
||||
|
||||
#define DEFINE_Apply_LoadMultimem_sum(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_sum(T, ptx_ty, PackSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncSum<T>, SIZEOF_BytePack_field_##pack_field> { \
|
||||
static constexpr int PackSize = SIZEOF_BytePack_field_##pack_field; \
|
||||
__device__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
BytePack<PackSize> ans; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
|
||||
struct Apply_LoadMultimem<FuncSum<T>, PackSize> { \
|
||||
__device__ __forceinline__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
BytePack<RegSize_for_size_##PackSize> reg; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add" PtxAcc_for_##ptx_ty "." #ptx_ty " %0, [%1];" \
|
||||
: "=" RegCode_for_size_##PackSize(reg.native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
BytePack<PackSize> ans; \
|
||||
ans.native = reg.native; \
|
||||
return ans; \
|
||||
} \
|
||||
};
|
||||
#define DEFINE_Apply_LoadMultimem_minmax(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_minmax(T, ptx_ty, PackSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncMinMax<T>, SIZEOF_BytePack_field_##pack_field> { \
|
||||
static constexpr int PackSize = SIZEOF_BytePack_field_##pack_field; \
|
||||
__device__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
BytePack<PackSize> ans; \
|
||||
struct Apply_LoadMultimem<FuncMinMax<T>, PackSize> { \
|
||||
__device__ __forceinline__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
BytePack<RegSize_for_size_##PackSize> reg; \
|
||||
if (fn.isMinNotMax) { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.min." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
|
||||
: "=" RegCode_for_size_##PackSize(reg.native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
} else { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.max." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
|
||||
: "=" RegCode_for_size_##PackSize(reg.native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
} \
|
||||
BytePack<PackSize> ans; \
|
||||
ans.native = reg.native; \
|
||||
return ans; \
|
||||
} \
|
||||
};
|
||||
|
||||
#define DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, VecEltSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncSum<T>, 4*(SIZEOF_BytePack_field_##pack_field)> { \
|
||||
static constexpr int PackSize = 4*(SIZEOF_BytePack_field_##pack_field); \
|
||||
__device__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
BytePack<PackSize> ans; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
|
||||
struct Apply_LoadMultimem<FuncSum<T>, 4*(VecEltSize)> { \
|
||||
static constexpr int PackSize = 4*(VecEltSize); \
|
||||
__device__ __forceinline__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
union { BytePack<PackSize> ans; BytePack<VecEltSize> elts[4]; }; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add" PtxAcc_for_##ptx_ty ".v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
|
||||
: "=" RegCode_for_size_##VecEltSize(elts[0].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[1].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[2].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[3].native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
return ans; \
|
||||
} \
|
||||
};
|
||||
#define DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, VecEltSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncMinMax<T>, 4*(SIZEOF_BytePack_field_##pack_field)> { \
|
||||
static constexpr int PackSize = 4*(SIZEOF_BytePack_field_##pack_field); \
|
||||
__device__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
BytePack<PackSize> ans; \
|
||||
struct Apply_LoadMultimem<FuncMinMax<T>, 4*(VecEltSize)> { \
|
||||
static constexpr int PackSize = 4*(VecEltSize); \
|
||||
__device__ __forceinline__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
union { BytePack<PackSize> ans; BytePack<VecEltSize> elts[4]; }; \
|
||||
if (fn.isMinNotMax) { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.min.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
|
||||
: "=" RegCode_for_size_##VecEltSize(elts[0].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[1].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[2].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[3].native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
} else { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.max.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
|
||||
: "=" RegCode_for_size_##VecEltSize(elts[0].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[1].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[2].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[3].native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
} \
|
||||
return ans; \
|
||||
} \
|
||||
};
|
||||
|
||||
#define DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(T, ptx_ty, pack_field) \
|
||||
DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(T, ptx_ty, VecEltSize) \
|
||||
DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, VecEltSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncSum<T>, sizeof(T)> { \
|
||||
__device__ static BytePack<sizeof(T)> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
BytePack<2*sizeof(T)> tmp; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
|
||||
: "l"(addr & -uintptr_t(2*sizeof(T))) : "memory"); \
|
||||
return tmp.half[(addr/sizeof(T))%2]; \
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
union { BytePack<VecEltSize> tmp; BytePack<sizeof(T)> elts[(VecEltSize)/sizeof(T)]; }; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add" PtxAcc_for_##ptx_ty "." #ptx_ty " %0, [%1];" \
|
||||
: "=" RegCode_for_size_##VecEltSize(tmp.native) \
|
||||
: "l"(addr & -uintptr_t(VecEltSize)) : "memory"); \
|
||||
return elts[(addr/sizeof(T))%((VecEltSize)/sizeof(T))]; \
|
||||
} \
|
||||
};
|
||||
#define DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(T, ptx_ty, pack_field) \
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(T, ptx_ty, VecEltSize) \
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, VecEltSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncMinMax<T>, sizeof(T)> { \
|
||||
__device__ static BytePack<sizeof(T)> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
BytePack<2*sizeof(T)> tmp; \
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
union { BytePack<VecEltSize> tmp; BytePack<sizeof(T)> elts[(VecEltSize)/sizeof(T)]; }; \
|
||||
if (fn.isMinNotMax) { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.min." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
|
||||
: "l"(addr & -uintptr_t(2*sizeof(T))) : "memory"); \
|
||||
: "=" RegCode_for_size_##VecEltSize(tmp.native) \
|
||||
: "l"(addr & -uintptr_t(VecEltSize)) : "memory"); \
|
||||
} else { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.max." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
|
||||
: "l"(addr & -uintptr_t(2*sizeof(T))) : "memory"); \
|
||||
: "=" RegCode_for_size_##VecEltSize(tmp.native) \
|
||||
: "l"(addr & -uintptr_t(VecEltSize)) : "memory"); \
|
||||
} \
|
||||
return tmp.half[(addr/sizeof(T))%2]; \
|
||||
return elts[(addr/sizeof(T))%((VecEltSize)/sizeof(T))]; \
|
||||
} \
|
||||
};
|
||||
|
||||
template<typename Fn, int BytePerPack>
|
||||
struct Apply_LoadMultimem {
|
||||
__device__ static BytePack<BytePerPack> load(Fn fn, uintptr_t addr) {
|
||||
__device__ __forceinline__ static BytePack<BytePerPack> load(Fn fn, uintptr_t addr) {
|
||||
__trap();
|
||||
return {};
|
||||
}
|
||||
@@ -826,29 +984,36 @@ struct Apply_LoadMultimem {
|
||||
/*multimem.ld_reduce not supported:*/ 0;
|
||||
};
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(uint32_t, u32, u32)
|
||||
DEFINE_Apply_LoadMultimem_minmax(uint32_t, u32, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum(uint32_t, u32, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax(uint32_t, u32, 4)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(int32_t, s32, u32)
|
||||
DEFINE_Apply_LoadMultimem_minmax(int32_t, s32, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum(int32_t, s32, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax(int32_t, s32, 4)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(uint64_t, u64, u64)
|
||||
DEFINE_Apply_LoadMultimem_minmax(uint64_t, u64, u64)
|
||||
DEFINE_Apply_LoadMultimem_sum(uint64_t, u64, 8)
|
||||
DEFINE_Apply_LoadMultimem_minmax(uint64_t, u64, 8)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(int64_t, u64, u64)
|
||||
DEFINE_Apply_LoadMultimem_minmax(int64_t, s64, u64)
|
||||
DEFINE_Apply_LoadMultimem_sum(int64_t, u64, 8)
|
||||
DEFINE_Apply_LoadMultimem_minmax(int64_t, s64, 8)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(float, f32, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4(float, f32, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum(float, f32, 4)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4(float, f32, 4)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(double, f64, u64)
|
||||
DEFINE_Apply_LoadMultimem_sum(double, f64, 8)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(half, f16x2, u32)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(half, f16x2, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(half, f16x2, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(half, f16x2, 4)
|
||||
|
||||
#if defined(__CUDA_BF16_TYPES_EXIST__)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(__nv_bfloat16, bf16x2, u32)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(__nv_bfloat16, bf16x2, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(__nv_bfloat16, bf16x2, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(__nv_bfloat16, bf16x2, 4)
|
||||
#endif
|
||||
|
||||
#if NCCL_CUDA_ARCH_FAMILY_SPECIFIC == 1000 || NCCL_CUDA_ARCH_FAMILY_SPECIFIC == 1010 || NCCL_CUDA_ARCH_SPECIFIC == 1200 || NCCL_CUDA_ARCH_SPECIFIC == 1210
|
||||
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(__nv_fp8_e4m3, e4m3x4, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(__nv_fp8_e4m3, e4m3x4, 4)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(__nv_fp8_e5m2, e5m2x4, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(__nv_fp8_e5m2, e5m2x4, 4)
|
||||
#endif
|
||||
#else
|
||||
template<typename Fn>
|
||||
@@ -860,11 +1025,29 @@ struct Apply_LoadMultimem {
|
||||
#undef DEFINE_Apply_LoadMultimem
|
||||
#undef DEFINE_Apply_LoadMultimem_v4
|
||||
#undef DEFINE_Apply_LoadMultimem_v4x2_and_subhalf
|
||||
#undef SIZEOF_BytePack_field_u64
|
||||
#undef PTX_REG_BytePack_field_u64
|
||||
#undef SIZEOF_BytePack_field_u32
|
||||
#undef PTX_REG_BytePack_field_u32
|
||||
#undef SIZEOF_BytePack_field_u16
|
||||
#undef PTX_REG_BytePack_field_u16
|
||||
|
||||
#undef RegCode_for_size_2
|
||||
#undef RegCode_for_size_4
|
||||
#undef RegCode_for_size_8
|
||||
|
||||
#undef RegSize_for_size_1
|
||||
#undef RegSize_for_size_2
|
||||
#undef RegSize_for_size_4
|
||||
#undef RegSize_for_size_8
|
||||
|
||||
#undef PtxAcc_for_u32
|
||||
#undef PtxAcc_for_s32
|
||||
#undef PtxAcc_for_s64
|
||||
#undef PtxAcc_for_u64
|
||||
#undef PtxAcc_for_f32
|
||||
#undef PtxAcc_for_f64
|
||||
#undef PtxAcc_for_f16
|
||||
#undef PtxAcc_for_bf16
|
||||
#undef PtxAcc_for_f16x2
|
||||
#undef PtxAcc_for_bf16x2
|
||||
#undef PtxAcc_for_e4m3
|
||||
#undef PtxAcc_for_e5m2
|
||||
#undef PtxAcc_for_e4m3x4
|
||||
#undef PtxAcc_for_e5m2x4
|
||||
|
||||
#endif // REDUCE_KERNEL_H_
|
||||
|
||||
Reference in New Issue
Block a user