| // Copyright 2023 The IREE Authors |
| // |
| // Licensed 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 |
| |
| #include <float.h> |
| #include <hip/hip_fp16.h> |
| #include <hip/hip_runtime.h> |
| |
| extern "C" __device__ __attribute__((const)) half __ockl_wfred_max_f16(half); |
| extern "C" __device__ |
| __attribute__((const)) int64_t __ockl_wfred_min_i64(int64_t); |
| extern "C" __device__ |
| __attribute__((const)) int32_t __ockl_wfred_min_i32(int32_t); |
| |
| /* |
| Constraint/Tiling note: |
| For simplicity, we distribute all parallel dim across different workgroup, and |
| only use single subgroup/warp per workgroup. This constraint is also set during |
| tiling phase in KernelConfig. |
| */ |
| |
| extern "C" __device__ void __iree_uk_rocm_argmax_F32I32(float *inputBuffer, |
| size_t input_offset, |
| int32_t *outputBuffer, |
| size_t output_offset, |
| size_t reductionSize) { |
| uint laneID = __builtin_amdgcn_workitem_id_x(); |
| // Set identity value to handle problem non divisible by subgroupSize. |
| float laneMax = |
| laneID >= reductionSize ? -FLT_MAX : inputBuffer[input_offset + laneID]; |
| int32_t laneResult = laneID; |
| |
| // NOTE: On F32 kernels with clang, reductionSize/blockDim.x has numerical |
| // inaccuracy. |
| uint numBatches = (reductionSize + warpSize - 1) / warpSize; |
| for (int i = 1; i < numBatches; ++i) { |
| uint idx = warpSize * i + laneID; |
| float newIn = |
| idx >= reductionSize ? -FLT_MAX : inputBuffer[input_offset + idx]; |
| if (newIn == laneMax) |
| continue; |
| laneMax = __ocml_fmax_f32(newIn, laneMax); |
| laneResult = newIn == laneMax ? idx : laneResult; |
| } |
| |
| // Final reduction with one subgroup |
| // NOTE: __ockl_wfred_max_f32 has correctness issue on gfx1100 documented on |
| // https://github.com/iree-org/iree/issues/16112. |
| float wgMax = laneMax; |
| for (int i = 1; i < warpSize; i *= 2) { |
| wgMax = __ocml_fmax_f32(__shfl_xor(wgMax, i), wgMax); |
| } |
| // Check if there are multiple max value holders. |
| uint64_t laneHasMaxValmask = __ballot(wgMax == laneMax); |
| // if there is only one max value holder, write and exit. |
| if (__popcll(laneHasMaxValmask) == 1) { |
| if (wgMax == laneMax) |
| outputBuffer[output_offset] = laneResult; |
| return; |
| } |
| // if there are multiple max value holder, find smallest index (argmax |
| // semantics). |
| int32_t indexVal = wgMax == laneMax ? laneResult : __INT32_MAX__; |
| laneResult = __ockl_wfred_min_i32(indexVal); |
| if (laneID == 0) |
| outputBuffer[output_offset] = laneResult; |
| } |
| |
| extern "C" __device__ void __iree_uk_rocm_argmax_F32I64(float *inputBuffer, |
| size_t input_offset, |
| int64_t *outputBuffer, |
| size_t output_offset, |
| size_t reductionSize) { |
| uint laneID = __builtin_amdgcn_workitem_id_x(); |
| // Set identity value to handle problem non divisible by subgroupSize. |
| float laneMax = |
| laneID >= reductionSize ? -FLT_MAX : inputBuffer[input_offset + laneID]; |
| int64_t laneResult = laneID; |
| |
| // NOTE: On F32 kernels with clang, reductionSize/blockDim.x has numerical |
| // inaccuracy. |
| uint numBatches = (reductionSize + warpSize - 1) / warpSize; |
| for (int i = 1; i < numBatches; ++i) { |
| uint idx = warpSize * i + laneID; |
| float newIn = |
| idx >= reductionSize ? -FLT_MAX : inputBuffer[input_offset + idx]; |
| if (newIn == laneMax) |
| continue; |
| laneMax = __ocml_fmax_f32(newIn, laneMax); |
| laneResult = newIn == laneMax ? idx : laneResult; |
| } |
| |
| // Final reduction with one subgroup |
| // NOTE: __ockl_wfred_max_f32 has correctness issue on gfx1100 documented on |
| // https://github.com/iree-org/iree/issues/16112. |
| float wgMax = laneMax; |
| for (int i = 1; i < warpSize; i *= 2) { |
| wgMax = __ocml_fmax_f32(__shfl_xor(wgMax, i), wgMax); |
| } |
| // Check if there are multiple max value holders. |
| uint64_t laneHasMaxValmask = __ballot(wgMax == laneMax); |
| // if there is only one max value holder, write and exit. |
| if (__popcll(laneHasMaxValmask) == 1) { |
| if (wgMax == laneMax) |
| outputBuffer[output_offset] = laneResult; |
| return; |
| } |
| // if there are multiple max value holder, find smallest index (argmax |
| // semantics). |
| int64_t indexVal = wgMax == laneMax ? laneResult : __INT64_MAX__; |
| laneResult = __ockl_wfred_min_i64(indexVal); |
| if (laneID == 0) |
| outputBuffer[output_offset] = laneResult; |
| } |
| |
| extern "C" __device__ void __iree_uk_rocm_argmax_F16I32(half *inputBuffer, |
| size_t input_offset, |
| int32_t *outputBuffer, |
| size_t output_offset, |
| size_t reductionSize) { |
| half NEG_F16_MAX = __float2half(-65504.0f); |
| uint laneID = __builtin_amdgcn_workitem_id_x(); |
| // Set identity value to handle problem non divisible by subgroupSize. |
| half laneMax = laneID >= reductionSize ? NEG_F16_MAX |
| : inputBuffer[input_offset + laneID]; |
| int32_t laneResult = laneID; |
| |
| uint numBatches = (reductionSize + warpSize - 1) / warpSize; |
| for (int i = 1; i < numBatches; ++i) { |
| uint idx = warpSize * i + laneID; |
| half newIn = |
| idx >= reductionSize ? NEG_F16_MAX : inputBuffer[input_offset + idx]; |
| if (newIn == laneMax) |
| continue; |
| laneMax = __ocml_fmax_f16(newIn, laneMax); |
| laneResult = newIn == laneMax ? idx : laneResult; |
| } |
| |
| // Final reduction with one subgroup |
| half wgMax = __ockl_wfred_max_f16(laneMax); |
| // Check if there are multiple max value holders. |
| uint64_t laneHasMaxValmask = __ballot(wgMax == laneMax); |
| // if there is only one max value holder, write and exit. |
| if (__popcll(laneHasMaxValmask) == 1) { |
| if (wgMax == laneMax) |
| outputBuffer[output_offset] = laneResult; |
| return; |
| } |
| // if there are multiple max value holder, find smallest index (argmax |
| // semantics). |
| int32_t indexVal = wgMax == laneMax ? laneResult : __INT32_MAX__; |
| laneResult = __ockl_wfred_min_i32(indexVal); |
| if (laneID == 0) |
| outputBuffer[output_offset] = laneResult; |
| } |
| |
| extern "C" __device__ void __iree_uk_rocm_argmax_F16I64(half *inputBuffer, |
| size_t input_offset, |
| int64_t *outputBuffer, |
| size_t output_offset, |
| size_t reductionSize) { |
| half NEG_F16_MAX = __float2half(-65504.0f); |
| uint laneID = __builtin_amdgcn_workitem_id_x(); |
| // Set identity value to handle problem non divisible by subgroupSize. |
| half laneMax = laneID >= reductionSize ? NEG_F16_MAX |
| : inputBuffer[input_offset + laneID]; |
| int64_t laneResult = laneID; |
| |
| uint numBatches = (reductionSize + warpSize - 1) / warpSize; |
| for (int i = 1; i < numBatches; ++i) { |
| uint idx = warpSize * i + laneID; |
| half newIn = |
| idx >= reductionSize ? NEG_F16_MAX : inputBuffer[input_offset + idx]; |
| if (newIn == laneMax) |
| continue; |
| laneMax = __ocml_fmax_f16(newIn, laneMax); |
| laneResult = newIn == laneMax ? idx : laneResult; |
| } |
| |
| // Final reduction with one subgroup |
| half wgMax = __ockl_wfred_max_f16(laneMax); |
| // Check if there are multiple max value holders. |
| uint64_t laneHasMaxValmask = __ballot(wgMax == laneMax); |
| // if there is only one max value holder, write and exit. |
| if (__popcll(laneHasMaxValmask) == 1) { |
| if (wgMax == laneMax) |
| outputBuffer[output_offset] = laneResult; |
| return; |
| } |
| // if there are multiple max value holder, find smallest index (argmax |
| // semantics). |
| int64_t indexVal = wgMax == laneMax ? laneResult : __INT64_MAX__; |
| laneResult = __ockl_wfred_min_i64(indexVal); |
| if (laneID == 0) |
| outputBuffer[output_offset] = laneResult; |
| } |
