warp_shuffle_op: avoid sub-int narrow↔wide roundtrip (IGC fp16 miscompile)

.ci/exclude.txt

@@ -14,15 +14,10 @@ rocprim\.device_reduce_by_key
 # Mitigation: CHIP_L0_FORCE_SIMD16 or wait for IGC fix.
 rocprim\.device_run_length_encode
 
-# IGC bug #402 (fp16 subgroup-shuffle narrowing miscompile).
-# Upstream: https://github.com/intel/intel-graphics-compiler/issues/402
-# Fix pending in PR [CHIP-SPV/rocPRIM#3]; remove these two lines when #3 merges.
-rocprim\.warp_reduce
-rocprim\.warp_scan
-
-# fp16/bf16 block_shuffle failures on PoCL CPU (pastrami). Related to the same
-# sub-int warp_shuffle roundtrip addressed by PR [CHIP-SPV/rocPRIM#3]; remove
-# once the fp16 fix lands and is verified on CPU.
+# IGC bug #402 fp16 tests — warp_reduce + warp_scan fixed by PR #3.
+# https://github.com/intel/intel-graphics-compiler/issues/402
+# block_shuffle bf16 still fails on PoCL CPU (pastrami) — separate code path
+# not covered by warp_shuffle_op fix.
 rocprim\.block_shuffle
 
 # Mali GPU (salami) lacks fp64 + SubgroupDispatch SPIR-V capabilities.

rocprim/include/rocprim/intrinsics/warp_shuffle.hpp

@@ -79,9 +79,53 @@ warp_shuffle_op(const T& input, ShuffleOp&& op)
 #endif
 }
 
+#if defined(__HIP_PLATFORM_SPIRV__) && !defined(__HIP_CPU_RT__)
+// Specialization for sub-int sized trivially-copyable types (e.g. __half).
+//
+// Works around an IGC miscompile on Intel Arc: the SPIR-V translator emits
+//
+//   %w = OpUConvert %uint %ushort_input      ; narrow -> wide
+//   %s = OpSubgroupShuffleINTEL %uint %w ...
+//   %n = OpUConvert %ushort %s               ; wide -> narrow
+//
+// and IGC miscompiles the post-shuffle narrowing UConvert when the consumer
+// lives in a block entered via OpBranchConditional (returns 0 instead of
+// the shuffled value). See upstream IGC issue (to be filed).
+//
+// We defeat the trigger by bouncing the shuffle result through a
+// memory-backed scratch slot and loading the low sizeof(T) bytes via a
+// byte-wise memcpy from a volatile pointer. This emits OpBitcast + OpLoad
+// of the narrow type instead of a scalar-register OpUConvert, which IGC
+// handles correctly.
 template<class T, class ShuffleOp>
 ROCPRIM_DEVICE ROCPRIM_INLINE
