Skip to content

ggml-cuda: HIP replace __shfl_xor_sync with dpp instructions #26466

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
thelittlefireman wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

ggml-cuda: HIP replace __shfl_xor_sync with dpp instructions #26466

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions ggml/src/ggml-cuda/argmax.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,8 @@ static __global__ void argmax_f32(const float * __restrict__ x, int32_t * __rest

#pragma unroll
for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) {
const float val = __shfl_xor_sync(0xFFFFFFFF, maxval, offset, WARP_SIZE);
const int col = __shfl_xor_sync(0xFFFFFFFF, argmax, offset, WARP_SIZE);
const float val = ggml_cuda_shfl_xor_sync(maxval, offset);
const int col = ggml_cuda_shfl_xor_sync(argmax, offset);
if (val > maxval) {
maxval = val;
argmax = col;
Expand Down Expand Up @@ -51,8 +51,8 @@ static __global__ void argmax_f32(const float * __restrict__ x, int32_t * __rest
}
#pragma unroll
for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) {
const float val = __shfl_xor_sync(0xFFFFFFFF, maxval, offset, WARP_SIZE);
const int col = __shfl_xor_sync(0xFFFFFFFF, argmax, offset, WARP_SIZE);
const float val = ggml_cuda_shfl_xor_sync(maxval, offset);
const int col = ggml_cuda_shfl_xor_sync(argmax, offset);
if (val > maxval) {
maxval = val;
argmax = col;
Expand Down
75 changes: 66 additions & 9 deletions ggml/src/ggml-cuda/common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -371,6 +371,63 @@ static bool ggml_cuda_is_aligned(const ggml_tensor * tensor, const size_t alignm
tensor->nb[3] % alignment == 0;
}

#ifdef GGML_USE_HIP
template <int dpp_ctrl, typename T, int row_mask = 0xf, int bank_mask = 0xf, bool bound_ctrl = true>
static __device__ __forceinline__ T hip_move_dpp(T old, T v) {
return __builtin_bit_cast(
T,
__builtin_amdgcn_update_dpp(
__builtin_bit_cast(int, old),
__builtin_bit_cast(int, v),
dpp_ctrl,
row_mask,
bank_mask,
bound_ctrl
)
);
}

template <int mask, typename T>
static __device__ __forceinline__ T hip_ds_swizzle(T v) {
return __builtin_bit_cast(T, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, v), mask));
}
#endif // GGML_USE_HIP

template<int width = WARP_SIZE, typename T>
static __device__ __forceinline__ T ggml_cuda_shfl_xor_sync(T x, int offset) {

#if defined(GGML_USE_HIP)
#if defined(__GFX9__)
static T old;
// clang (v20) will not unroll loops with just the plain `offset` in switch
switch (~offset) {
// subgroups (width) should not make a difference for a butterfly shuffle pattern
case ~1: return hip_move_dpp<0xB1>(old, x); // quad_perm:[1,0,3,2]
case ~2: return hip_move_dpp<0x4E>(old, x); // quad_perm:[2,3,0,1]
case ~4: return hip_ds_swizzle<0x101F>(x); // ds_swizzle AND mask = 0x1F; OR mask = 0; XOR mask = 4
case ~8: return hip_move_dpp<0x128>(old, x); // row_ror:8
case ~16: return hip_ds_swizzle<0x401f>(x); // swap neighboring groups of 16
default: return __shfl_xor(x, offset, width);
}
#else
static T old;
// clang (v20) will not unroll loops with just the plain `offset` in switch
switch (~offset) {
// subgroups (width) should not make a difference for a butterfly shuffle pattern
case ~1: return hip_move_dpp<0x160 + 1>(old, x); // row_xor_mask: offset
case ~2: return hip_move_dpp<0x160 + 2>(old, x);
case ~4: return hip_move_dpp<0x160 + 4>(old, x);
case ~8: return hip_move_dpp<0x160 + 8>(old, x);
case ~16: return hip_ds_swizzle<0x401f>(x); // swap neighboring groups of 16
default: return __shfl_xor(x, offset, width);
}
#endif // GCN
#else
return __shfl_xor_sync(0xffffffff, x, offset, width);
#endif // defined(GGML_USE_HIP)
}


static constexpr __device__ int ggml_cuda_get_physical_warp_size() {
#if defined(GGML_USE_HIP) && (defined(__GFX9__) || defined(__GFX8__))
return 64;
Expand Down Expand Up @@ -446,7 +503,7 @@ static __device__ __forceinline__ int warp_reduce_sum(int x) {
#else
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
x += __shfl_xor_sync(0xffffffff, x, offset, width);
x += ggml_cuda_shfl_xor_sync<width>(x, offset);
}
return x;
#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
Expand All @@ -456,7 +513,7 @@ template<int width = WARP_SIZE>
static __device__ __forceinline__ float warp_reduce_sum(float x) {
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
x += __shfl_xor_sync(0xffffffff, x, offset, width);
x += ggml_cuda_shfl_xor_sync<width>(x, offset);
}
return x;
}
Expand All @@ -465,8 +522,8 @@ template<int width = WARP_SIZE>
static __device__ __forceinline__ float2 warp_reduce_sum(float2 a) {
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
a.x += __shfl_xor_sync(0xffffffff, a.x, offset, width);
a.y += __shfl_xor_sync(0xffffffff, a.y, offset, width);
a.x += ggml_cuda_shfl_xor_sync<width>(a.x, offset);
a.y += ggml_cuda_shfl_xor_sync<width>(a.y, offset);
}
return a;
}
Expand All @@ -476,7 +533,7 @@ static __device__ __forceinline__ half2 warp_reduce_sum(half2 a) {
#ifdef FP16_AVAILABLE
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
a = __hadd2(a, __shfl_xor_sync(0xffffffff, a, offset, width));
a = __hadd2(a, ggml_cuda_shfl_xor_sync<width>(a, offset));
}
return a;

