AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Modifying ReduceOrCopyMulti to accept number of preOp source to support clique-based kernels

Browse Source
This commit is contained in:
Gilbert Lee
2021-08-11 11:00:34 -05:00
parent a4929465c5
commit ae13d2a354
4 changed files with 21 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/collectives/device/common_kernel.h
+15 -9
View File
@@ -503,7 +503,7 @@ inline __device__ void Store128(Pack128* p, Pack128& v) {
template<class FUNC, typename T, int UNROLL, int MINSRCS, int MAXSRCS, int MINDSTS, int MAXDSTS>
__device__ __forceinline__ void ReduceCopyMulti(const int w, const int nw, const int t,
FUNC fn, bool preOpSrc0, bool postOp, int nsrcs, const T** s, int ndsts, T** d, const int elemOffset, const int Nelem
FUNC fn, int const numPreOpSrcs, bool postOp, int nsrcs, const T** s, int ndsts, T** d, const int elemOffset, const int Nelem
) {
const int inc = nw * UNROLL * WARP_SIZE;
int offset = w * UNROLL * WARP_SIZE + t;
@@ -517,7 +517,7 @@ __device__ __forceinline__ void ReduceCopyMulti(const int w, const int nw, const
T vals[UNROLL];
// Load and reduce
for (int u = 0; u < UNROLL; ++u) vals[u] = vFetch(srcs[0]+u*WARP_SIZE);
if (preOpSrc0) {
if (numPreOpSrcs) {
for (int u = 0; u < UNROLL; ++u) vals[u] = FuncTraits<FUNC>().preOp(fn, vals[u]);
}
@@ -525,6 +525,9 @@ __device__ __forceinline__ void ReduceCopyMulti(const int w, const int nw, const
for (int i=1; i<MINSRCS; i++) {
T vals2[UNROLL];
for (int u = 0; u < UNROLL; ++u) vals2[u] = vFetch(srcs[i]+u*WARP_SIZE);
if (i < numPreOpSrcs) {
for (int u = 0; u < UNROLL; ++u) vals2[u] = FuncTraits<FUNC>().preOp(fn, vals2[u]);
}
for (int u = 0; u < UNROLL; ++u) vals[u] = fn(vals[u], vals2[u]);
}
#pragma unroll
@@ -560,7 +563,7 @@ __device__ __forceinline__ void ReduceCopyMulti(const int w, const int nw, const
template<class FUNC, typename T, int UNROLL, int MINSRCS, int MAXSRCS, int MINDSTS, int MAXDSTS>
__device__ __forceinline__ void ReduceCopy128bMulti(const int w, const int nw, const int t,
FUNC fn, bool preOpSrc0, bool postOp, int nsrcs, const T** s, int ndsts, T** d, const int elemOffset, const int Npack
FUNC fn, int numPreOpSrcs, bool postOp, int nsrcs, const T** s, int ndsts, T** d, const int elemOffset, const int Npack
) {
const int inc = nw * UNROLL * WARP_SIZE;
int offset = w * UNROLL * WARP_SIZE + t;
@@ -574,7 +577,7 @@ __device__ __forceinline__ void ReduceCopy128bMulti(const int w, const int nw, c
Pack128 vals[UNROLL];
// Load and reduce
for (int u = 0; u < UNROLL; ++u) Fetch128(vals[u], srcs[0]+u*WARP_SIZE);
if (preOpSrc0) {
if (numPreOpSrcs) {
for (int u = 0; u < UNROLL; ++u) MULTI128<FUNC, T>().preOp(fn, vals[u]);
}
@@ -582,6 +585,9 @@ __device__ __forceinline__ void ReduceCopy128bMulti(const int w, const int nw, c
for (int i=1; i<MINSRCS; i++) {
Pack128 vals2[UNROLL];
for (int u = 0; u < UNROLL; ++u) Fetch128(vals2[u], srcs[i]+u*WARP_SIZE);
if (i < numPreOpSrcs) {
for (int u = 0; u < UNROLL; ++u) MULTI128<FUNC, T>().preOp(fn, vals2[u]);
}
for (int u = 0; u < UNROLL; ++u) MULTI128<FUNC, T>()(fn, vals[u], vals2[u]);
}
#pragma unroll
@@ -623,7 +629,7 @@ __device__ int ptrAlign128(T* ptr) { return (uint64_t)ptr % alignof(uint32_t); }
template<int UNROLL, class FUNC, typename T, int MINSRCS, int MAXSRCS, int MINDSTS, int MAXDSTS>
__device__ __forceinline__ void ReduceOrCopyMulti(
const int tid, const int nthreads, FUNC fn, bool preOpSrc0, bool postOp, int nsrcs, const T** srcs, int ndsts, T** dsts, int N
const int tid, const int nthreads, FUNC fn, int numPreOpSrcs, bool postOp, int nsrcs, const T** srcs, int ndsts, T** dsts, int N
) {
int Nrem = N;
if (Nrem <= 0) return;
@@ -651,7 +657,7 @@ __device__ __forceinline__ void ReduceOrCopyMulti(
int Nelem = Npack * PACKELEMS;
ReduceCopy128bMulti<FUNC, T, AUTOUNROLL, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>
(w, nw, t, fn, preOpSrc0, postOp, nsrcs, srcs, ndsts, dsts, offset, Npack);
(w, nw, t, fn, numPreOpSrcs, postOp, nsrcs, srcs, ndsts, dsts, offset, Npack);
Nrem -= Nelem;
if (Nrem == 0) return;
@@ -662,7 +668,7 @@ __device__ __forceinline__ void ReduceOrCopyMulti(
Nelem = Npack * PACKELEMS;
ReduceCopy128bMulti<FUNC, T, 1, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>
(w, nw, t, fn, preOpSrc0, postOp, nsrcs, srcs, ndsts, dsts, offset, Npack);
(w, nw, t, fn, numPreOpSrcs, postOp, nsrcs, srcs, ndsts, dsts, offset, Npack);
Nrem -= Nelem;
if (Nrem == 0) return;
@@ -673,7 +679,7 @@ __device__ __forceinline__ void ReduceOrCopyMulti(
int Nelem = (Nrem / (AUTOUNROLL*PACKELEMS/2*WARP_SIZE)) * (AUTOUNROLL*PACKELEMS/2*WARP_SIZE); // round down
ReduceCopyMulti<FUNC, T, AUTOUNROLL*PACKELEMS/2, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>
(w, nw, t, fn, preOpSrc0, postOp, nsrcs, srcs, ndsts, dsts, offset, Nelem);
(w, nw, t, fn, numPreOpSrcs, postOp, nsrcs, srcs, ndsts, dsts, offset, Nelem);
Nrem -= Nelem;
if (Nrem == 0) return;
@@ -681,7 +687,7 @@ __device__ __forceinline__ void ReduceOrCopyMulti(
// no unroll, by type. Should finish what's remaining.
ReduceCopyMulti<FUNC, T, 1, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>
(w, nw, t, fn, preOpSrc0, postOp, nsrcs, srcs, ndsts, dsts, offset, Nrem);
(w, nw, t, fn, numPreOpSrcs, postOp, nsrcs, srcs, ndsts, dsts, offset, Nrem);
}
#endif // COMMON_KERNEL_H_