Validate gfx1201 WMMA layout and compile first K kernel

kernels/common/hip_device_minimal.h

@@ -0,0 +1,32 @@
+#pragma once
+
+// Device-only Windows HIP build shim. The diagnostic/kernel code objects are
+// compiled with -nogpuinc/-nogpulib; no host HIP or C++ standard library is
+// needed. Definitions below mirror public HIP qualifiers and LLVM AMDGPU
+// work-item/grid builtins. The normal Linux HIP include path is untouched.
+#define __device__ __attribute__((device))
+#define __global__ __attribute__((global))
+#define __shared__ __attribute__((shared))
+#define __forceinline__ __attribute__((always_inline)) inline
+#define __launch_bounds__(n) __attribute__((amdgpu_flat_work_group_size(1, n)))
+
+struct d4r_device_dim3 { unsigned x, y, z; };
+__device__ __forceinline__ d4r_device_dim3 d4r_thread_idx() {
+    return {__builtin_amdgcn_workitem_id_x(), __builtin_amdgcn_workitem_id_y(),
+        __builtin_amdgcn_workitem_id_z()};
+}
+__device__ __forceinline__ d4r_device_dim3 d4r_block_idx() {
+    return {__builtin_amdgcn_workgroup_id_x(), __builtin_amdgcn_workgroup_id_y(),
+        __builtin_amdgcn_workgroup_id_z()};
+}
+__device__ __forceinline__ d4r_device_dim3 d4r_grid_dim() {
+    return {__builtin_amdgcn_grid_size_x() / __builtin_amdgcn_workgroup_size_x(),
+        __builtin_amdgcn_grid_size_y() / __builtin_amdgcn_workgroup_size_y(),
+        __builtin_amdgcn_grid_size_z() / __builtin_amdgcn_workgroup_size_z()};
+}
+#define threadIdx d4r_thread_idx()
+#define blockIdx d4r_block_idx()
+#define gridDim d4r_grid_dim()
+__device__ __attribute__((convergent, always_inline)) inline void __syncthreads() {
+    __builtin_amdgcn_s_barrier();
+}

kernels/common/wmma_backend.h

@@ -0,0 +1,57 @@
+#pragma once
+#include <stdint.h>
+
+// d4r's gfx11 interface: each lane supplies 16 FP16 K operands. The 8 FP32
+// outputs have logical row 2*i + lane/16 and logical column lane%16.
+// gfx12 consumes only 8 FP16 operands per lane and emits 8 consecutive rows
+// per lane group. Keep the old interface for K/M/texture call sites until their
+// surrounding fragments are ported and validated individually.
+using d4r_wmma_h16 = _Float16 __attribute__((ext_vector_type(16)));
+using d4r_wmma_h8 = _Float16 __attribute__((ext_vector_type(8)));
+using d4r_wmma_f8 = float __attribute__((ext_vector_type(8)));
+
+__attribute__((device, always_inline)) static inline float d4r_wmma_other_half(float value)
+{
+    const uint32_t bits = __builtin_bit_cast(uint32_t, value);
+    const uint32_t swapped = __builtin_amdgcn_permlanex16(
+        bits, bits, 0x76543210u, 0xfedcba98u, false, false);
+    return __builtin_bit_cast(float, swapped);
+}
+
+__attribute__((device, always_inline)) static inline d4r_wmma_f8
+d4r_wmma_legacy_layout(d4r_wmma_h16 a, d4r_wmma_h16 b, d4r_wmma_f8 c)
+{
+#if defined(__gfx1200__) || defined(__gfx1201__)
+    const uint32_t group = __builtin_amdgcn_mbcnt_lo(~0u, 0u) >> 4;
+    d4r_wmma_h8 a12, b12;
+    d4r_wmma_f8 c12;
+    d4r_wmma_f8 remote_c;
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        remote_c[i] = d4r_wmma_other_half(c[i]);
+#pragma unroll
+    for (int i = 0; i < 8; ++i) {
+        a12[i] = a[8u * group + i];
+        b12[i] = b[8u * group + i];
+        const uint32_t row = 8u * group + i;
+        const uint32_t old_index = row >> 1;
+        c12[i] = (row & 1u) == group ? c[old_index] : remote_c[old_index];
+    }
+    const d4r_wmma_f8 d12 = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(a12, b12, c12);
+    d4r_wmma_f8 result, remote_d;
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        remote_d[i] = d4r_wmma_other_half(d12[i]);
+#pragma unroll
+    for (int i = 0; i < 8; ++i) {
+        const uint32_t row = 2u * i + group;
+        const uint32_t new_index = row & 7u;
+        result[i] = (row >> 3) == group ? d12[new_index] : remote_d[new_index];
+    }
+    return result;
+#elif defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__)
+    return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
+#else
+#error "d4r WMMA backend requires a validated gfx11 or gfx12 target"
+#endif
+}

kernels/k/pwin_common.h

@@ -6,7 +6,11 @@
 // and the activations as the B operand (lane l: token l & 15, 16 K values). The f32 result then has
 // lane l = token l & 15 and VGPR i = output channel 2i + (l >> 4) of the 16-channel tile ("D^T").
 #pragma once
+#ifdef D4R_DEVICE_ONLY_MINIMAL
+#include "../common/hip_device_minimal.h"
+#else
 #include <hip/hip_runtime.h>
+#endif
 #include <stdint.h>
 
 #pragma clang fp contract(off)
@@ -17,6 +21,7 @@ typedef _Float16 h16 __attribute__((ext_vector_type(16)));
 typedef float f8v __attribute__((ext_vector_type(8)));
 typedef uint32_t u8v __attribute__((ext_vector_type(8)));
 typedef uint32_t u4v __attribute__((ext_vector_type(4)));
+#include "../common/wmma_backend.h"
 
 __device__ __forceinline__ uint32_t lane_id()
 {
@@ -25,7 +30,7 @@ __device__ __forceinline__ uint32_t lane_id()
 
 __device__ __forceinline__ f8v wmma(u8v a, u8v b, f8v c)
 {
-    return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
+    return d4r_wmma_legacy_layout(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
 }
 
 // one k16 step with the f16 accumulator of NVIDIA's f16 wmma (rounded after the step)