Expand All @@ -493,7 +550,7 @@ static __device__ __forceinline__ int warp_reduce_all(int x) {
} else {
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
x = __shfl_xor_sync(0xffffffff, x, offset, width) && x;
x = ggml_cuda_shfl_xor_sync<width>(x, offset) && x;
}
return x;
}
Expand All @@ -506,7 +563,7 @@ static __device__ __forceinline__ int warp_reduce_any(int x) {
} else {
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
x = __shfl_xor_sync(0xffffffff, x, offset, width) || x;
x = ggml_cuda_shfl_xor_sync<width>(x, offset) || x;
}
return x;
}
Expand All @@ -516,7 +573,7 @@ template<int width = WARP_SIZE>
static __device__ __forceinline__ float warp_reduce_max(float x) {
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset, width));
x = fmaxf(x, ggml_cuda_shfl_xor_sync<width>(x, offset));
}
return x;
}
Expand Down Expand Up @@ -682,7 +739,7 @@ static __device__ __forceinline__ half2 warp_reduce_max(half2 x) {
#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL || defined(GGML_USE_HIP)
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
x = ggml_cuda_hmax2(x, __shfl_xor_sync(0xffffffff, x, offset, width));
x = ggml_cuda_hmax2(x, ggml_cuda_shfl_xor_sync<width>(x, offset));
}
return x;
#else
Expand Down
4 changes: 2 additions & 2 deletions ggml/src/ggml-cuda/fattn-common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -347,8 +347,8 @@ static __device__ __forceinline__ void quantize_q8_1_to_shared(
}
#pragma unroll
for (int mask = QI8_1/2; mask > 0; mask >>= 1) {
amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, 32));
sum += __shfl_xor_sync(0xFFFFFFFF, sum, mask, 32);
amax = fmaxf(amax, ggml_cuda_shfl_xor_sync<32>(amax, mask));
sum += ggml_cuda_shfl_xor_sync<32>(sum, mask);
}

