vulkan: support GPUs without hardware DP4A in mul_mat_vec_ptq1_0 and eliminate dequant divergence

ggml/src/ggml-vulkan/vulkan-shaders/dequant_ptq1_0.comp

@@ -7,48 +7,57 @@ layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
 layout (binding = 0) readonly buffer A {block_ptq1_0 data_a[];};
 layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
 
-// One invocation per 8 elements: 16 invocations cover a 128-weight block. Element
-// order follows the CPU codec (16-byte qs chunk, 8-byte qs chunk, then qh), so the
-// index maths lives in one place rather than being reproduced per shader.
-float ptq1_0_trit_l(uint ib, uint e) {
-    uint b;
-    uint n;
-    if (e < 80u) {
-        b = uint(data_a[ib].qs[e & 15u]);
-        n = e >> 4u;
-    } else if (e < 120u) {
-        const uint t = e - 80u;
-        b = uint(data_a[ib].qs[16u + (t & 7u)]);
-        n = t >> 3u;
-    } else {
-        const uint t = e - 120u;
-        b = uint(data_a[ib].qh[t & 1u]);
-        n = t >> 1u;
-    }
-
-    uint v = b;
-    for (uint i = 0u; i < n; ++i) {
-        v = (v * 3u) & 0xFFu;
-    }
-    return float(int((v * 3u) >> 8u) - 1);
-}
+// Powers of 3: 3^0, 3^1, 3^2, 3^3, 3^4 (mod 256)
+const uint pow3[5] = uint[5](1u, 3u, 9u, 27u, 81u);
 
 void main() {
     const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64;
 
     const uint tid = gl_LocalInvocationID.x % 64;
-    const uint il  = tid / 4;
-    const uint ir  = tid % 4;
+    const uint il = tid / 4;
+    const uint ir = tid % 4;
     const uint ib = 4*i + ir;
     if (ib >= p.nel / 128) {
         return;
     }
 
     const uint b_idx = 512*i + 128*ir + 8*il;
-
     const float d = float(data_a[ib].d);
 
-    [[unroll]] for (uint l = 0; l < 8; ++l) {
-        data_b[b_idx + l] = D_TYPE(ptq1_0_trit_l(ib, 8*il + l) * d);
+    if (il < 10u) {
+        // e in [0..79]: qs[0..15], n is uniform for all 8 elements
+        const uint n = il >> 1u;
+        const uint p_val = pow3[n];
+        const uint qs_offset = (il & 1u) != 0u ? 8u : 0u;
+
+        [[unroll]] for (uint l = 0u; l < 8u; ++l) {
+            const uint b = uint(data_a[ib].qs[qs_offset + l]);
+            const uint v = (b * p_val) & 0xFFu;
+            const float trit = float(int((v * 3u) >> 8u) - 1);
+            data_b[b_idx + l] = D_TYPE(trit * d);
+        }
+    } else if (il < 15u) {
+        // e in [80..119]: qs[16..23], n is uniform for all 8 
+        const uint n = il - 10u;
+        const uint p_val = pow3[n];
+
+        [[unroll]] for (uint l = 0u; l < 8u; ++l) {
+            const uint b = uint(data_a[ib].qs[16u + l]);
+            const uint v = (b * p_val) & 0xFFu;
+            const float trit = float(int((v * 3u) >> 8u) - 1);
+            data_b[b_idx + l] = D_TYPE(trit * d);
+        }
+    } else {
+        // il == 15 (e in [120..127]): read both qh bytes once into registers
+        const uint qh0 = uint(data_a[ib].qh[0]);
+        const uint qh1 = uint(data_a[ib].qh[1]);
+
+        [[unroll]] for (uint l = 0u; l < 8u; ++l) {
+            const uint b = (l & 1u) != 0u ? qh1 : qh0;
+            const uint n = l >> 1u;
+            const uint v = (b * pow3[n]) & 0xFFu;
+            const float trit = float(int((v * 3u) >> 8u) - 1);
+            data_b[b_idx + l] = D_TYPE(trit * d);
+        }
     }
 }

ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq_ptq1_0.comp

@@ -12,9 +12,11 @@
 // Several columns (batched decode, ALU bound): passes of CG columns; per block decode all rows once per level, pack the trits and dot them with every column of the pass.
 
 #extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
-#extension GL_EXT_integer_dot_product : require
 #extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
 
+// Enable integer dot product optionally: allows compilation on GPUs without hardware DP4A (e.g. AMD Polaris / Vega)
+#extension GL_EXT_integer_dot_product : enable
+
 #define B_TYPE block_q8_1_x4
 
 #include "mul_mat_vec_base.glsl"
@@ -38,18 +40,27 @@ layout (constant_id = 3) const uint CG = 3;
 
 uint a_offset, b_offset, d_offset;
 
-#if defined(PTQ1_0_MUL3_SHIFT)
-#define MUL3(v) (((v) << 1u) + (v))
-#elif defined(PTQ1_0_MUL3_PK16)
+#if defined(PTQ1_0_MUL3_PK16)
 #define MUL3(v) pack32(unpack16(v) * u16vec2(3us, 3us))
 #else
-#define MUL3(v) ((v) * 3u)
+// (v << 1u) + (v) maps to a single-cycle shift-add (v_lshl_add_u32) on AMD GCN/RDNA
+#define MUL3(v) (((v) << 1u) + (v))
 #endif
 
-#ifdef PTQ1_0_DOT_NOSAT
+#if defined(GL_EXT_integer_dot_product)
+// Hardware DP4A: saturation is mathematically impossible (max sum in 32 trits is 8128 << 2^31-1),
+// so non-saturating dotPacked4x8EXT maps directly to single-cycle native instructions (v_dot4_i32_i8).
 #define DOT4(a, b, acc) ((acc) + dotPacked4x8EXT((a), (b)))
 #else
-#define DOT4(a, b, acc) dotPacked4x8AccSatEXT((a), (b), (acc))
+// Universal fallback for GPUs without hardware DP4A (AMD Polaris RX 400/500, Vega, older Intel/Nvidia):
+// bitfieldExtract maps directly to AMD GCN's v_bfe_i32 instruction.
+int32_t dot4_emulated(const int32_t a, const int32_t b, const int32_t acc) {
+    return acc + bitfieldExtract(a, 0, 8)  * bitfieldExtract(b, 0, 8)
+               + bitfieldExtract(a, 8, 8)  * bitfieldExtract(b, 8, 8)
+               + bitfieldExtract(a, 16, 8) * bitfieldExtract(b, 16, 8)
+               + bitfieldExtract(a, 24, 8) * bitfieldExtract(b, 24, 8);
+}
+#define DOT4(a, b, acc) dot4_emulated((a), (b), (acc))
 #endif
 
 // One trit level from the 16-bit lanes vl (bytes 0,2) and vh (bytes 1,3): after w = v*3 the trit is in byte 1/3 and the residue in byte 0/2.