perf(joint): infer on burn-flex (burn 0.21) — ~5.6x faster cold inference

Upgrade burn 0.19.1 -> 0.21.0 and switch the joint tagger's runtime
inference backend from burn-ndarray to burn-flex, a fast pure-Rust CPU
backend (gemm + SIMD). flex is now the only inference backend, on every
target including wasm32 (no rayon + critical-section atomics shim there).

  - cold inference ~5.6x faster (0.53ms vs 2.97ms/sentence), byte-identical
    output; real-prose parse_essay 5.9x faster, lint_essay_uncached 3.4x.
  - retrain the embedded model on 0.21 (UPOS 0.9463 / NP F1 0.8872); the
    0.19 model.mpk can't load on 0.21 (version guard + new LSTM
    gate_activation field).
  - burn 0.21 moved Backend::Device to a BackendTypes supertrait; name the
    concrete device type in runtime.rs instead.
  - regenerate_english_joint: AutoGraphicsApi instead of hardcoded Metal so
    training runs on any wgpu GPU (Metal/Vulkan/DX12).
  - save_then_load_preserves_forward: rewrite as an fp16->fp16 round-trip
    equality check (the old fp32-vs-fp16 threshold failed on both backends
    under 0.21's RNG, not a flex issue).
  - char/suffix vocab JSON now pretty-printed with ordered keys for
    diff-friendly artifacts; key order is cosmetic (ids are the values).

burn-ndarray remains only for the legacy BurnChunker and tests.

Author: marcbachmann

Cargo.lock

@@ -11,7 +11,7 @@ harper-pos-utils = { path = "../harper-pos-utils/", version = "2.0.0" }
 serde_json = "1.0.150"
 # Only pulled in by the `training` feature (GPU regeneration tooling). Inference
 # and tests never need it — the trained model is embedded as bytes and run on CPU.
-burn = { version = "0.19.1", default-features = false, optional = true }
+burn = { version = "0.21.0", default-features = false, optional = true }
 
 [build-dependencies]
 serde_json = "1.0.150"

harper-brill/Cargo.toml

@@ -25,7 +25,7 @@ use std::process::Command;
 use std::rc::Rc;
 use std::time::Instant;
 
-use burn::backend::wgpu::{RuntimeOptions, WgpuDevice, graphics::Metal, init_setup};
+use burn::backend::wgpu::{RuntimeOptions, WgpuDevice, graphics::AutoGraphicsApi, init_setup};
 use burn::backend::{Autodiff, wgpu::Wgpu};
 
 use harper_pos_utils::conllu_utils::extract_records_from_files;
@@ -176,9 +176,13 @@ fn main() {
         suffix_vocab.len()
     );
 
-    // Explicit Metal init must precede any tensor operations on macOS.
+    // Explicit wgpu init must precede any tensor operations on macOS.
+    // `AutoGraphicsApi` selects the platform's native graphics API at runtime —
+    // Metal on macOS, Vulkan on Linux (NVIDIA/AMD/Intel), DX12 on Windows — so the
+    // same example runs on any GPU wgpu supports, not just Apple Metal. Override
+    // in tests via the `AUTO_GRAPHICS_BACKEND` env var.
     let device = WgpuDevice::DefaultDevice;
-    init_setup::<Metal>(&device, RuntimeOptions::default());
+    init_setup::<AutoGraphicsApi>(&device, RuntimeOptions::default());
 
     let cfg = JointTrainConfig {
         char_dim: ARCH.char_dim,

harper-brill/examples/regenerate_english_joint.rs

@@ -1 +1,218 @@
-{"ي":132,"V":59,">":97,"²":211,"é":81,"T":32,"±":112,"ι":155,"è":169,"ę":198,"U":44,"³":90,"…":106,"况":182,"0":53,"í":110,"澹":188,"S":19,";":68,"?":73,"×":190,"樋":137,"<":96,"R":58,"Q":28,"=":93,"k":17,"3":38,"ə":149,"ë":215,"½":119,"ɛ":115,"•":172,"l":3,"2":40,"1":50,"m":7,"6":52,"h":20,"ف":193,"W":49,"̃":123,"奈":141,"j":63,"4":67,"5":69,"w":27,"°":202,"i":12,"ح":130,"v":43,"#":78,"ε":164,"ʻ":171,"u":22,"!":42,"æ":173,"ن":194,"\"":74,"t":24,"一":139,"ʃ":121,"c":13,"n":8,"ç":105,"/":76,"ه":128,"津":142,"β":154,"o":15,"b":21,"Ἱ":163,"远":186,"ʁ":116,"a":6,".":30,"`":88,"f":14,"8":62,"ú":147,"ο":156,"Σ":158,"谿":179,"'":57,"¢":175,"§":189,"e":10,"à":126,":":9,"d":18,"9":51,"£":109,"ˈ":117,"%":79,"y":29,"Ḥ":134,"á":103,"夏":143,"μ":167,"ū":146,"z":56,"$":82,"g":33,"ʊ":118,"x":71,"&":80,"ī":135,"ν":162,"â":170,"和":183,"î":216,"O":64,"Ã":89,"{":98,"–":108,"(":65,"N":60,"@":94,"ا":129,")":66,"}":102,"­":107,"ό":161,"*":87,"|":99,"清":185,"ω":159,"+":91,"空":177,"å":205,"~":100,"—":84,"口":138,"ź":197,"A":2,"C":70,"Ε":150,"p":25,"B":39,"-":4,"q":26,"J":46,"ρ":120,"€":191,"I":37,"r":11,",":23,"ô":125,"ς":157,"ɔ":168,"古":187,"s":16,"σ":152,"✉":208,"ö":217,"琴":181,"M":41,"د":131,"L":75,"H":48,"ó":111,"子":144,"έ":153,"ώ":165,"K":61,"ñ":113,"ɑ":122,"司":176,"ò":114,"山":180,"ż":199,"ś":201,"ὐ":151,"图":178,"ä":213,"[":31,"‘":92,"ō":136,"’":83,"7":72,"Ō":145,"φ":160,"ز":127,"É":124,"υ":166,"Ā":195,"ø":206,"☏":207,"º":210,"ß":212,"ā":133,"葉":140,"☎":203,"ص":192,"P":36,"E":47,"”":86,"·":101,"ł":200,"D":35,"_":77,"“":85,"ʒ":148,"Ł":196,"^":104,"F":45,"Ø":204,"œ":214,"]":34,"G":55,"静":184,"X":95,"Æ":209,"Y":54,"č":174,"Z":5}
+{
+  "!": 42,
+  "\"": 74,
+  "#": 78,
+  "$": 82,
+  "%": 79,
+  "&": 80,
+  "'": 57,
+  "(": 65,
+  ")": 66,
+  "*": 87,
+  "+": 91,
+  ",": 23,
+  "-": 4,
+  ".": 30,
+  "/": 76,
+  "0": 53,
+  "1": 50,
+  "2": 40,
+  "3": 38,
+  "4": 67,
+  "5": 69,
+  "6": 52,
+  "7": 72,
+  "8": 62,
+  "9": 51,
+  ":": 9,
+  ";": 68,
+  "<": 96,
+  "=": 93,
+  ">": 97,
+  "?": 73,
+  "@": 94,
+  "A": 2,
+  "B": 39,
+  "C": 70,
+  "D": 35,
+  "E": 47,
+  "F": 45,
+  "G": 55,
+  "H": 48,
+  "I": 37,
+  "J": 46,
+  "K": 61,
+  "L": 75,
+  "M": 41,
+  "N": 60,
+  "O": 64,
+  "P": 36,
+  "Q": 28,
+  "R": 58,
+  "S": 19,
+  "T": 32,
+  "U": 44,
+  "V": 59,
+  "W": 49,
+  "X": 95,
+  "Y": 54,
+  "Z": 5,
+  "[": 31,
+  "]": 34,
+  "^": 104,
+  "_": 77,
+  "`": 88,
+  "a": 6,
+  "b": 21,
+  "c": 13,
+  "d": 18,
+  "e": 10,
+  "f": 14,
+  "g": 33,
+  "h": 20,
+  "i": 12,
+  "j": 63,
+  "k": 17,
+  "l": 3,
+  "m": 7,
+  "n": 8,
+  "o": 15,
+  "p": 25,
+  "q": 26,
+  "r": 11,
+  "s": 16,
+  "t": 24,
+  "u": 22,
+  "v": 43,
+  "w": 27,
+  "x": 71,
+  "y": 29,
+  "z": 56,
+  "{": 98,
+  "|": 99,
+  "}": 102,
+  "~": 100,
+  "¢": 175,
+  "£": 109,
+  "§": 189,
+  "­": 107,
+  "°": 202,
+  "±": 112,
+  "²": 211,
+  "³": 90,
+  "·": 101,
+  "º": 210,
+  "½": 119,
+  "Ã": 89,
+  "Æ": 209,
+  "É": 124,
+  "×": 190,
+  "Ø": 204,
+  "ß": 212,
+  "à": 126,
+  "á": 103,
+  "â": 170,
+  "ä": 213,
+  "å": 205,
+  "æ": 173,
+  "ç": 105,
+  "è": 169,
+  "é": 81,
+  "ë": 215,
+  "í": 110,
+  "î": 216,
+  "ñ": 113,
+  "ò": 114,
+  "ó": 111,
+  "ô": 125,
+  "ö": 217,
+  "ø": 206,
+  "ú": 147,
+  "Ā": 195,
+  "ā": 133,
+  "č": 174,
+  "ę": 198,
+  "ī": 135,
+  "Ł": 196,
+  "ł": 200,
+  "Ō": 145,
+  "ō": 136,
+  "œ": 214,
+  "ś": 201,
+  "ū": 146,
+  "ź": 197,
+  "ż": 199,
+  "ɑ": 122,
+  "ɔ": 168,
+  "ə": 149,
+  "ɛ": 115,
+  "ʁ": 116,
+  "ʃ": 121,
+  "ʊ": 118,
+  "ʒ": 148,
+  "ʻ": 171,
+  "ˈ": 117,
+  "̃": 123,
+  "Ε": 150,
+  "Σ": 158,
+  "έ": 153,
+  "β": 154,
+  "ε": 164,
+  "ι": 155,
+  "μ": 167,
+  "ν": 162,
+  "ο": 156,
+  "ρ": 120,
+  "ς": 157,
+  "σ": 152,
+  "υ": 166,
+  "φ": 160,
+  "ω": 159,
+  "ό": 161,
+  "ώ": 165,
+  "ا": 129,
+  "ح": 130,
+  "د": 131,
+  "ز": 127,
+  "ص": 192,
+  "ف": 193,
+  "ن": 194,
+  "ه": 128,
+  "ي": 132,
+  "Ḥ": 134,
+  "Ἱ": 163,
+  "ὐ": 151,
+  "–": 108,
+  "—": 84,
+  "‘": 92,
+  "’": 83,
+  "“": 85,
+  "”": 86,
+  "•": 172,
+  "…": 106,
+  "€": 191,
+  "☎": 203,
+  "☏": 207,
+  "✉": 208,
+  "一": 139,
+  "况": 182,
+  "口": 138,
+  "古": 187,
+  "司": 176,
+  "和": 183,
+  "图": 178,
+  "夏": 143,
+  "奈": 141,
+  "子": 144,
+  "山": 180,
+  "樋": 137,
+  "津": 142,
+  "清": 185,
+  "澹": 188,
+  "琴": 181,
+  "空": 177,
+  "葉": 140,
+  "谿": 179,
+  "远": 186,
+  "静": 184
+}

harper-brill/finished_joint/char_vocab.json

@@ -27,8 +27,9 @@ const JOINT_ARCH: JointArch = JointArch {
 };
 
 thread_local! {
-    // burn's NdArray model is `!Sync`, so the runtime is a per-thread `Rc`
-    // rather than a process-wide `Arc`. Inference is memoized inside the runtime.
+    // The runtime is a per-thread `Rc` (not a process-wide `Arc`): burn CPU
+    // models are not guaranteed `Sync`, and a per-thread instance also keeps the
+    // memoization cache lock-free across threads. Inference is memoized inside.
     static JOINT: Rc<JointRuntime> = Rc::new(build_joint());
 }
 

harper-brill/finished_joint/model.mpk

@@ -14,18 +14,31 @@ strum_macros = "0.28.0"
 serde = { version = "1.0.228", features = ["derive"] }
 is-macro = "0.3.7"
 rand = { version = "0.10.1", optional = true }
-burn = { version = "0.19.1", default-features = false, features = ["std"] }
-burn-ndarray = { version = "0.19.0", default-features = false }
+burn = { version = "0.21.0", default-features = false, features = ["std"] }
+burn-ndarray = { version = "0.21.0", default-features = false }
 serde_json = "1.0.150"
 smallvec = "1.15.1"
 lru = "0.18.0"
 rayon = { version = "1.12.0", optional = true }
 
-# burn-ndarray's explicit SIMD kernels (~1.8x faster cold inference) require the
-# `simd` feature, which isn't no_std-clean in 0.19.1 and so also pulls in `std`.
-# Gate both to native targets; the wasm32 build keeps the no_std scalar path.
+# The joint tagger infers on burn-flex (0.21+) on every target: a fast pure-Rust
+# CPU backend (gemm + SIMD) — ~5.6x faster cold inference than burn-ndarray, at
+# byte-identical output. Native uses the default features (incl. rayon).
+# burn-ndarray remains only for the legacy `BurnChunker` (burn_chunker.rs) and
+# tests; its explicit SIMD kernels need the `simd` feature (not no_std-clean, so
+# it also pulls `std`), hence the native-only entry below.
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
-burn-ndarray = { version = "0.19.0", default-features = false, features = ["simd", "std"] }
+burn-ndarray = { version = "0.21.0", default-features = false, features = ["simd", "std"] }
+burn-flex = { version = "0.21.0" }
+
+# wasm32 has no threads (drop `rayon`) and no atomic CAS (flex requires the
+# `critical-section` shim there). SIMD uses wasm's simd128.
+[target.'cfg(target_arch = "wasm32")'.dependencies]
+burn-flex = { version = "0.21.0", default-features = false, features = ["std", "simd", "critical-section"] }
+# burn-flex's RNG init pulls getrandom 0.4, which needs its `wasm_js` backend
+# feature enabled to build for wasm32-unknown-unknown (mirrors harper-wasm's own
+# getrandom 0.3 wasm_js setup; the two major versions coexist independently).
+getrandom = { version = "0.4", default-features = false, features = ["wasm_js"] }
 
 [features]
 default = []

harper-brill/finished_joint/suffix_vocab.json

@@ -3,6 +3,7 @@
 use crate::joint::{CHAR_PAD, CHAR_UNK};
 use hashbrown::HashMap;
 use serde::{Deserialize, Serialize};
+use std::collections::BTreeMap;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CharVocab {
@@ -53,7 +54,11 @@ impl CharVocab {
     }
 
     pub fn to_json(&self) -> String {
-        serde_json::to_string(&self.map).expect("serialize char vocab")
+        // Ordered (BTreeMap by char) + pretty-printed so the committed artifact
+        // has a stable, diff-friendly layout. The ids are the values; key order
+        // is cosmetic and does not affect decoding.
+        let ordered: BTreeMap<char, usize> = self.map.iter().map(|(&k, &v)| (k, v)).collect();
+        serde_json::to_string_pretty(&ordered).expect("serialize char vocab")
     }
 
     pub fn from_json(s: &str) -> Self {