hip: fast packed-byte intrinsics for RDNA3/RDNA4 (v_perm_b32 + SWAR)

CMakeLists.txt

@@ -373,6 +373,12 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
       add_test(NAME hip_intrinsics COMMAND hip_intrinsics)
       list(APPEND _strata_hip_tests hip_intrinsics)
     endif()
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/hip/fast_intrinsics_parity.cpp")
+      add_executable(hip_fast_intrinsics_parity tests/hip/fast_intrinsics_parity.cpp)
+      target_link_libraries(hip_fast_intrinsics_parity PRIVATE strata_hip_runtime)
+      add_test(NAME hip_fast_intrinsics_parity COMMAND hip_fast_intrinsics_parity)
+      list(APPEND _strata_hip_tests hip_fast_intrinsics_parity)
+    endif()
     if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/hip/handoff.cpp")
       add_executable(hip_handoff tests/hip/handoff.cpp)
       target_link_libraries(hip_handoff PRIVATE strata_kernels strata_hip_runtime)

cmake/hip_backend.cmake

@@ -71,7 +71,7 @@ file(GLOB_RECURSE _strata_hip_sources CONFIGURE_DEPENDS
 if(_strata_hip_sources)
   set_source_files_properties(${_strata_hip_sources} PROPERTIES LANGUAGE HIP)
 endif()
-foreach(_source IN ITEMS tests/hip/intrinsics.cpp tests/hip/native_qsa_score.cpp)
+foreach(_source IN ITEMS tests/hip/intrinsics.cpp tests/hip/native_qsa_score.cpp tests/hip/fast_intrinsics_parity.cpp)
   if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${_source}")
     set_source_files_properties("${_source}" PROPERTIES LANGUAGE HIP)
   endif()

include/strata/hip_compat/intrinsics.hpp

@@ -32,45 +32,46 @@ __device__ __forceinline__ int dp4a(int a, int b, int c) {
 #endif
 }
 
-// CUDA's packed byte subtract wraps independently in each unsigned byte lane.
+// CUDA's __byte_perm (default mode): result byte i is byte s.nibble[i] & 7 of the pair {y:x}, x the low word.
+// HIP's own version indexes a byte array in private memory, which becomes scratch traffic in the i-quant table
+// lookups; v_perm_b32 is the same selection in one instruction (selector bytes 0-3 pick from its second
+// operand, 4-7 from its first), so the nibble selector is spread to bytes and masked to 0-7.
+__device__ __forceinline__ uint32_t byte_perm(uint32_t x, uint32_t y, uint32_t s) {
+    const uint32_t sel = (s & 0x7u) | ((s & 0x70u) << 4) | ((s & 0x700u) << 8) | ((s & 0x7000u) << 12);
+    return __builtin_amdgcn_perm(y, x, sel);
+}
+
+// The packed byte operations below work on all four lanes at once (SWAR, Hacker's Delight 2-18): gfx11/gfx12 have
+// no packed 8-bit subtract, and a per-lane loop costs the Q3_K/Q6_K dots and the i-quant sign expansion several
+// times the instructions.  kHigh is each lane's top bit.
+constexpr uint32_t kHigh = 0x80808080u;
+
+// CUDA's packed byte subtract wraps independently in each unsigned byte lane: subtract the low seven bits with the
+// minuend's top bit forced on (so no borrow crosses a lane), then fix each top bit by XOR.
 __device__ __forceinline__ int vsub4(int a, int b) {
     const uint32_t ua = static_cast<uint32_t>(a);
     const uint32_t ub = static_cast<uint32_t>(b);
-    uint32_t out = 0;
-#pragma unroll
-    for (int lane = 0; lane < 4; ++lane) {
-        const uint32_t x = (ua >> (lane * 8)) & 0xffu;
-        const uint32_t y = (ub >> (lane * 8)) & 0xffu;
-        out |= ((x - y) & 0xffu) << (lane * 8);
-    }
-    return static_cast<int>(out);
+    return static_cast<int>(((ua | kHigh) - (ub & ~kHigh)) ^ ((ua ^ ~ub) & kHigh));
 }
 
-// CUDA's packed signed-byte saturating subtract.
+// CUDA's packed signed-byte saturating subtract: the wrapping difference, and in each lane that overflowed (operand
+// signs differ and the result's sign differs from the minuend's) the bound on the minuend's side, 0x7f or 0x80.
 __device__ __forceinline__ int vsubss4(int a, int b) {
     const uint32_t ua = static_cast<uint32_t>(a);
     const uint32_t ub = static_cast<uint32_t>(b);
-    uint32_t out = 0;
-#pragma unroll
-    for (int lane = 0; lane < 4; ++lane) {
-        int value = signed_byte(ua, lane) - signed_byte(ub, lane);
-        value = value < -128 ? -128 : (value > 127 ? 127 : value);
-        out |= (static_cast<uint32_t>(value) & 0xffu) << (lane * 8);
-    }
-    return static_cast<int>(out);
+    const uint32_t d = static_cast<uint32_t>(vsub4(a, b));
+    const uint32_t overflow = (ua ^ ub) & (ua ^ d) & kHigh;
+    const uint32_t mask = (overflow >> 7) * 0xffu;
+    const uint32_t bound = 0x7f7f7f7fu + ((ua & kHigh) >> 7);
+    return static_cast<int>((d & ~mask) | (bound & mask));
 }
 
-// CUDA's four-lane byte compare, returning 0xff for each unequal lane and 0 otherwise.
+// CUDA's four-lane byte compare, returning 0xff for each unequal lane and 0 otherwise: a lane of a ^ b is nonzero
+// when its low seven bits carry into the top bit on adding 0x7f, or its top bit is already set.
 __device__ __forceinline__ int vcmpne4(int a, int b) {
-    const uint32_t ua = static_cast<uint32_t>(a);
-    const uint32_t ub = static_cast<uint32_t>(b);
-    uint32_t out = 0;
-#pragma unroll
-    for (int lane = 0; lane < 4; ++lane) {
-        if (((ua >> (lane * 8)) & 0xffu) != ((ub >> (lane * 8)) & 0xffu))
-            out |= 0xffu << (lane * 8);
-    }
-    return static_cast<int>(out);
+    const uint32_t t = static_cast<uint32_t>(a) ^ static_cast<uint32_t>(b);
+    const uint32_t nonzero = (((t & ~kHigh) + ~kHigh) | t) & kHigh;
+    return static_cast<int>((nonzero >> 7) * 0xffu);
 }
 
 // CUDA's mask argument describes participating lanes. The current kernel set uses full
@@ -111,6 +112,7 @@ __device__ __forceinline__ unsigned ballot_sync(uint32_t mask, int predicate) {
 }  // namespace strata::hip_compat
 
 #define __dp4a(a, b, c) (::strata::hip_compat::dp4a((a), (b), (c)))
+#define __byte_perm(x, y, s) (::strata::hip_compat::byte_perm((x), (y), (s)))
 #define __vsub4(a, b) (::strata::hip_compat::vsub4((a), (b)))
 #define __vsubss4(a, b) (::strata::hip_compat::vsubss4((a), (b)))
 #define __vcmpne4(a, b) (::strata::hip_compat::vcmpne4((a), (b)))

tests/hip/fast_intrinsics_parity.cpp

@@ -0,0 +1,75 @@
+// Independent exhaustive-byte-pair and randomized mixed-lane parity.
+#include <hip/hip_runtime.h>
+#include "strata/hip_compat/intrinsics.hpp"
+#include <algorithm>
+#include <cstdint>
+#include <cstdio>
+#include <random>
+#include <vector>
+
+#define CHECK(call) do { auto error = (call); if (error != hipSuccess) { \
+    std::fprintf(stderr, "%s: %s\n", #call, hipGetErrorString(error)); return 2; } } while (0)
+
+struct Input { uint32_t a, b, selector; };
+__global__ void probe(const Input* input, uint32_t* output, size_t count) {
+    size_t i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= count) return;
+    auto x = input[i];
+    output[i * 4] = __byte_perm(x.a, x.b, x.selector);
+    output[i * 4 + 1] = static_cast<uint32_t>(__vsub4(static_cast<int>(x.a), static_cast<int>(x.b)));
+    output[i * 4 + 2] = static_cast<uint32_t>(__vsubss4(static_cast<int>(x.a), static_cast<int>(x.b)));
+    output[i * 4 + 3] = static_cast<uint32_t>(__vcmpne4(static_cast<int>(x.a), static_cast<int>(x.b)));
+}
+
+int main() {
+    constexpr size_t count = 262144, exhaustive = 65536, guard = 16;
+    constexpr uint32_t sentinel = 0xdeadbeefu;
+    std::vector<Input> input(count);
+    std::mt19937 random(20260930);
+    for (size_t i = 0; i < count; ++i) {
+        input[i].a = i < exhaustive ? static_cast<uint32_t>(i >> 8) * 0x01010101u : random();
+        input[i].b = i < exhaustive ? static_cast<uint32_t>(i & 255u) * 0x01010101u : random();
+        input[i].selector = static_cast<uint32_t>(i);
+    }
+    std::vector<uint32_t> output(count * 4 + guard, sentinel);
+    Input* device_input = nullptr;
+    uint32_t* device_output = nullptr;
+    CHECK(hipMalloc(reinterpret_cast<void**>(&device_input), input.size() * sizeof(Input)));
+    CHECK(hipMalloc(reinterpret_cast<void**>(&device_output), output.size() * sizeof(uint32_t)));
+    CHECK(hipMemcpy(device_input, input.data(), input.size() * sizeof(Input), hipMemcpyHostToDevice));
+    CHECK(hipMemcpy(device_output, output.data(), output.size() * sizeof(uint32_t), hipMemcpyHostToDevice));
+    hipLaunchKernelGGL(probe, dim3((count + 255) / 256), dim3(256), 0, 0, device_input, device_output, count);
+    CHECK(hipGetLastError());
+    CHECK(hipDeviceSynchronize());
+    CHECK(hipMemcpy(output.data(), device_output, output.size() * sizeof(uint32_t), hipMemcpyDeviceToHost));
+    CHECK(hipFree(device_output));
+    CHECK(hipFree(device_input));
+    for (size_t i = 0; i < count; ++i) {
+        uint32_t expected[4] = {};
+        uint64_t pair = (static_cast<uint64_t>(input[i].b) << 32) | input[i].a;
+        for (unsigned lane = 0; lane < 4; ++lane) {
+            unsigned shift = lane * 8;
+            unsigned a = (input[i].a >> shift) & 255u, b = (input[i].b >> shift) & 255u;
+            unsigned selected = (input[i].selector >> (lane * 4)) & 7u;
+            expected[0] |= static_cast<uint32_t>((pair >> (selected * 8)) & 255u) << shift;
+            expected[1] |= ((a - b) & 255u) << shift;
+            int signed_a = a < 128 ? static_cast<int>(a) : static_cast<int>(a) - 256;
+            int signed_b = b < 128 ? static_cast<int>(b) : static_cast<int>(b) - 256;
+            expected[2] |= (static_cast<uint32_t>(std::clamp(signed_a - signed_b, -128, 127)) & 255u) << shift;
+            expected[3] |= (a != b ? 255u : 0u) << shift;
+        }
+        for (unsigned operation = 0; operation < 4; ++operation) {
+            if (output[i * 4 + operation] != expected[operation]) {
+                std::fprintf(stderr, "Mismatch case=%zu operation=%u a=%08x b=%08x selector=%04x actual=%08x expected=%08x\n",
+                             i, operation, input[i].a, input[i].b, input[i].selector,
+                             output[i * 4 + operation], expected[operation]);
+                return 1;
+            }
+        }
+    }
+    for (size_t i = count * 4; i < output.size(); ++i) {
+        if (output[i] != sentinel) { std::fprintf(stderr, "Output guard changed\n"); return 1; }
+    }
+    std::printf("Fast intrinsics parity PASS: %zu cases, %zu operation results, all 65536 byte pairs, 4096 selectors, 16 guards\n", count, count * 4);
+    return 0;
+}