const float d = amax / 127;
Expand Down
10 changes: 5 additions & 5 deletions ggml/src/ggml-cuda/fattn-mma-f16.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -730,7 +730,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
for (int col = 0; col < cols_per_thread; ++col) {
#pragma unroll
for (int offset = 16; offset >= 4; offset >>= 1) {
KQ_max_new[col] = fmaxf(KQ_max_new[col], __shfl_xor_sync(0xFFFFFFFF, KQ_max_new[col], offset, warp_size));
KQ_max_new[col] = fmaxf(KQ_max_new[col], ggml_cuda_shfl_xor_sync(KQ_max_new[col], offset));
}
}

Expand Down Expand Up @@ -823,7 +823,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
#endif // defined(TURING_MMA_AVAILABLE)
#pragma unroll
for (int offset = offset_first; offset >= offset_last; offset >>= 1) {
KQ_max_new[col] = fmaxf(KQ_max_new[col], __shfl_xor_sync(0xFFFFFFFF, KQ_max_new[col], offset, warp_size));
KQ_max_new[col] = fmaxf(KQ_max_new[col], ggml_cuda_shfl_xor_sync(KQ_max_new[col], offset));
}
}

Expand Down Expand Up @@ -1341,7 +1341,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
for (int col = 0; col < cols_per_thread; ++col) {
#pragma unroll
for (int offset = offset_first; offset >= offset_last; offset >>= 1) {
KQ_rowsum[col] += __shfl_xor_sync(0xFFFFFFFF, KQ_rowsum[col], offset, warp_size);
KQ_rowsum[col] += ggml_cuda_shfl_xor_sync(KQ_rowsum[col], offset);
}
}
}
Expand Down Expand Up @@ -1513,7 +1513,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
#pragma unroll
for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) {
if (offset < warp_size) {
KQ_cmn = fmaxf(KQ_cmn, __shfl_xor_sync(0xFFFFFFFF, KQ_cmn, offset, warp_size));
KQ_cmn = fmaxf(KQ_cmn, ggml_cuda_shfl_xor_sync(KQ_cmn, offset));
}
}

Expand All @@ -1531,7 +1531,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
#pragma unroll
for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) {
if (offset < warp_size) {
KQ_crs += __shfl_xor_sync(0xFFFFFFFF, KQ_crs, offset, warp_size);
KQ_crs += ggml_cuda_shfl_xor_sync(KQ_crs, offset);
}
}

Expand Down
2 changes: 1 addition & 1 deletion ggml/src/ggml-cuda/fattn-vec.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -293,7 +293,7 @@ static __global__ void flash_attn_ext_vec(
for (int j = 0; j < ncols; ++j) {
#pragma unroll
for (int offset = nthreads_KQ; offset < WARP_SIZE; offset <<= 1) {
KQ_max_new[j] = fmaxf(KQ_max_new[j], __shfl_xor_sync(0xFFFFFFFF, KQ_max_new[j], offset, WARP_SIZE));
KQ_max_new[j] = fmaxf(KQ_max_new[j], ggml_cuda_shfl_xor_sync(KQ_max_new[j], offset));
}
const float KQ_max_scale = expf(KQ_max[j] - KQ_max_new[j]);
KQ_max[j] = KQ_max_new[j];
Expand Down
2 changes: 1 addition & 1 deletion ggml/src/ggml-cuda/fwht.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ __global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows,
#pragma unroll
for (int j = 0; j < el_w; j++) {
const float val = reg[j];
const float val2 = __shfl_xor_sync(0xFFFFFFFF, val, h, warp_size);
const float val2 = ggml_cuda_shfl_xor_sync<warp_size>(val, h);

reg[j] = (lane & h) == 0 ? val + val2 : val2 - val;
}
Expand Down
6 changes: 3 additions & 3 deletions ggml/src/ggml-cuda/mma.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -736,7 +736,7 @@ namespace ggml_cuda_mma {
int i = threadIdx.x / 16;
tmp[i] = tile_float.x[l];
i ^= 1;
tmp[i] = __shfl_xor_sync(0xFFFFFFFF, tile_float.x[l], 16, WARP_SIZE);
tmp[i] = ggml_cuda_shfl_xor_sync(tile_float.x[l], 16);
ret.x[l] = make_half2(tmp[0], tmp[1]);
}
return ret;
Expand Down Expand Up @@ -767,8 +767,8 @@ namespace ggml_cuda_mma {

// On Volta FP16 and FP32 tiles have a different memory layout,
// for the conversion threads with an offset of 2 need to exchange half their values:
ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)] = __shfl_xor_sync(
0xFFFFFFFF, ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)], 2, WARP_SIZE);
ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)] = ggml_cuda_shfl_xor_sync(
ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)], 2);
}
return ret;
}
Expand Down
6 changes: 3 additions & 3 deletions ggml/src/ggml-cuda/quantize.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -390,7 +390,7 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
float amax = fabsf(xi);
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, WARP_SIZE));
amax = fmaxf(amax, ggml_cuda_shfl_xor_sync(amax, mask));
}

