fixes to gelu and swiglu approx

Author: gbenroscience

parser-ng-bench/src/main/java/com/github/gbenroscience/parser/ng/bench/ActivationBench.java

@@ -1,9 +1,5 @@
 package com.github.gbenroscience.parser.ng.bench;
 
-/**
- *
- * @author GBEMIRO
- */
 import com.github.gbenroscience.parser.MathExpression;
 import com.github.gbenroscience.simd.turbo.tools.SIMDVectorTurboEvaluator;
 import org.openjdk.jmh.annotations.*;
@@ -20,45 +16,41 @@
 import org.openjdk.jmh.runner.options.TimeValue;
 
 /**
- * Use this to build in the parser-ng-pro directory mvn clean install -Dgpg.skip
- * Use this to run the benchmarks
- *
- * java -jar target/benchmarks.jar SIMDTurboBench -prof perfasm
+ * Use this to build in the parser-ng-pro directory: mvn clean install -Dgpg.skip
+ * Use this to run the benchmarks:
  *
- * java -jar target/benchmarks.jar SIMDTurboBench -prof perfasm >
- * perf_output.txt Now find your specific method's assembly grep -A 50
- * "parserNG" perf_output.txt
+ * java -jar target/benchmarks.jar ActivationBench -prof perfasm
  *
  * @author GBEMIRO
  */
-@BenchmarkMode(Mode.AverageTime) // ns/op
+@BenchmarkMode(Mode.AverageTime) 
 @OutputTimeUnit(TimeUnit.NANOSECONDS)
 @State(Scope.Benchmark)
 @Fork(value = 3, jvmArgsAppend = {
     "-Xms4g", "-Xmx4g",
     "-XX:+UnlockDiagnosticVMOptions",
     "--add-modules", "jdk.incubator.vector", "-XX:+UnlockDiagnosticVMOptions"
-// "-XX:+PrintAssembly", "-XX:CompileCommand=print,*SIMDCompositeExpression.applyMatrixKernel" // uncomment if you have hsdis
 })
 @Warmup(iterations = 5, time = 2)
 @Measurement(iterations = 5, time = 2)
