From 17c1828f97277d207d83c62765cac13dc9eda603 Mon Sep 17 00:00:00 2001
From: woutermans <woutertje554@gmail.com>
Date: Thu, 15 Oct 2026 02:49:16 +0200
Subject: [PATCH] engine/rr, vulkan: honour PTX alignment for 32-bit and shared
 vector accesses

---
 engine/rr/ptx2rr.py | 21 +++++++++++++--------
 vulkan/ptx2glsl.py  |  3 +++
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/engine/rr/ptx2rr.py b/engine/rr/ptx2rr.py
index dde5d11..d454057 100755
--- a/engine/rr/ptx2rr.py
+++ b/engine/rr/ptx2rr.py
@@ -98,7 +98,8 @@ void OOB(uint k, uint a, uint size) { if (a + 2u > size) { atomicAdd(dbg.d[3u *
 #define CHECKS(a)
 #endif
 uint GW16(uint a) { CHECKW(a); return uint(wbuf.wh[a >> 1]); }
-uint GW32(uint a) { CHECKW(a + 2u); return uint(wbuf.wh[a >> 1]) | (uint(wbuf.wh[(a >> 1) + 1u]) << 16); }
+// a 32-bit access is 4-byte aligned (PTX ISA): one word of the same buffer's word view
+uint GW32(uint a) { CHECKW(a + 2u); return wbuf32.ww[a >> 2]; }
 uint GA16(uint a) { CHECKR(a); return uint(abuf.ah[a >> 1]); }
 uint GA32(uint a) { CHECKR(a + 2u); return uint(abuf.ah[a >> 1]) | (uint(abuf.ah[(a >> 1) + 1u]) << 16); }
 void SA16(uint a, uint v) { CHECKS(a); abuf.ah[a >> 1] = uint16_t(v & 0xffffu); }
@@ -238,13 +239,15 @@ uint movm_trans(uint src)
     return ((w1 >> s) & 0xffffu) | (((w2 >> s) & 0xffffu) << 16);
 }
 
-// PTX min/max: a NaN operand yields the other one (GLSL leaves NaN undefined, so the result depended on codegen)
-float16_t pmin(float16_t a, float16_t b) { return isnan(a) ? b : isnan(b) ? a : min(a, b); }
-float16_t pmax(float16_t a, float16_t b) { return isnan(a) ? b : isnan(b) ? a : max(a, b); }
-f16vec2 pmin2(f16vec2 a, f16vec2 b) { return f16vec2(pmin(a.x, b.x), pmin(a.y, b.y)); }
-f16vec2 pmax2(f16vec2 a, f16vec2 b) { return f16vec2(pmax(a.x, b.x), pmax(a.y, b.y)); }
-float pminf(float a, float b) { return isnan(a) ? b : isnan(b) ? a : min(a, b); }
-float pmaxf(float a, float b) { return isnan(a) ? b : isnan(b) ? a : max(a, b); }
+// PTX min/max: a NaN operand yields the other one. GLSL's min/max leave NaN undefined (the result depended on
+// codegen); SPIR-V's NMin/NMax (GLSL.std.450 79/80) are defined that way, and compile to the bare v_min/v_max
+// (checked on RADV / RDNA4 for f16, f16x2 and f32 with NaN, infinity and signed zero operands, 2026-10-15)
+spirv_instruction(set = "GLSL.std.450", id = 79) float16_t pmin(float16_t a, float16_t b);
+spirv_instruction(set = "GLSL.std.450", id = 80) float16_t pmax(float16_t a, float16_t b);
+spirv_instruction(set = "GLSL.std.450", id = 79) f16vec2 pmin2(f16vec2 a, f16vec2 b);
+spirv_instruction(set = "GLSL.std.450", id = 80) f16vec2 pmax2(f16vec2 a, f16vec2 b);
+spirv_instruction(set = "GLSL.std.450", id = 79) float pminf(float a, float b);
+spirv_instruction(set = "GLSL.std.450", id = 80) float pmaxf(float a, float b);
 // e4m3 <-> f16 (PTX cvt semantics; RNE, satfinite; NaN -> 0x7f e4m3 / 0x7e00 f16).
 uint e4m3_to_h(uint b)
 {
@@ -801,10 +804,12 @@ class RRT(base.T):
         w('#extension GL_KHR_cooperative_matrix : require')
         w('#extension GL_KHR_memory_scope_semantics : require')
         w('#extension GL_EXT_float_e4m3 : require')
+        w('#extension GL_EXT_spirv_intrinsics : require')
         w(f'// Generated by engine/rr/ptx2rr.py from {self.kernel}. Model data derived from the user\'s DLL; do not commit.')
         w(f'layout(local_size_x = {self.bx * self.by * self.bz}, local_size_y = 1, local_size_z = 1) in;')
         w('layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint w[]; } param;')
         w('layout(set = 0, binding = 1, std430) readonly buffer WeightBuf { uint16_t wh[]; } wbuf;')
+        w('layout(set = 0, binding = 1, std430) readonly buffer WeightBuf32 { uint ww[]; } wbuf32;')
         w('layout(set = 0, binding = 2, std430) buffer ArenaBuf { uint16_t ah[]; } abuf;')
         n = 3
         for off in sorted(self.textures):
diff --git a/vulkan/ptx2glsl.py b/vulkan/ptx2glsl.py
index 1df4404..3ba673e 100644
--- a/vulkan/ptx2glsl.py
+++ b/vulkan/ptx2glsl.py
@@ -667,6 +667,9 @@ class T:
                 if po is None:
                     raise Err('param address ' + a[1])
                 e = f'PARAM{size * 8}({po + i * size}u)'
+            elif 'shared' in mods and size == 2 and vec > 1:
+                # a vector access is aligned to its whole size (PTX ISA): its 16-bit halves have fixed places in words
+                e = f'(sh[({V(base, "u32")} + {off + i * size}u) >> 2] ' + ('& 0xffffu)' if i % 2 == 0 else '>> 16u)')
             elif 'shared' in mods:
                 e = f'SH{size * 8}({V(base, "u32")} + {off + i * size}u)'
             elif 'local' in mods: