feat(simd): add mag2/4, magsq2/4, move/extract inline fns, update tes…

…ts, readme

packages/simd/README.md

@@ -48,6 +48,10 @@ for sources:
 - `invsqrt4_f32` (!)
 - `madd4_f32`
 - `maddn4_f32`
+- `mag2_f32_aos`
+- `mag4_f32_aos`
+- `magsq2_f32_aos`
+- `magsq4_f32_aos`
 - `max4_f32`
 - `min4_f32`
 - `mix4_f32`

packages/simd/assembly/dot.ts

@@ -1,4 +1,4 @@
-import { hadd2_f32, hadd4_f32 } from "./hadd";
+import { __hadd2_f32, __hadd4_f32 } from "./inline/hadd";
 
 /**
  * Takes two densely packed vec2 AOS buffers `a` and `b`, computes their
@@ -23,7 +23,7 @@ export function dot2_f32_aos(
     const res = out;
     num >>= 1;
     for (; num-- > 0; ) {
-        const m = hadd2_f32(f32x4.mul(v128.load(a), v128.load(b)));
+        const m = __hadd2_f32(f32x4.mul(v128.load(a), v128.load(b)));
         f32.store(out, f32x4.extract_lane(m, 0));
         f32.store(out, f32x4.extract_lane(m, 2), 4);
         out += 8;
@@ -62,7 +62,7 @@ export function dot4_f32_aos(
     sb <<= 2;
     // a1*b1 + a2*b2 + a3*b3 + a4*b4
     for (; num-- > 0; ) {
-        f32.store(out, hadd4_f32(f32x4.mul(v128.load(a), v128.load(b))));
+        f32.store(out, __hadd4_f32(f32x4.mul(v128.load(a), v128.load(b))));
         out += so;
         a += sa;
         b += sb;

packages/simd/assembly/index.ts

@@ -10,6 +10,8 @@ export * from "./clamp";
 export * from "./dot";
 export * from "./madd";
 export * from "./maddn";
+export * from "./mag";
+export * from "./magsq";
 export * from "./max";
 export * from "./min";
 export * from "./mix";

packages/simd/assembly/hadd.ts → packages/simd/assembly/inline/hadd.ts

@@ -9,7 +9,7 @@
  */
 // @ts-ignore: decorator
 @inline
