RDNA4 phase 1: make HIP_ARCH work, attention wave-size helpers

AGENTS.md

@@ -29,7 +29,7 @@ ctest -R '^utest_softmax_cu$'                    # or via ctest, from build/
 ctest -R '_cpp$'                                 # CPU-backend tests only (no GPU needed at runtime)
 ```
 
-Key CMake options: `ENABLE_CPU` (ON), `ENABLE_CPU_MULTITHREAD` (ON if the `thread-pool` submodule is checked out, OFF otherwise; forcing it ON without the submodule is a `FATAL_ERROR`; when ON the FKL target defines `FK_CPU_MULTITHREAD_ENABLED`, adds `thread-pool/include` (installed as `include/fused_kernel/thread_pool`) and links `Threads::Threads`), `ENABLE_CUDA` (ON if nvcc found), `ENABLE_HIP` (ON if amdclang is found), `BUILD_TEST` (ON), `BUILD_UTEST` (ON), `ENABLE_BENCHMARK` (OFF), `CUDA_ARCH` ("native", passed verbatim to `CUDA_ARCHITECTURES` — no arch filtering exists), gfx1200 by default for HIP, `ARCH_FLAGS` (CPU SIMD), `ENABLE_NVTX`, `ENABLE_DEBUG`, `TEMPLATE_DEPTH` (1000).
+Key CMake options: `ENABLE_CPU` (ON), `ENABLE_CPU_MULTITHREAD` (ON if the `thread-pool` submodule is checked out, OFF otherwise; forcing it ON without the submodule is a `FATAL_ERROR`; when ON the FKL target defines `FK_CPU_MULTITHREAD_ENABLED`, adds `thread-pool/include` (installed as `include/fused_kernel/thread_pool`) and links `Threads::Threads`), `ENABLE_CUDA` (ON if nvcc found), `ENABLE_HIP` (ON if amdclang is found), `BUILD_TEST` (ON), `BUILD_UTEST` (ON), `ENABLE_BENCHMARK` (OFF), `CUDA_ARCH` ("native", passed verbatim to `CUDA_ARCHITECTURES` — no arch filtering exists), HIP: the build machine's GPU by default (`CMAKE_HIP_ARCHITECTURES` as detected by CMake), or `HIP_ARCH` (e.g. `gfx1200;gfx1201`, RX 9060 and RX 9070, which can't load each other's code objects), which seeds `CMAKE_HIP_ARCHITECTURES` on the first configure of a build directory; an explicit `CMAKE_HIP_ARCHITECTURES` wins, and `hip::device` follows it through `GPU_TARGETS`, `ARCH_FLAGS` (CPU SIMD), `ENABLE_NVTX`, `ENABLE_DEBUG`, `TEMPLATE_DEPTH` (1000).
 
 `cmake --install build --prefix <dir>` installs the headers (`include/fused_kernel`, plus `include/fused_kernel/thread_pool` with the backend) and the `FKL` package (`share/cmake/FKL`): consumers use `find_package(FKL)` and link `FKL::FKL`, as with `add_subdirectory`. `FKL::FKL` requires C++20 with the `cxx_std_20` compile feature, which CMake also applies to CUDA and HIP sources that don't set `CUDA_STANDARD`/`HIP_STANDARD` (never add `cuda_std_20`/`hip_std_20`: with `add_subdirectory` from a C++-only project they fail with "No known features for CUDA compiler"); the project using FKL must enable CXX, and both the root `CMakeLists.txt` (as a subproject) and the package configuration stop with a clear message otherwise. Without `-DCMAKE_INSTALL_PREFIX` the install goes to `<build>/install`. `lib/export/` is not the installed package: it is the package that projects with FKL as a git submodule (cvGPUSpeedup, FastNPP) find without building FKL, and configure keeps its version file up to date; don't remove it while they use it. The package configuration template is `cmake/FKLConfig.cmake.in`, which `find_dependency(Threads)` when the backend is enabled; the version file accepts the same minor version before 1.0 (`find_package(FKL 0.2)` finds only 0.2.x) and the same major version from 1.0 on.
 

cmake/libs/hip/archs.cmake

@@ -1,9 +1,4 @@
-set(HIP_ARCH "gfx1200" CACHE STRING "HIP architecture(s) to build")
-
+# HIP_ARCH and the CMAKE_HIP_ARCHITECTURES it seeds are set in hip_init.cmake, before enable_language(HIP).
 function(set_target_hip_arch_flags TARGET_NAME)
-    if(DEFINED CMAKE_HIP_ARCHITECTURES)
-        set_target_properties(${TARGET_NAME} PROPERTIES HIP_ARCHITECTURES "${CMAKE_HIP_ARCHITECTURES}")
-    else()
-        set_target_properties(${TARGET_NAME} PROPERTIES HIP_ARCHITECTURES "${HIP_ARCH}")
-    endif()
+    set_target_properties(${TARGET_NAME} PROPERTIES HIP_ARCHITECTURES "${CMAKE_HIP_ARCHITECTURES}")
 endfunction()

cmake/libs/hip/hip.cmake

@@ -1,6 +1,12 @@
 include(cmake/libs/hip/debug.cmake)
 include(cmake/libs/hip/target_generation.cmake)
 
+# hip::device compiles the C++ sources of a target with --offload-arch=${GPU_TARGETS}, and hip-config sets GPU_TARGETS
+# to the GPU of this machine when it is empty: make it follow CMAKE_HIP_ARCHITECTURES, as the HIP sources do.
+# hip-config reads it through the cache variable GPU_BUILD_TARGETS, which it only creates, so drop that variable to
+# update it on a reconfigure.
+set(GPU_TARGETS "${CMAKE_HIP_ARCHITECTURES}")
+unset(GPU_BUILD_TARGETS CACHE)
 find_package(HIP CONFIG REQUIRED)
 option(ENABLE_HIP_LINE_INFO "Enable line info for HIP kernels compilation" ON)
 

cmake/libs/hip/hip_init.cmake

@@ -9,6 +9,14 @@ if (NOT CMAKE_HIP_COMPILER)
 endif()
 
 set(CMAKE_CXX_SCAN_FOR_MODULES OFF) #clang with HIP does not come with  clang-scan-deps
+
+# By default enable_language(HIP) sets CMAKE_HIP_ARCHITECTURES to the GPU of this machine (or the compiler default).
+# -DHIP_ARCH=<list> (e.g. "gfx1200;gfx1201": Navi 44 and Navi 48 can't load each other's code objects) seeds it instead
+# on the first configure of a build directory; an explicit -DCMAKE_HIP_ARCHITECTURES wins over HIP_ARCH.
+
+if(DEFINED HIP_ARCH AND NOT DEFINED CACHE{CMAKE_HIP_ARCHITECTURES})
+  set(CMAKE_HIP_ARCHITECTURES "${HIP_ARCH}" CACHE STRING "HIP architectures")
+endif()
 enable_language(HIP)
 
 include(cmake/libs/hip/hip.cmake)

include/fused_kernel/algorithms/attention/attention_wave.h

@@ -0,0 +1,69 @@
+/* Copyright 2026 Oscar Amoros Huguet
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License. */
+
+#ifndef FK_ATTENTION_ATTENTION_WAVE_H
+#define FK_ATTENTION_ATTENTION_WAVE_H
+
+/* Wave (warp) size of the attention DPPs: they map one row, or one fragment
+ * row group, onto exactly 32 lanes (NVIDIA warps, and RDNA waves compiled as
+ * wave32, the HIP default on gfx12). Wave64 is not supported.
+ *
+ * ROCm's clang no longer exposes the wave size at compile time
+ * (__AMDGCN_WAVEFRONT_SIZE__ was removed, and neither warpSize nor
+ * __builtin_amdgcn_wavefrontsize() is a constant expression), so:
+ *  - with a compiler that still defines __AMDGCN_WAVEFRONT_SIZE__, a wave64
+ *    device compilation fails here with #error;
+ *  - otherwise the kernels call attnRequireWave32(), which the optimizer folds
+ *    away in a wave32 build and turns into a trap in a wave64 one. */
+
+#include <fused_kernel/core/utils/utils.h>
+
+namespace fk {
+
+constexpr int ATTN_WAVE_SIZE = 32;
+
+#if defined(__HIPCC__) && defined(__HIP_DEVICE_COMPILE__) && defined(__AMDGCN_WAVEFRONT_SIZE__)
+#if __AMDGCN_WAVEFRONT_SIZE__ != 32
+#error "The attention DPPs need wave32: don't compile them with -mwavefrontsize64"
+#endif
+#endif
+
+#if defined(__NVCC__) || defined(__HIPCC__)
+
+// Traps when the kernel runs as wave64. Compiles to nothing in a wave32 (and in
+// any CUDA) build.
+__device__ __forceinline__ void attnRequireWave32() {
+#if defined(__HIPCC__) && defined(__HIP_DEVICE_COMPILE__)
+    if (__builtin_amdgcn_wavefrontsize() != ATTN_WAVE_SIZE) {
+        __builtin_trap();
+    }
+#endif
+}
+
+// Butterfly exchange across the 32 lanes of the wave: returns the value of
+// lane (lane ^ laneMask).
+template <typename T>
+__device__ __forceinline__ T attnShflXor(const T value, const int laneMask) {
+#if defined(__HIPCC__)
+    return __shfl_xor(value, laneMask, ATTN_WAVE_SIZE);
+#else
+    return __shfl_xor_sync(0xffffffffu, value, laneMask);
+#endif
+}
+
+#endif // defined(__NVCC__) || defined(__HIPCC__)
+
+} // namespace fk
+
+#endif // FK_ATTENTION_ATTENTION_WAVE_H

utests/algorithm/attention/utest_attention_wave.h

@@ -0,0 +1,95 @@
+/* Copyright 2026 Oscar Amoros Huguet
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License. */
+
+/* attention_wave.h: attnShflXor exchanges values inside each 32-lane wave
+ * (never across waves) and attnRequireWave32 lets a wave32 kernel run. Two
+ * waves per block: each lane holds its thread index, a butterfly sum must give
+ * the sum of its own wave, and a single exchange must return lane ^ mask. */
+
+#include <tests/main.h>
+
+#include <fused_kernel/algorithms/attention/attention_wave.h>
+#include <fused_kernel/core/data/ptr_nd.h>
+#include <fused_kernel/core/execution_model/stream.h>
+
+#include <iostream>
+
+#if defined(__NVCC__) || defined(__HIPCC__)
+
+namespace fk {
+
+constexpr int WAVE_TEST_THREADS = 2 * ATTN_WAVE_SIZE;
+constexpr int WAVE_TEST_MASK = 5;
+
+__global__ void attnWaveTestKernel(const RawPtr<ND::_1D, float> sums, const RawPtr<ND::_1D, int> exchanged) {
+    attnRequireWave32();
+    const int tid = static_cast<int>(threadIdx.x);
+    float sum = static_cast<float>(tid);
+    for (int offset = ATTN_WAVE_SIZE / 2; offset > 0; offset >>= 1) {
+        sum += attnShflXor(sum, offset);
+    }
+    sums.data[tid] = sum;
+    exchanged.data[tid] = attnShflXor(tid, WAVE_TEST_MASK);
+}
+
+int launch_impl() {
+    Ptr1D<float> sums(WAVE_TEST_THREADS);
+    Ptr1D<int> exchanged(WAVE_TEST_THREADS);
+    Stream stream;
+#if defined(__HIPCC__)
+    attnWaveTestKernel<<<1, WAVE_TEST_THREADS, 0, stream.getHIPStream()>>>(sums.ptr(), exchanged.ptr());
+    gpuErrchk(hipGetLastError());
+#else
+    attnWaveTestKernel<<<1, WAVE_TEST_THREADS, 0, stream.getCUDAStream()>>>(sums.ptr(), exchanged.ptr());
+    gpuErrchk(cudaGetLastError());
+#endif
+    sums.download(stream);
+    exchanged.download(stream);
+    stream.sync();
+
+    int failures = 0;
+    for (int tid = 0; tid < WAVE_TEST_THREADS; ++tid) {
+        const int first = tid - tid % ATTN_WAVE_SIZE;
+        const int expectedIntSum = ATTN_WAVE_SIZE * first + ATTN_WAVE_SIZE * (ATTN_WAVE_SIZE - 1) / 2;
+        const float expectedSum = static_cast<float>(expectedIntSum);
+        const int expectedExchange = tid ^ WAVE_TEST_MASK;
+        const float gotSum = sums.ptrPinned().data[tid];
+        const int gotExchange = exchanged.ptrPinned().data[tid];
+        if (gotSum != expectedSum || gotExchange != expectedExchange) {
+            std::cout << "FAIL attnShflXor lane " << tid << ": sum " << gotSum << " (expected " << expectedSum
+                      << "), exchange " << gotExchange << " (expected " << expectedExchange << ")" << std::endl;
+            ++failures;
+        }
+    }
+    if (failures == 0) {
+        std::cout << "Running test attnShflXor butterfly sum and exchange, 2 waves: Success!!" << std::endl;
+        return 0;
+    }
+    return -1;
+}
+
+} // namespace fk
+
+int launch() {
+    return fk::launch_impl();
+}
+
+#else
+
+int launch() {
+    std::cout << "attention_wave.h has no CPU code to test" << std::endl;
+    return 0;
+}
+
+#endif // defined(__NVCC__) || defined(__HIPCC__)