loopless string serialization (for the remainder)

JairusSW · Dec 27, 2024 · 45a0612 · 45a0612
1 parent bd96a29
commit 45a0612
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 11 deletions.
diff --git a/assembly/__benches__/misc.bench.ts b/assembly/__benches__/misc.bench.ts
@@ -1,8 +1,8 @@
 import { deserializeString_SIMD } from "../deserialize/simd/string";
 import { serializeString_SIMD } from "../serialize/simd/string";
 import { bench } from "as-bench/assembly/index"
-const str = '""""""""';
-const str2 = '"\"\"\"\"\"\"\"\""';
+const str = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!@#$%^&*()\\\"\t\r\f\n\u0000';
+const str2 = '"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!@#$%^&*()\\\\\\"\\t\\r\\f\\n\\u0000"';
 
 // bench("Serialize String (Simple)", () => {
 //   serializeString(str);
@@ -13,11 +13,13 @@ const str2 = '"\"\"\"\"\"\"\"\""';
 //     bs.reset();
 // });
 
-const out = new ArrayBuffer(16);
+const out = new ArrayBuffer(256);
 bench("Serialize String (SIMD)", () => {
+  // ~5.07GB/s
   serializeString_SIMD(str, changetype<usize>(out));
 });
 
 bench("Deserialize String (SIMD)", () => {
+  // ~4.03GB/s
   deserializeString_SIMD(str2, changetype<usize>(out));
 });
diff --git a/assembly/serialize/simd/string.ts b/assembly/serialize/simd/string.ts
@@ -43,6 +43,7 @@ const SPLAT_34 = i16x8.splat(34); /* " */
 const SPLAT_92 = i16x8.splat(92); /* \ */
 
 const SPLAT_32 = i16x8.splat(32); /* [ESC] */
+const SPLAT_0 = i16x8.splat(0); /* 0 */
 
 /**
  * Serializes strings into their JSON counterparts using SIMD operations
@@ -91,11 +92,94 @@ export function serializeString_SIMD(src: string, dst: usize): usize {
       }
     }
 
-    src_ptr += 16;
-    dst_ptr += 16;
+    src_ptr += 16; dst_ptr += 16;
   }
 
-  while (src_ptr < src_end) {
+  let rem = src_end - src_ptr;
+
+  if (rem & 8) {
+    const block = v128.load64_zero(src_ptr);
+    v128.store64_lane(dst_ptr, block, 0);
+
+    const backslash_indices = i16x8.eq(block, SPLAT_92);
+    const quote_indices = i16x8.eq(block, SPLAT_34);
+    const escape_indices = i16x8.lt_u(block, SPLAT_32);
+    const zero_indices = i16x8.eq(block, SPLAT_0);
+    const sieve = v128.and(v128.or(v128.or(backslash_indices, quote_indices), escape_indices), v128.not(zero_indices));
+
+    let mask = i16x8.bitmask(sieve);
+    while (mask != 0) {
+      let lane_index = ctz(mask) << 1;
+      const dst_offset = dst_ptr + lane_index;
+      const src_offset = src_ptr + lane_index;
+      const code = load<u16>(src_offset) << 2;
+      const escaped = load<u32>(ESCAPE_TABLE + code);
+      mask &= mask - 1;
+
+      if ((escaped & 0xFFFF) != BACK_SLASH) {
+        store<u64>(dst_offset, 13511005048209500);
+        store<u32>(dst_offset, escaped, 8);
+        while (lane_index < 6) {
+          store<u8>(dst_ptr + lane_index, load<u8>(src_ptr + lane_index, 2), 12);
+          lane_index += 2;
+        }
+        dst_ptr += 10;
+      } else {
+        store<u32>(dst_offset, escaped);
+
+        while (lane_index < 6) {
+          store<u8>(dst_ptr + lane_index, load<u8>(src_ptr + lane_index, 2), 4);
+          lane_index += 2;
+        }
+        dst_ptr += 2;
+      }
+    }
+
+    dst_ptr += 8; src_ptr += 8;
+  }
+  if (rem & 4) {
+    const block = load<u32>(src_ptr);
+
+    const codeA = block & 0xFFFF;
+    const codeB = block >> 16;
+
+    if (codeA == 92 || codeA == 34 || codeA < 32) {
+      const escaped = load<u32>(ESCAPE_TABLE + (codeA << 2));
+
+      if ((escaped & 0xFFFF) != BACK_SLASH) {
+        store<u64>(dst_ptr, 13511005048209500);
+        store<u32>(dst_ptr, escaped, 8);
+        dst_ptr += 12;
+      } else {
+        store<u32>(dst_ptr, escaped);
+        dst_ptr += 4;
+      }
+
+    } else {
+      store<u16>(dst_ptr, codeA);
+      dst_ptr += 2;
+    }
+
+    if (codeB == 92 || codeB == 34 || codeB < 32) {
+      const escaped = load<u32>(ESCAPE_TABLE + (codeB << 2));
+
+      if ((escaped & 0xFFFF) != BACK_SLASH) {
+        store<u64>(dst_ptr, 13511005048209500);
+        store<u32>(dst_ptr, escaped, 8);
+        dst_ptr += 12;
+      } else {
+        store<u32>(dst_ptr, escaped);
+        dst_ptr += 4;
+      }
+
+    } else {
+      store<u16>(dst_ptr, codeB);
+      dst_ptr += 2;
+    }
+
+    src_ptr += 4;
+  }
+  if (rem & 2) {
     const code = load<u16>(src_ptr);
     if (code == 92 || code == 34 || code < 32) {
       const escaped = load<u32>(ESCAPE_TABLE + (code << 2));
@@ -104,17 +188,14 @@ export function serializeString_SIMD(src: string, dst: usize): usize {
         store<u64>(dst_ptr, 13511005048209500);
         store<u32>(dst_ptr, escaped, 8);
         dst_ptr += 12;
-        src_ptr += 2;
       } else {
         store<u32>(dst_ptr, escaped);
         dst_ptr += 4;
-        src_ptr += 2;
       }
 
     } else {
       store<u16>(dst_ptr, code);
       dst_ptr += 2;
-      src_ptr += 2;
     }
   }
 

diff --git a/assembly/test.ts b/assembly/test.ts
@@ -10,13 +10,15 @@ import { JSON } from ".";
 //   public z: T;
 // }
 
-let a = "\u0000\u0001\u0002\u0003\u0004\u0005\u0006\u0007\u0008\u0009\u000a\u000b\u000c\u000d\u000e\u000f\u000f\u0011\u0012\u0013\u0014\u0015\u0016\u0017\u0018\u0019\u001a\u001b\u001c\u001d\u001e\u001f";
+let a = "12\"345678123456";
 const b = "world";
-JSON.stringifyTo(a, a)
+let c = "000000000000000000000000000000000000000";
+JSON.stringifyTo(a, c)
 // console.log(JSON.stringifyTo(a, a));
 
 console.log("A: " + a.toString());
 console.log("B: " + b.toString());
+console.log("C: " + c.toString());
 
 // console.log(new Vec3<i32>().__SERIALIZE())