-typename std::enable_if<!(std::is_trivially_copyable<T>::value && (sizeof(T) % sizeof(int) == 0)), T>::type
+typename std::enable_if<
+    std::is_trivially_copyable<T>::value && (sizeof(T) < sizeof(int)), T>::type
+warp_shuffle_op(const T& input, ShuffleOp&& op)
+{
+    int word = 0; // zero-init upper bytes that the memcpy below leaves untouched
+    __builtin_memcpy(&word, &input, sizeof(T));
+    word = op(word);
+
+    // Store the shuffle result into a volatile stack slot to force the
+    // narrowing to happen via a memory load rather than a register-level
+    // i32->i16 trunc (which the translator lowers to OpUConvert).
+    volatile int word_scratch = word;
+    T output;
+    __builtin_memcpy(&output,
+                     const_cast<const int*>(&word_scratch),
+                     sizeof(T));
+    return output;
+}
+#endif
+
+template<class T, class ShuffleOp>
+ROCPRIM_DEVICE ROCPRIM_INLINE
+typename std::enable_if<!(std::is_trivially_copyable<T>::value && (sizeof(T) % sizeof(int) == 0))
+#if defined(__HIP_PLATFORM_SPIRV__) && !defined(__HIP_CPU_RT__)
+                        && !(std::is_trivially_copyable<T>::value && (sizeof(T) < sizeof(int)))
+#endif
+                        , T>::type
 warp_shuffle_op(const T& input, ShuffleOp&& op)
 {
     constexpr int words_no = (sizeof(T) + sizeof(int) - 1) / sizeof(int);

test/rocprim/CMakeLists.txt

@@ -383,6 +383,7 @@ add_rocprim_test("rocprim.warp_reduce" test_warp_reduce.cpp)
 add_rocprim_test("rocprim.warp_scan" test_warp_scan.cpp)
 add_rocprim_test("rocprim.warp_sort" test_warp_sort.cpp)
 add_rocprim_test("rocprim.warp_store" test_warp_store.cpp)
+add_rocprim_test("rocprim.warp_shuffle_fp16" test_warp_shuffle_fp16.cpp)
 
 # GPU: 3/4 pass, TransformReduce fail.
 # CPU: likely stuck

test/rocprim/test_warp_shuffle_fp16.cpp

@@ -0,0 +1,119 @@
+// MIT License
+//
+// Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+// Regression test for IGC fp16 narrow<->wide shuffle miscompile.
+//
+// Reproduces the pattern emitted by rocprim::warp_shuffle_op for
+// sizeof(T) < sizeof(int): the SPIR-V translator inserts an
+// OpUConvert %uint <- %ushort before OpSubgroupShuffleINTEL and a
+// matching OpUConvert %ushort <- %uint afterwards. When the result
+// is consumed inside a divergent block (OpBranchConditional), IGC
+// on Intel Arc miscompiles the chain and the second shuffle yields 0.
+
+#include "../common_test_header.hpp"
+#include "test_utils.hpp"
+
+#include <rocprim/intrinsics/warp_shuffle.hpp>
+#include <rocprim/types.hpp>
+
+// Two chained warp_shuffles on __half followed by a divergent store.
+// This is the minimum shape that triggers the IGC bug; a single
+// unconditional shuffle passes.
+template<unsigned int LogicalWarpSize>
+__global__ __launch_bounds__(64)
+void chained_half_shuffle_divergent_kernel(const __half* in, __half* out)
+{
+    const unsigned int tid = threadIdx.x;
+    __half v = in[tid];
+
+    // Mirror the rocprim::warp_reduce_shuffle inner loop:
+    // for(offset=1; offset<WarpSize; offset*=2) value = shuffle_down(output, offset);
+    // This is the exact shape that emits `UConvert uint <- ushort` before and
+    // `UConvert ushort <- uint` after OpSubgroupShuffleINTEL and triggers the
+    // IGC miscompile on Arc.
+    __half output = v;
+    __half s1 = rocprim::warp_shuffle_down(output, 1, LogicalWarpSize);
+    output = output + s1;
+
+    __half s2 = rocprim::warp_shuffle_down(output, 2, LogicalWarpSize);
+    output = output + s2;
+
+    __half a2 = output;
+
+    // Divergent consumer: only lane 0 of every logical warp writes.
+    if ((tid % LogicalWarpSize) == 0)
+    {
+        out[tid / LogicalWarpSize] = a2;
+    }
+}
+
+TEST(RocprimWarpShuffleFp16, ChainedDivergentHalfShuffle)
+{
+    int device_id = test_common_utils::obtain_device_from_ctest();
+    HIP_CHECK(hipSetDevice(device_id));
+
+    constexpr unsigned int LogicalWarpSize = 4;
+    constexpr unsigned int BlockSize       = 64;
+    constexpr unsigned int NumWarps        = BlockSize / LogicalWarpSize;
+
+    std::vector<__half> h_in(BlockSize);
+    for(unsigned int i = 0; i < BlockSize; ++i)
+    {
+        // Values 2,3,4,5 per 4-lane group -> expected sum = 14.
+        h_in[i] = __half(static_cast<float>((i % LogicalWarpSize) + 2));
+    }
+    std::vector<__half> h_out(NumWarps, __half(0.0f));
+
+    __half* d_in  = nullptr;
+    __half* d_out = nullptr;
+    HIP_CHECK(hipMalloc(&d_in,  h_in.size()  * sizeof(__half)));
+    HIP_CHECK(hipMalloc(&d_out, h_out.size() * sizeof(__half)));
+    HIP_CHECK(hipMemcpy(d_in,  h_in.data(),  h_in.size()  * sizeof(__half),
+                        hipMemcpyHostToDevice));
+    HIP_CHECK(hipMemcpy(d_out, h_out.data(), h_out.size() * sizeof(__half),
+                        hipMemcpyHostToDevice));
+
+    hipLaunchKernelGGL(HIP_KERNEL_NAME(chained_half_shuffle_divergent_kernel<LogicalWarpSize>),
+                       dim3(1), dim3(BlockSize), 0, 0, d_in, d_out);
+    HIP_CHECK(hipGetLastError());
+    HIP_CHECK(hipDeviceSynchronize());
+
+    HIP_CHECK(hipMemcpy(h_out.data(), d_out, h_out.size() * sizeof(__half),
+                        hipMemcpyDeviceToHost));
+
+    constexpr float expected = 2.0f + 3.0f + 4.0f + 5.0f; // 14
+    unsigned int mismatches = 0;
+    for(unsigned int w = 0; w < NumWarps; ++w)
+    {
+        const float got = static_cast<float>(h_out[w]);
+        if(std::abs(got - expected) > 0.01f)
+        {
+            ++mismatches;
+        }
+    }
+    EXPECT_EQ(mismatches, 0u)
+        << "IGC fp16 narrow<->wide chained-shuffle miscompile triggered "
+        << "(got 0 for second shuffle under divergent consumer).";
+
+    HIP_CHECK(hipFree(d_in));
+    HIP_CHECK(hipFree(d_out));
+}