const uint8_t e = compute_e8m0_scale(amax);
Expand Down Expand Up @@ -498,7 +498,7 @@ static __global__ void quantize_mmq_q8_1(
// Exchange max. abs. value between vals_per_scale/4 threads.
#pragma unroll
for (int offset = vals_per_scale/8; offset > 0; offset >>= 1) {
amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE));
amax = fmaxf(amax, ggml_cuda_shfl_xor_sync(amax, offset));
}

float sum;
Expand All @@ -508,7 +508,7 @@ static __global__ void quantize_mmq_q8_1(
// Calculate sums across vals_per_sum/4 threads.
#pragma unroll
for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE);
sum += ggml_cuda_shfl_xor_sync(sum, offset);
}
}

Expand Down
10 changes: 5 additions & 5 deletions ggml/src/ggml-cuda/topk-moe.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,9 +195,9 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *

#pragma unroll
for (int mask = WARP_SIZE / 2; mask > 0; mask /= 2) {
const float val = __shfl_xor_sync(0xFFFFFFFF, max_val, mask, WARP_SIZE);
const float val_s = __shfl_xor_sync(0xFFFFFFFF, max_val_s, mask, WARP_SIZE);
const int expert = __shfl_xor_sync(0xFFFFFFFF, max_expert, mask, WARP_SIZE);
const float val = ggml_cuda_shfl_xor_sync(max_val, mask);
const float val_s = ggml_cuda_shfl_xor_sync(max_val_s, mask);
const int expert = ggml_cuda_shfl_xor_sync(max_expert, mask);
if (val_s > max_val_s || (val_s == max_val_s && expert < max_expert)) {
max_val = val;
max_val_s = val_s;
Expand All @@ -220,8 +220,8 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *

#pragma unroll
for (int mask = WARP_SIZE / 2; mask > 0; mask /= 2) {
const float val = __shfl_xor_sync(0xFFFFFFFF, max_val, mask, WARP_SIZE);
const int expert = __shfl_xor_sync(0xFFFFFFFF, max_expert, mask, WARP_SIZE);
const float val = ggml_cuda_shfl_xor_sync(max_val, mask);
const int expert = ggml_cuda_shfl_xor_sync(max_expert, mask);
if (val > max_val || (val == max_val && expert < max_expert)) {
max_val = val;
max_expert = expert;
Expand Down
1 change: 0 additions & 1 deletion ggml/src/ggml-cuda/vendors/hip.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@
#define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }}
#define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width)
#define __shfl_up_sync(mask, var, laneMask, width) __shfl_up(var, laneMask, width)
#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width)
#define __all_sync(mask, var) __all(var)
#define __any_sync(mask, var) __any(var)
#define cublasStrsmBatched hipblasStrsmBatched
Expand Down