+@OperationsPerInvocation(ActivationBench.N) // Tells JMH to divide total time by N to output ns/element
 public class ActivationBench {
 
-    @Param({"512", "1024", "2048", "4096"})
-    public int sz;
-
+    // Must be a compile-time constant to be used in the annotation above.
+    // 1048576 elements represents a square matrix of 1024 x 1024.
+    public static final int N = 2048 * 2048; 
+    
+    private final int sz = (int) Math.sqrt(N);
     private SIMDCompositeExpression evaluator;
     private FlatMatrixF in1;
     private FlatMatrixF in2;
     private FlatMatrixF out;
-    private int N;
 
     @Setup(Level.Trial)
     public void setup() throws Throwable {
         MathExpression me = new MathExpression("x * 0.5 * (1 + tanh(0.79788456 * (x + 0.044715 * x * x * x)))");
         evaluator = (SIMDCompositeExpression) new SIMDVectorTurboEvaluator(me).compile();
 
-        N = sz * sz;
         in1 = new FlatMatrixF(sz, sz);
         in2 = new FlatMatrixF(sz, sz);
         out = new FlatMatrixF(sz, sz);
@@ -68,30 +60,23 @@ public void setup() throws Throwable {
 
     @Setup(Level.Iteration) // C2 can't optimize across iterations
     public void dirtyInput() {
-        // JMH will call this before each measurement iteration
-        // This makes the input unprovable without timing it
+        // JMH calls this before each measurement iteration to make inputs unpredictable
         for (int j = 0; j < N; j++) {
             in1.data[j] = ThreadLocalRandom.current().nextFloat();
-            in2.data[j] = in1.data[j]*2.13f;
+            in2.data[j] = in1.data[j] * 2.13f;
         }
     }
 
     @Benchmark
     public void gelu(Blackhole bh) {
         evaluator.applyMatrixKernel(new FlatMatrixF[]{in1}, out, "gelu");
-        // Consume the whole output so C2 can't DCE
-        /*
-        for (int j = 0; j < N; j++) {
-            bh.consume(Float.floatToRawIntBits(out.data[j]));
-        }
-        */
-         bh.consume(out.data[0]);
+        bh.consume(out); // Completely safe, zero-overhead way to prevent Dead Code Elimination
     }
-        @Benchmark
+
+    @Benchmark
     public void swiglu(Blackhole bh) {
         evaluator.applyMatrixKernel(new FlatMatrixF[]{in1, in2}, out, "swiglu");
-        // Consume the whole output so C2 can't DCE
-               bh.consume(out.data[0]);
+        bh.consume(out); 
     }
 
     public static void main(String[] args) throws RunnerException {
@@ -112,4 +97,4 @@ public static void main(String[] args) throws RunnerException {
 
         new Runner(configurations).run();
     }
-}
+}

parser-ng-simd/src/main/java/com/github/gbenroscience/simd/turbo/tools/FlatMatrixF.java

@@ -271,28 +271,84 @@ public void geluInPlace() {
         }
     }
 
-    public void geluInPlaceHighSpeed() {
-        int len = rows * cols, idx = 0;
-        if (HAS_VECTOR && isContiguous()) {
+    /**
+     * GeGLU in-place: output = inputs[0] ⊙ GELU(inputs[1])
+     *
+     * @param gate
+     */
+    public void gegluInPlaceOld(FlatMatrixF gate) {   // gate is inputs[1]
+        if (rows != gate.rows || cols != gate.cols) {
+            throw new IllegalArgumentException("Shape mismatch for GeGLU");
+        }
+
+        int len = rows * cols;
+        int idx = 0;
+
+        if (HAS_VECTOR && isContiguous() && gate.isContiguous()) {
             for (; idx < F_SPECIES.loopBound(len); idx += VF_LEN) {
-                var v = FloatVector.fromArray(F_SPECIES, data, offset + idx);
+                var value = FloatVector.fromArray(F_SPECIES, this.data, this.offset + idx);
+                var g = FloatVector.fromArray(F_SPECIES, gate.data, gate.offset + idx);
 
-                // x + 0.044715 * x^3
-                var x2 = v.mul(v);
-                var x3 = x2.mul(v);
-                var inner = v.add(x3.mul(V_GELU_C2)).mul(V_GELU_C1);
+                // GELU on gate
+                var g2 = g.mul(g);
+                var g3 = g2.mul(g);
+                var inner = g.add(g3.mul(V_GELU_C2)).mul(V_GELU_C1);
+                var t = inner.lanewise(VectorOperators.TANH);
+                var geluGate = g.mul(V_HALF).mul(t.add(V_ONE));
 
-                // tanh_poly(inner) = inner * (27 + inner^2) / (27 + 9*inner^2)
-                var t = tanh256(inner); // <- No scalar calls here
+                // value * GELU(gate)
+                value.mul(geluGate).intoArray(this.data, this.offset + idx);
+            }
+        }
 
-                v.mul(V_HALF).mul(t.add(V_ONE)).intoArray(data, offset + idx);
+        // Scalar fallback
+        for (; idx < len; idx++) {
+            float value = this.data[this.offset + idx];
+            float g = gate.data[gate.offset + idx];
+
+            float g3 = g * g * g;
+            float t = (float) Math.tanh(GELU_C1 * (g + GELU_C2 * g3));
+            float geluGate = 0.5f * g * (1.0f + t);
+
+            this.data[this.offset + idx] = value * geluGate;
+        }
+    }
+
+    public void gegluInPlace(FlatMatrixF gate) {
+        if (rows != gate.rows || cols != gate.cols) {
+            throw new IllegalArgumentException("Shape mismatch for GeGLU");
+        }
+
+        int len = rows * cols;
+        int idx = 0;
+
+        if (HAS_VECTOR && isContiguous() && gate.isContiguous()) {
+            for (; idx < F_SPECIES.loopBound(len); idx += VF_LEN) {
+                var value = FloatVector.fromArray(F_SPECIES, this.data, this.offset + idx);
+                var g = FloatVector.fromArray(F_SPECIES, gate.data, gate.offset + idx);
+
+                // GELU(gate)
+                var g2 = g.mul(g);
+                var g3 = g2.mul(g);
+                var inner = g.add(g3.mul(V_GELU_C2)).mul(V_GELU_C1);
+                var t = inner.lanewise(VectorOperators.TANH);
+                var geluGate = g.mul(V_HALF).mul(t.add(V_ONE));
+
+                // value * GELU(gate)
+                value.mul(geluGate).intoArray(this.data, this.offset + idx);
             }
         }
+
+        // Scalar fallback
         for (; idx < len; idx++) {
-            float x = data[offset + idx];
-            float x3 = x * x * x;
-            float t = (float) Math.tanh(GELU_C1 * (x + GELU_C2 * x3));
-            data[offset + idx] = 0.5f * x * (1.0f + t);
+            float value = this.data[this.offset + idx];
+            float g = gate.data[gate.offset + idx];
+
+            float g3 = g * g * g;
+            float t = (float) Math.tanh(GELU_C1 * (g + GELU_C2 * g3));
+            float geluGate = 0.5f * g * (1.0f + t);
+
+            this.data[this.offset + idx] = value * geluGate;
         }
     }
 
@@ -658,9 +714,13 @@ private static void matmulInPlaceScalar(FlatMatrixF A, FlatMatrixF B, FlatMatrix
     }
 
     public static void randomFill(FlatMatrixF A) {
+        randomFill(A, 101);
+    }//end method randomFill
+
+    public static void randomFill(FlatMatrixF A, int max) {
         ThreadLocalRandom ran = ThreadLocalRandom.current();
         for (int i = 0; i < A.data.length; i++) {
-            A.data[i] = 1 + ran.nextFloat(101);
+            A.data[i] = 1 + ran.nextFloat(max);
         }
     }//end method randomFill
 
@@ -1009,66 +1069,77 @@ private static void weightedSum(float[] scores, FlatMatrixF V, FlatMatrixF out,
         }
     }
 
-    // 100% Vector ops. Compiles to vmulpd vaddpd vdivpd only
-    private static FloatVector tanh256(FloatVector x) {
-        var x2 = x.mul(x);
-        var num = x2.add(27.0f);          // 27 + x^2
-        var den = x2.mul(9.0f).add(27.0f); // 27 + 9*x^2
-        return x.mul(num).div(den);       // x * num / den
-    }
-
-    /**
-     * Fast sigmoid: 1 / (1 + exp(-x)) Max error ~1e-3. ~3x faster than Math.exp
-     * scalar. Uses expm1 + rational approx to stay stable.
-     *
-     * @param x
-     * @return
-     */
-    public static FloatVector sigmoid(FloatVector x) {
-        // 1. Clamp: sig(-inf)=0, sig(+inf)=1
-        x = x.max(MIN_X).min(MAX_X);
-
-        // 2. Use 1 / (1 + exp(-x)) = exp(x) / (1 + exp(x))
-        // For x>=0: exp(-x) is small, more stable
-        var negX = x.neg();
-        var expNegX = fastExp(negX); // exp(-x)
-        return ONE.div(ONE.add(expNegX));
-    }
-
     /**
      * Fast exp approx. Max error ~2e-3 on [-10, 10] Based on Pade approx:
      * exp(x) = 2^(x * log2e)
+     *
+     * @param x
      */
     public static FloatVector fastExp(FloatVector x) {
+        // Clamp input to avoid float underflow/overflow bounds [-88.0f, 88.0f]
+        x = x.lanewise(VectorOperators.MAX, -88.0f)
+                .lanewise(VectorOperators.MIN, 88.0f);
+
         final FloatVector LOG2E = FloatVector.broadcast(F_SPECIES, 1.44269504f);
         final FloatVector LN2 = FloatVector.broadcast(F_SPECIES, 0.69314718f);
+        final FloatVector ONE = FloatVector.broadcast(F_SPECIES, 1.0f);
 
+        // x * log2(e)
         var fx = x.mul(LOG2E);
 
-        // float -> int, truncate toward 0
+        // STEP 1: Perform manual floor via F2I truncation toward zero
         IntVector ni = (IntVector) fx.convertShape(VectorOperators.F2I, I_SPECIES, 0);
-
-        // F2I truncates, we need floor. Correct if fx < int(fx)
         FloatVector ni_f = (FloatVector) ni.convertShape(VectorOperators.I2F, F_SPECIES, 0);
+
+        // For negative values, truncation goes the wrong way. 
+        // If fx < ni_f, we must subtract 1 from the integer vector.
         var needsDec = fx.compare(VectorOperators.LT, ni_f);
-        ni = ni.sub(needsDec.toVector().reinterpretAsInts().lanewise(VectorOperators.ASHR, 31)); // mask->int, subtract 1 where true
+        var oneInt = IntVector.broadcast(I_SPECIES, 1);
+
+        // CORRECT FIX: Cast the Float mask to an Integer mask matching I_SPECIES
+        VectorMask<Integer> intMask = needsDec.cast(I_SPECIES);
+        ni = ni.sub(oneInt, intMask);
 
-        var f = fx.sub(ni_f); // fractional part in [0,1)
+        // STEP 2: Recompute the corrected float representation of the floor
+        ni_f = (FloatVector) ni.convertShape(VectorOperators.I2F, F_SPECIES, 0);
 
-        // 2^f poly: 1 + y + y^2/2 + y^3/6, y = f*ln2
+        // Fractional part is now guaranteed to be in [0, 1)
+        var f = fx.sub(ni_f);
+
+        // STEP 3: 4th-order Polynomial approximation for 2^f
         var y = f.mul(LN2);
         var y2 = y.mul(y);
         var y3 = y2.mul(y);
-        var poly = ONE.add(y).add(y2.mul(0.5f)).add(y3.mul(0.16666667f));
+        var y4 = y3.mul(y);
+
+        var poly = ONE.add(y)
+                .add(y2.mul(0.5f))
+                .add(y3.mul(0.16666667f))
+                .add(y4.mul(0.041666668f));
 
-        // 2^n via bit cast: (n+127)<<23
-        IntVector biasedExp = ni.add(127);
-        IntVector expBits = biasedExp.lanewise(VectorOperators.LSHL, 23);
-        FloatVector pow2n = (FloatVector) expBits.convertShape(VectorOperators.I2F, F_SPECIES, 0);
+        // STEP 4: 2^ni via raw bit manipulation: (ni + 127) << 23
+        var biased = ni.add(127);
+        var bits = biased.lanewise(VectorOperators.LSHL, 23);
+
+        // Pure bit-cast from integer registers to float registers
+        FloatVector pow2n = bits.reinterpretAsFloats();
 
         return poly.mul(pow2n);
     }
 
+    /**
+     * Fast sigmoid: 1 / (1 + exp(-x)) Max error ~1e-3. ~3x faster than Math.exp
+     * scalar. Uses expm1 + rational approx to stay stable.
+     *
+     * @param x
+     * @return
+     */
+    public static FloatVector sigmoid(FloatVector x) {
+        final FloatVector ONE = FloatVector.broadcast(F_SPECIES, 1.0f);
+        var negX = x.neg();
+        return ONE.div(ONE.add(fastExp(negX)));
+    }
+
 
 //////////////////////////////MHA-ATTENTION-DONE/////////////////////////////////////////////
 }

parser-ng-simd/src/main/java/com/github/gbenroscience/simd/turbo/tools/VectorTurboEvaluator.java

@@ -1713,7 +1713,7 @@ private void evaluateTile(double[][] _2DVariables, int dataSize, double[] output
          */
         @Override
         public void applyBulkParallel(double[] flatVariables, double[] output) {
-           
+
             final int length = output.length;
 
             if (length < PARALLEL_OPS_THRESHOLD) {
@@ -1785,8 +1785,22 @@ public void applyMatrixKernel(FlatMatrixF[] inputs, FlatMatrixF output, String o
                     if (inputs[0] != output) {
                         System.arraycopy(inputs[0].data, inputs[0].offset, output.data, output.offset, output.rows * output.cols);
                     }
-                    //output.geluInPlace();
-                    output.geluInPlaceHighSpeed();
+                    output.geluInPlace();
+                    //output.geluInPlaceHighSpeed();
+                }
+                case "geglu", "geglu_in_place" -> {
+                    if (inputs.length < 2) {
+                        throw new IllegalArgumentException("GeGLU requires 2 inputs");
+                    }
+
+                    // Copy first input to output if they are different
+                    if (inputs[0] != output) {
+                        System.arraycopy(inputs[0].data, inputs[0].offset,
+                                output.data, output.offset,
+                                output.rows * output.cols);
+                    }
+
+                    output.gegluInPlace(inputs[1]);   // <- gate
                 }
                 // === New Q8 + Attention kernels ===
                 case "q8_quantize" -> {
@@ -2222,7 +2236,7 @@ private boolean isMatrixKernel(String name, int arity) {
             // BLAS
             case "matmul", "gemm", "mm" ->
                 arity == 3;
-            case "gemv", "matvec" ->
+            case "gemv", "matvec","geglu" ->
                 arity == 2;
             case "dot", "inner", "axpy" ->
                 arity == 3;