From c13b0027f3820d26076bf6e97d45ddb0a531f0c2 Mon Sep 17 00:00:00 2001
From: illuzen <illuzen@users.noreply.github.com>
Date: Thu, 11 Dec 2025 14:55:14 +0800
Subject: [PATCH] optimized mds

---
 crates/engine-gpu/src/mining.wgsl | 51 ++++++++++++-------------------
 1 file changed, 20 insertions(+), 31 deletions(-)

diff --git a/crates/engine-gpu/src/mining.wgsl b/crates/engine-gpu/src/mining.wgsl
index 3d99c4c..6a2819d 100644
--- a/crates/engine-gpu/src/mining.wgsl
+++ b/crates/engine-gpu/src/mining.wgsl
@@ -597,44 +597,33 @@ fn external_linear_layer(state: ptr<function, array<GoldilocksField, 12>>) {
     // First apply the 4x4 MDS matrix to each consecutive 4 elements
     for (var chunk = 0u; chunk < 3u; chunk++) {
         let offset = chunk * 4u;
-        var chunk_state: array<GoldilocksField, 4>;
+        var x: array<GoldilocksField, 4>;
 
         // Copy chunk
         for (var i = 0u; i < 4u; i++) {
-            chunk_state[i] = (*state)[offset + i];
+            x[i] = (*state)[offset + i];
         }
 
-        // Apply 4x4 MDS matrix transformation: [[2, 3, 1, 1], [1, 2, 3, 1], [1, 1, 2, 3], [3, 1, 1, 2]]
-        let two = gf_from_limbs(2u, 0u, 0u, 0u);
-        let three = gf_from_limbs(3u, 0u, 0u, 0u);
+        // Optimized 4x4 MDS application using only additions and doubles,
+        // matching apply_external_linear_layer_to_chunk in the CPU implementation.
+        let t01 = gf_add(x[0], x[1]);
+        let t23 = gf_add(x[2], x[3]);
+        let t0123 = gf_add(t01, t23);
+        let t01123 = gf_add(t0123, x[1]);
+        let t01233 = gf_add(t0123, x[3]);
 
-        // new[0] = 2*x[0] + 3*x[1] + 1*x[2] + 1*x[3]
-        let new_0 = gf_add(gf_add(gf_add(
-            gf_mul(chunk_state[0], two),
-            gf_mul(chunk_state[1], three)),
-            chunk_state[2]),
-            chunk_state[3]);
+        let two_x0 = gf_add(x[0], x[0]);
+        let two_x2 = gf_add(x[2], x[2]);
 
-        // new[1] = 1*x[0] + 2*x[1] + 3*x[2] + 1*x[3]
-        let new_1 = gf_add(gf_add(gf_add(
-            chunk_state[0],
-            gf_mul(chunk_state[1], two)),
-            gf_mul(chunk_state[2], three)),
-            chunk_state[3]);
-
-        // new[2] = 1*x[0] + 1*x[1] + 2*x[2] + 3*x[3]
-        let new_2 = gf_add(gf_add(gf_add(
-            chunk_state[0],
-            chunk_state[1]),
-            gf_mul(chunk_state[2], two)),
-            gf_mul(chunk_state[3], three));
-
-        // new[3] = 3*x[0] + 1*x[1] + 1*x[2] + 2*x[3]
-        let new_3 = gf_add(gf_add(gf_add(
-            gf_mul(chunk_state[0], three),
-            chunk_state[1]),
-            chunk_state[2]),
-            gf_mul(chunk_state[3], two));
+        // The order of updates matches the reference algorithm:
+        // x[3] = t01233 + 2*x[0]
+        // x[1] = t01123 + 2*x[2]
+        // x[0] = t01123 + t01
+        // x[2] = t01233 + t23
+        let new_3 = gf_add(t01233, two_x0);
+        let new_1 = gf_add(t01123, two_x2);
+        let new_0 = gf_add(t01123, t01);
+        let new_2 = gf_add(t01233, t23);
 
         // Copy back
         (*state)[offset + 0u] = new_0;