d4r/patches/zluda/0013-gfx12-f16-reference-accumulation.patch

diff --git a/llvm_zluda/src/compile.rs b/llvm_zluda/src/compile.rs
index 5a54417..90c0034 100644
--- a/llvm_zluda/src/compile.rs
+++ b/llvm_zluda/src/compile.rs
@@ -158,6 +158,10 @@ fn wmma_rewrite_flags(gcn_arch: &str) -> Option<u32> {
     let mut flags = if wmma_fp8_rewrite_requested() { 1 } else { 0 };
     if family == 12 {
         flags |= 2;
+        // Explicit diagnostic/correctness mode; retain upstream's fast default.
+        if std::env::var("D4R_ZLUDA_WMMA_F16_REFERENCE").is_ok_and(|value| value.trim() == "1") {
+            flags |= 8;
+        }
         // native FP8 needs gfx12 hardware; the layout shim on gfx11 takes it only for tests (D4R_WMMA12_SHIM)
         if flags & 1 != 0 && wmma_fp8_native_requested() {
             flags |= 4;
diff --git a/llvm_zluda/src/lib.cpp b/llvm_zluda/src/lib.cpp
index 07238db..d262fc6 100644
--- a/llvm_zluda/src/lib.cpp
+++ b/llvm_zluda/src/lib.cpp
@@ -970,6 +970,7 @@ static unsigned widenMmaK8(Module *M)
 // Flags bit 0: also lower m16n8k32 e4m3 MMAs (D4R_ZLUDA_WMMA_FP8=1).
 // Flags bit 1: gfx12 (RDNA4) WMMA layout: the d4r_wmma12_* helpers.
 // Flags bit 2 (with bit 1): e4m3 MMAs on native FP8 WMMA: the d4r_wmma12f8_* e4m3 helpers.
+// Flags bit 3 (with bit 1): f64 sum / one f16 round for f16 accumulator MMAs.
 // The helper sets share one contract (see zluda_ptx_impl.cpp), so only their names change here.
 unsigned LLVMZludaRewriteMmaF16Wmma(LLVMModuleRef ModuleRef, unsigned Flags)
 {
@@ -999,14 +1000,15 @@ unsigned LLVMZludaRewriteMmaF16Wmma(LLVMModuleRef ModuleRef, unsigned Flags)
     const std::string Arch = Gfx12 ? "__zluda_ptx_impl_d4r_wmma12_" : "__zluda_ptx_impl_d4r_wmma_";
     const std::string E4m3 = NativeFp8 ? "__zluda_ptx_impl_d4r_wmma12f8_" : Arch;
     const std::string Widen = NativeFp8 ? E4m3 : "__zluda_ptx_impl_d4r_wmma_";
+    const std::string F16 = Arch + (Gfx12 && (Flags & 8) != 0 ? "t_mma_f16_reference" : "t_mma_f16");
     Helpers H;
     H.AFromB = M->getOrInsertFunction((Arch + "t_a_from_b").c_str(), Attributes, V8, V2);
     H.BFromA = M->getOrInsertFunction((Arch + "t_b_from_a").c_str(), Attributes, V8, V4);
-    H.Mma = M->getOrInsertFunction((Arch + "t_mma_f16").c_str(), Attributes, V2, V8, V8, V2);
+    H.Mma = M->getOrInsertFunction(F16.c_str(), Attributes, V2, V8, V8, V2);
     H.FromNv = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_from_nv", Attributes, V2, V2);
     H.ToNv = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_to_nv", Attributes, V2, V2);
     H.AFromBPair = M->getOrInsertFunction((Arch + "t_a_from_b_pair").c_str(), Attributes, V8, V2, V2);
-    H.MmaPair = M->getOrInsertFunction((Arch + "t_mma_f16_pair").c_str(), Attributes, V4, V8, V8, V2, V2);
+    H.MmaPair = M->getOrInsertFunction((F16 + "_pair").c_str(), Attributes, V4, V8, V8, V2, V2);
     {
         Type *F = Type::getFloatTy(Context);
         Type *V4F = FixedVectorType::get(F, 4);
diff --git a/ptx/lib/zluda_ptx_impl.cpp b/ptx/lib/zluda_ptx_impl.cpp
index 4ff25dd..36bd6b1 100644
--- a/ptx/lib/zluda_ptx_impl.cpp
+++ b/ptx/lib/zluda_ptx_impl.cpp
@@ -1801,6 +1801,78 @@ __device__ static wmma_float8 wmma12_f16(wmma_uint8 a, wmma_uint8 b, wmma_float8
 #endif
 }
 
+// Correctness baseline for f16-accumulator MMAs. f32 WMMA accumulation may
+// reach the other f16 neighbour around cancellation-heavy rounding midpoints.
+// Keep the existing gfx12 gathers/layout; sum their f16 operands in f64 and
+// round directly to f16 once, matching d4r's independent NumPy/PTX reference.
+// The e4m3 and f32-shadow paths continue using native WMMA.
+__device__ static double wmma12_half_value(uint32_t bits)
+{
+    const uint32_t magnitude = bits & 0x7fffu, exponent = magnitude >> 10;
+    const double value = exponent == 0 ? double(magnitude) * 0x1p-24 :
+        double(std::bit_cast<float>(exponent == 31 ? (0x7f800000u | ((magnitude & 1023u) << 13)) :
+                                     ((magnitude << 13) + 0x38000000u)));
+    return bits & 0x8000u ? -value : value;
+}
+
+// Preserve f16 subnormal result bits even when the target's default f16
+// arithmetic mode flushes them. No intermediate f32 rounding is involved.
+__device__ static uint32_t wmma12_half_round(double value)
+{
+    const uint64_t bits = std::bit_cast<uint64_t>(value);
+    const uint32_t sign = uint32_t(bits >> 48) & 0x8000u;
+    const uint32_t exponent = uint32_t((bits >> 52) & 2047u);
+    const uint64_t fraction = bits & 0xfffffffffffffu;
+    if (exponent == 2047u) return sign | (fraction ? 0x7e00u : 0x7c00u);
+    const int power = int(exponent) - 1023;
+    if (power < -25) return sign;
+    if (power > 15) return sign | 0x7c00u;
+    const unsigned shift = power < -14 ? unsigned(28 - power) : 42u;
+    const uint64_t significand = fraction | 0x10000000000000ull;
+    uint32_t code = uint32_t(significand >> shift);
+    const uint64_t tail = significand & ((uint64_t(1) << shift) - 1);
+    const uint64_t midpoint = uint64_t(1) << (shift - 1);
+    code += tail > midpoint || (tail == midpoint && (code & 1u));
+    if (power >= -14) code += uint32_t(power + 14) << 10;
+    return sign | code;
+}
+
+__device__ static wmma_uint4 wmma12_reference_c(uint2::Native_vec_ c0, uint2::Native_vec_ c1, bool paired)
+{
+    const bool high = wmma12_half() != 0;
+    const uint2::Native_vec_ own = paired && high ? c1 : c0;
+    const uint2::Native_vec_ send = paired && !high ? c1 : c0;
+    const uint2::Native_vec_ other = {wmma12_other_half(send.x), wmma12_other_half(send.y)};
+    const uint2::Native_vec_ even = paired && high ? other : own;
+    const uint2::Native_vec_ odd = paired && high ? own : other;
+    return {(even.x & 0xffffu) | (odd.x << 16), (even.x >> 16) | (odd.x & 0xffff0000u),
+            (even.y & 0xffffu) | (odd.y << 16), (even.y >> 16) | (odd.y & 0xffff0000u)};
+}
+
+__device__ static wmma_uint4 wmma12_f16_reference(wmma_uint8 a, wmma_uint8 b, wmma_uint4 c)
+{
+    // x/y: packed even rows, z/w: packed odd rows, in this half-wave.
+    wmma_uint4 result = {};
+    const uint32_t half = wmma12_half();
+    for (uint32_t i = 0; i < 8; ++i)
+    {
+        const uint32_t row = i + 8 * half;
+        double sum = wmma12_half_value((c[i / 2] >> (16 * (i & 1))) & 0xffffu);
+        for (uint32_t khalf = 0; khalf < 2; ++khalf)
+            for (uint32_t j = 0; j < 4; ++j)
+            {
+                const uint32_t av = wmma_from_lane(row + 16 * khalf, a[j]);
+                const uint32_t other = wmma12_other_half(b[j]);
+                const uint32_t bv = half == khalf ? b[j] : other;
+                sum += wmma12_half_value(av & 0xffffu) * wmma12_half_value(bv & 0xffffu);
+                sum += wmma12_half_value(av >> 16) * wmma12_half_value(bv >> 16);
+            }
+        // Do not round through f32: that can reintroduce the midpoint error.
+        result[(i & 1) * 2 + i / 4] |= wmma12_half_round(sum) << (16 * ((i / 2) & 1));
+    }
+    return result;
+}
+
 // f32 D' = A' B' + C' for operands whose dwords hold two e4m3 bytes in their low 16 bits (the d4r_wmma12f8
 // e4m3 "widening" below keeps the bytes): true e4m3 values, no scaling
 __device__ static wmma_float8 wmma12_e4m3(wmma_uint8 a, wmma_uint8 b, wmma_float8 c)
@@ -1959,6 +2031,22 @@ extern "C"
         return wmma12_t_from_d_pair(wmma12_f16(a_t, b_t, wmma12_c_from_t_pair(c0_t, c1_t)));
     }
 
+    uint2::Native_vec_ FUNC(d4r_wmma12_t_mma_f16_reference)(wmma_uint8 a_t, wmma_uint8 b_t, uint2::Native_vec_ c_t)
+    {
+        const wmma_uint4 rows = wmma12_f16_reference(a_t, b_t, wmma12_reference_c(c_t, c_t, false));
+        const uint2::Native_vec_ other = {wmma12_other_half(rows.z), wmma12_other_half(rows.w)};
+        return wmma12_half() ? other : uint2::Native_vec_{rows.x, rows.y};
+    }
+
+    uint4::Native_vec_ FUNC(d4r_wmma12_t_mma_f16_reference_pair)(wmma_uint8 a_t, wmma_uint8 b_t, uint2::Native_vec_ c0_t, uint2::Native_vec_ c1_t)
+    {
+        const wmma_uint4 rows = wmma12_f16_reference(a_t, b_t, wmma12_reference_c(c0_t, c1_t, true));
+        const bool high = wmma12_half() != 0;
+        const uint32_t r0 = wmma12_other_half(high ? rows.x : rows.z);
+        const uint32_t r1 = wmma12_other_half(high ? rows.y : rows.w);
+        return high ? uint4::Native_vec_{r0, r1, rows.z, rows.w} : uint4::Native_vec_{rows.x, rows.y, r0, r1};
+    }
+
     wmma_float4 FUNC(d4r_wmma12_s_mma_f16)(wmma_uint8 a_t, wmma_uint8 b_t, wmma_float4 c, float up, float down)
     {
         return wmma12_s_from_d(wmma12_f16(a_t, b_t, wmma12_c_from_s(c * up))) * down;