-export function hadd2_f32(v: v128): v128 {
+export function __hadd2_f32(v: v128): v128 {
     return f32x4.add(v, v128.shuffle<f32>(v, v, 1, 0, 3, 2));
 }
 
@@ -23,7 +23,7 @@ export function hadd2_f32(v: v128): v128 {
  */
 // @ts-ignore: decorator
 @inline
-export function hadd4_f32(v: v128): f32 {
+export function __hadd4_f32(v: v128): f32 {
     v = f32x4.add(v, v128.shuffle<f32>(v, v, 2, 3, 0, 1));
     return f32x4.extract_lane(v, 0) + f32x4.extract_lane(v, 1);
 }

packages/simd/assembly/inline/magsq.ts

@@ -0,0 +1,13 @@
+import { __hadd2_f32, __hadd4_f32 } from "./hadd";
+
+// @ts-ignore: decorator
+@inline
+export function __magsq2(v: v128): v128 {
+    return __hadd2_f32(f32x4.mul(v, v));
+}
+
+// @ts-ignore: decorator
+@inline
+export function __magsq4(v: v128): f32 {
+    return __hadd4_f32(f32x4.mul(v, v));
+}

packages/simd/assembly/mag.ts

@@ -0,0 +1,32 @@
+import { __magsq2, __magsq4 } from "./inline/magsq";
+
+export function mag2_f32_aos(out: usize, a: usize, num: usize): usize {
+    const res = out;
+    num >>= 1;
+    for (; num-- > 0; ) {
+        const v = __magsq2(v128.load(a));
+        f32.store(out, sqrt(f32x4.extract_lane(v, 0)));
+        f32.store(out, sqrt(f32x4.extract_lane(v, 2)), 4);
+        out += 8;
+        a += 16;
+    }
+    return res;
+}
+
+export function mag4_f32_aos(
+    out: usize,
+    a: usize,
+    num: usize,
+    so: usize,
+    sa: usize
+): usize {
+    const res = out;
+    so <<= 2;
+    sa <<= 2;
+    for (; num-- > 0; ) {
+        f32.store(out, sqrt(__magsq4(v128.load(a))));
+        out += so;
+        a += sa;
+    }
+    return res;
+}

packages/simd/assembly/magsq.ts

@@ -0,0 +1,32 @@
+import { __magsq2, __magsq4 } from "./inline/magsq";
+
+export function magsq2_f32_aos(out: usize, a: usize, num: usize): usize {
+    const res = out;
+    num >>= 1;
+    for (; num-- > 0; ) {
+        const v = __magsq2(v128.load(a));
+        f32.store(out, f32x4.extract_lane(v, 0));
+        f32.store(out, f32x4.extract_lane(v, 2), 4);
+        out += 8;
+        a += 16;
+    }
+    return res;
+}
+
+export function magsq4_f32_aos(
+    out: usize,
+    a: usize,
+    num: usize,
+    so: usize,
+    sa: usize
+): usize {
+    const res = out;
+    so <<= 2;
+    sa <<= 2;
+    for (; num-- > 0; ) {
+        f32.store(out, __magsq4(v128.load(a)));
+        out += so;
+        a += sa;
+    }
+    return res;
+}

packages/simd/assembly/normalize.ts

@@ -1,4 +1,4 @@
-import { hadd2_f32, hadd4_f32 } from "./hadd";
+import { __magsq2, __magsq4 } from "./inline/magsq";
 
 export function normalize2_f32_aos(
     out: usize,
@@ -10,7 +10,7 @@ export function normalize2_f32_aos(
     num >>= 1;
     for (; num-- > 0; ) {
         const v = v128.load(a);
-        let vm = hadd2_f32(f32x4.mul(v, v));
+        let vm = __magsq2(v);
         const m1 = f32x4.extract_lane(vm, 0);
         const m2 = f32x4.extract_lane(vm, 2);
         vm = f32x4.replace_lane(
@@ -43,7 +43,7 @@ export function normalize4_f32_aos(
     const res = out;
     for (; num-- > 0; ) {
         const v = v128.load(a);
-        const mag = hadd4_f32(f32x4.mul(v, v));
+        const mag = __magsq4(v);
         v128.store(
             out,
             mag > f32.EPSILON

packages/simd/assembly/sum.ts

@@ -1,4 +1,4 @@
-import { hadd4_f32 } from "./hadd";
+import { __hadd4_f32 } from "./inline/hadd";
 
 export function sum4_f32(a: usize, num: usize, sa: usize): f32 {
     sa <<= 2;
@@ -7,5 +7,5 @@ export function sum4_f32(a: usize, num: usize, sa: usize): f32 {
         acc = f32x4.add(acc, v128.load(a));
         a += sa;
     }
-    return hadd4_f32(acc);
+    return __hadd4_f32(acc);
 }

packages/simd/src/api.ts

@@ -110,6 +110,16 @@ export interface SIMD {
     // prettier-ignore
     maddn4_f32(out: number, a: number, b: number, c: number, num: number, so: number, sa: number, sc: number): number;
 
+    mag2_f32_aos(out: number, a: number, num: number): number;
+
+    magsq2_f32_aos(out: number, a: number, num: number): number;
+
+    // prettier-ignore
+    mag4_f32_aos( out: number, a: number, num: number, so: number, sa: number): number;
+
+    // prettier-ignore
+    magsq4_f32_aos( out: number, a: number, num: number, so: number, sa: number): number;
+
     // prettier-ignore
     max4_f32(out: number, a: number, b: number, num: number, so: number, sa: number, sb: number): number;
 

packages/simd/test/index.ts

@@ -170,6 +170,29 @@ simd.msubn4_f32(1024, 0, 11, 32, 2, 4, 4, 4);
 // prettier-ignore
 assertEqual(res_f32(1024, 8), [-89, -178, -267, -356, -445, -534, -623, -712]);
 
+// magsq2
+// magsq4
+simd.f32.set([1, 2, 10, 20, -100, 200, 100, -200]);
+simd.magsq2_f32_aos(1024, 0, 4);
+assertEqualDelta(res_f32(1024, 4), [
+    1 * 1 + 2 * 2,
+    10 * 10 + 20 * 20,
+    100 * 100 + 200 * 200,
+    100 * 100 + 200 * 200
+]);
+simd.mag2_f32_aos(1024, 0, 4);
+assertEqualDelta(res_f32(1024, 4), [
+    Math.sqrt(5),
+    Math.sqrt(500),
+    Math.sqrt(50000),
+    Math.sqrt(50000)
+]);
+
+simd.magsq4_f32_aos(1024, 0, 2, 1, 4);
+assertEqualDelta(res_f32(1024, 2), [505, 100000]);
+simd.mag4_f32_aos(1024, 0, 2, 1, 4);
+assertEqualDelta(res_f32(1024, 2), [Math.sqrt(505), Math.sqrt(100000)]);
+
 // mix4_f32
 // mixn4_f32
 // prettier-ignore