Android: opt-in causal NPU/GPU execution with generic QNN discovery

JetlinkKit/Package.swift

@@ -48,6 +48,7 @@ let package = Package(
     .library(name: "JetlinkLiteRT", targets: ["JetlinkLiteRT"]),
     .library(name: "jetlink", type: .dynamic, targets: ["JetlinkAndroid"]),
     .executable(name: "jetlink-server", targets: ["jetlink-server"]),
+    .executable(name: "jetlink-prefill-base", targets: ["jetlink-prefill-base"]),
   ],
   dependencies: [
     .package(url: "https://github.com/apple/swift-crypto.git", "3.0.0"..<"5.0.0"),
@@ -125,6 +126,7 @@ let package = Package(
         .target(name: "JetlinkLinux", condition: .when(platforms: [.linux])),
         .product(name: "ArgumentParser", package: "swift-argument-parser", condition: .when(platforms: [.macOS, .linux])),
       ]),
+    .executableTarget(name: "jetlink-prefill-base", dependencies: ["JetlinkONNX"]),
     // Helpers more than one test target uses: fixtures, scratch directories, waits, a fake network.
     .target(name: "JetlinkTestSupport", dependencies: ["JetlinkRegistry", "JetlinkServer"], path: "Tests/JetlinkTestSupport"),
     .testTarget(name: "JetlinkKitTests", dependencies: ["JetlinkKit", "JetlinkUI", "JetlinkTestSupport"], resources: [.copy("Fixtures")]),

JetlinkKit/Sources/COrt/include/jl_ort.h

@@ -36,6 +36,7 @@ enum {
 typedef struct jl_env jl_env;
 typedef struct jl_session jl_session;
 typedef struct jl_binding jl_binding;
+typedef struct jl_shared_allocator jl_shared_allocator;
 
 typedef struct {
   const char *key;
@@ -77,10 +78,19 @@ char *jl_session_io_info(const jl_session *session, int output, size_t index,
                          char *name, size_t name_cap, int *elem_type,
                          int64_t *dims, size_t dims_cap, size_t *rank);
 
+// QNN's host-accessible RPC memory. The session must outlive this allocator,
+// and every allocation and binding must be freed before releasing it.
+char *jl_shared_allocator_create(jl_session *session, jl_shared_allocator **out);
+char *jl_shared_allocator_alloc(jl_shared_allocator *allocator, size_t bytes, void **out);
+void jl_shared_allocator_free(jl_shared_allocator *allocator, void *data);
+void jl_shared_allocator_release(jl_shared_allocator *allocator);
+size_t jl_shared_allocator_total_bytes(const jl_shared_allocator *allocator);
+
 // Binds caller-owned buffers to a session's inputs and outputs once, for
 // jl_binding_run to run with. The buffers must outlive the binding.
 char *jl_binding_create(jl_session *session, const jl_tensor *inputs, size_t n_inputs,
-                        const jl_tensor *outputs, size_t n_outputs, jl_binding **out);
+                        const jl_tensor *outputs, size_t n_outputs,
+                        const jl_shared_allocator *shared, jl_binding **out);
 char *jl_binding_run(jl_binding *binding);
 void jl_binding_release(jl_binding *binding);
 

JetlinkKit/Sources/COrt/jl_ort.c

@@ -2,6 +2,7 @@
 
 #include <stdlib.h>
 #include <string.h>
+#include <stdio.h>
 
 #include <onnxruntime/onnxruntime_c_api.h>
 
@@ -24,6 +25,13 @@ struct jl_binding {
   size_t n_values;
 };
 
+struct jl_shared_allocator {
+  OrtSession *session;
+  OrtAllocator *allocator;
+  OrtMemoryInfo *memory;
+  size_t total_bytes;
+};
+
 #ifdef JL_ORT_DLOPEN
 // NULL when the library or its entry point is missing; jl_env_create says why.
 static const OrtApiBase *base(void) {
@@ -274,14 +282,76 @@ void jl_binding_release(jl_binding *binding) {
   free(binding);
 }
 
+void jl_shared_allocator_release(jl_shared_allocator *allocator) {
+  if (allocator == NULL) return;
+  if (allocator->allocator != NULL) api()->ReleaseAllocator(allocator->allocator);
+  if (allocator->memory != NULL) api()->ReleaseMemoryInfo(allocator->memory);
+  free(allocator);
+}
+
+char *jl_shared_allocator_create(jl_session *session, jl_shared_allocator **out) {
+  if (session == NULL || out == NULL) return copy("shared allocator needs a session and output");
+  *out = NULL;
+  jl_shared_allocator *allocator = calloc(1, sizeof(*allocator));
+  if (allocator == NULL) return copy("out of memory");
+  allocator->session = session->session;
+  char *error = take(api()->CreateMemoryInfo("QnnHtpShared", OrtDeviceAllocator, 0,
+                                             OrtMemTypeDefault, &allocator->memory));
+  if (error == NULL) error = take(api()->CreateAllocator(session->session, allocator->memory, &allocator->allocator));
+  const OrtMemoryInfo *actual = NULL;
+  if (error == NULL) error = take(api()->AllocatorGetInfo(allocator->allocator, &actual));
+  int different = 1;
+  if (error == NULL) error = take(api()->CompareMemoryInfo(allocator->memory, actual, &different));
+  if (error == NULL && different != 0) error = copy("session did not return QNN shared memory");
+  if (error != NULL) {
+    jl_shared_allocator_release(allocator);
+    return error;
+  }
+  *out = allocator;
+  return NULL;
+}
+
+char *jl_shared_allocator_alloc(jl_shared_allocator *allocator, size_t bytes, void **out) {
+  if (allocator == NULL || out == NULL) return copy("invalid shared allocation");
+  *out = NULL;
+  if (bytes > SIZE_MAX - allocator->total_bytes) return copy("shared allocation byte count overflow");
+  char *error = take(api()->AllocatorAlloc(allocator->allocator, bytes, out));
+  if (error == NULL && *out == NULL) error = copy("QNN shared allocator returned no memory");
+  if (error == NULL) allocator->total_bytes += bytes;
+  return error;
+}
+
+void jl_shared_allocator_free(jl_shared_allocator *allocator, void *data) {
+  if (allocator == NULL || data == NULL) return;
+  char *error = take(api()->AllocatorFree(allocator->allocator, data));
+  if (error != NULL) {
+    fprintf(stderr, "Jetlink QNN shared memory free: %s\n", error);
+    free(error);
+  }
+}
+
+size_t jl_shared_allocator_total_bytes(const jl_shared_allocator *allocator) {
+  return allocator == NULL ? 0 : allocator->total_bytes;
+}
+
 static char *wrap(const OrtMemoryInfo *memory, const jl_tensor *t, OrtValue **out) {
   return take(api()->CreateTensorWithDataAsOrtValue(memory, t->data, t->nbytes, t->dims, t->rank,
                                                     (ONNXTensorElementDataType)t->elem_type, out));
 }
 
 char *jl_binding_create(jl_session *session, const jl_tensor *inputs, size_t n_inputs,
-                        const jl_tensor *outputs, size_t n_outputs, jl_binding **out) {
+                        const jl_tensor *outputs, size_t n_outputs,
+                        const jl_shared_allocator *shared, jl_binding **out) {
   *out = NULL;
+  if (shared != NULL && shared->session != session->session) {
+    // ORT's QNN allocator tracks RPC allocations globally. A second HTP
+    // context can register the same host-accessible allocation. Verify that
+    // this target session really exposes QnnHtpShared; CPU/GPU may not claim it.
+    jl_shared_allocator *compatible = NULL;
+    char *error = jl_shared_allocator_create(session, &compatible);
+    if (error != NULL) return error;
+    jl_shared_allocator_release(compatible);
+  }
   jl_binding *binding = calloc(1, sizeof(jl_binding));
   if (binding == NULL) {
     return copy("out of memory");
@@ -292,8 +362,13 @@ char *jl_binding_create(jl_session *session, const jl_tensor *inputs, size_t n_i
     free(binding);
     return copy("out of memory");
   }
-  OrtMemoryInfo *memory = NULL;
-  char *error = take(api()->CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &memory));
+  OrtMemoryInfo *cpu_memory = NULL;
+  const OrtMemoryInfo *memory = shared != NULL ? shared->memory : NULL;
+  char *error = NULL;
+  if (memory == NULL) {
+    error = take(api()->CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &cpu_memory));
+    memory = cpu_memory;
+  }
   if (error == NULL) {
     error = take(api()->CreateIoBinding(session->session, &binding->binding));
   }
@@ -311,8 +386,8 @@ char *jl_binding_create(jl_session *session, const jl_tensor *inputs, size_t n_i
       error = take(api()->BindOutput(binding->binding, outputs[i].name, *value));
     }
   }
-  if (memory != NULL) {
-    api()->ReleaseMemoryInfo(memory);
+  if (cpu_memory != NULL) {
+    api()->ReleaseMemoryInfo(cpu_memory);
   }
   if (error != NULL) {
     jl_binding_release(binding);

JetlinkKit/Sources/JetlinkAndroid/Bridge.swift

@@ -232,7 +232,9 @@
         }
         return OrtBackend(
           profile: profile, preparer: ONNXPreparer(), keepAlive: (config["keep_alive"] as? Bool) ?? true,
-          keepCPUWarm: (config["keep_cpu_warm"] as? Bool) ?? false, chip: chip)
+          keepCPUWarm: (config["keep_cpu_warm"] as? Bool) ?? false, chip: chip,
+            qnn: try QnnPolicy(configuration: config),
+            causalBundle: (config["parallel_bundle"] as? String).map { URL(fileURLWithPath: $0, isDirectory: true) })
       case let other:
         throw HostFailure("unknown backend \(other)")
       }

JetlinkKit/Sources/JetlinkONNX/CoreMLPreparation.swift

@@ -42,6 +42,8 @@ public enum CoreMLPreparation {
     /// (`--device ane-whole`): model.onnx, with the policy's norms prescaled
     /// and the vision heads in fp32.
     case aneWhole
+    /// Offline base for the hash-bound Cinque V3 causal-prefill compiler.
+    case qnnPrefillBase
     /// One model for TensorRT or the CPU: model.onnx, tinygrad's ops
     /// stripped and the images fp16, nothing else changed.
     case plain
@@ -119,7 +121,7 @@ public enum CoreMLPreparation {
     var tiles = 0
     if layout != .plain {
       gathers = try Patches.normalizeGatherIndices(&g, src)
-      gemms = try Patches.gemmWithTransposedWeight(&g, src)
+      if layout != .qnnPrefillBase { gemms = try Patches.gemmWithTransposedWeight(&g, src) }
       tiles = try Patches.expandToTile(&g, src)
     }
     var norms = 0
@@ -128,14 +130,15 @@ public enum CoreMLPreparation {
       norms = try Patches.prescaleLayerNorm(&g)
       heads = try Patches.headsInFP32(&g, src)
     }
+    if layout == .qnnPrefillBase { norms = try QnnPreparation.fusedLayerNorms(&g) }
     model.graph = g
 
     var parts: [(name: String, model: Model)]
     switch layout {
     case .split:
       let (vision, policy) = try Split.visionPolicy(model)
       parts = [("vision", vision), ("policy", policy)]
-    case .whole, .aneWhole, .plain:
+    case .whole, .aneWhole, .qnnPrefillBase, .plain:
       parts = [("model", model)]
     }
     for i in parts.indices where layout != .plain {

JetlinkKit/Sources/JetlinkONNX/QnnPreparation.swift

@@ -0,0 +1,51 @@
+import Foundation
+
+/// Range-safe fp16 preparation for the validated experimental QNN profile.
+public enum QnnPreparation {
+  /// Keep the native LayerNormalization kernel instead of expanding it into
+  /// seventeen elementwise/reduction operations. A power-of-two input scale
+  /// and epsilon/k^2 preserve the real-valued expression; device parity is
+  /// still required because the HTP's fp16 arithmetic differs from the CPU.
+  static func fusedLayerNorms(_ graph: inout Graph, k: Float = 4) throws -> Int {
+    let types = Patches.staticTypes(graph)
+    let dimensions = Patches.staticDims(graph)
+    let constant = "__qnn_fused_norm_scale4"
+    guard !graph.initializers.contains(where: { $0.key == constant }) else {
+      throw OnnxError("reserved QNN normalization constant already present")
+    }
+    var rewritten: [Node] = []
+    var count = 0
+    for (index, original) in graph.nodes.enumerated() {
+      var node = original
+      if node.op == "LayerNormalization", let input = node.inputs.first, types[input] == DataType.float16 {
+        guard node.outputs.dropFirst().allSatisfy({ $0.isEmpty }) else {
+          throw OnnxError("experimental fused normalization does not rewrite optional statistics outputs")
+        }
+        let scaled = "__qnn_fused_norm_\(index)"
+        rewritten.append(Node(inputs: [input, constant], outputs: [scaled], name: scaled, opType: "Mul"))
+        node.inputs[0] = scaled
+        let epsilon = (node.attribute("epsilon")?.f ?? 1e-5) / (k * k)
+        node.attributes.removeAll { $0.name == "epsilon" }
+        var encoded = Encoded()
+        encoded.stringField(1, "epsilon")
+        encoded.tag(2, .fixed32)
+        let bits = epsilon.bitPattern
+        for shift in stride(from: 0, to: 32, by: 8) { encoded.byte(UInt8(truncatingIfNeeded: bits >> shift)) }
+        encoded.intField(20, 1)
+        node.attributes.append(Attribute(bytes: .owned(encoded.tail), name: "epsilon", type: 1, i: 0, f: epsilon, ints: [], s: "", floats: [], t: nil))
+        if let shape = dimensions[input] { graph.valueInfo.append(.tensor(scaled, DataType.float16, shape)) }
+        count += 1
+      }
+      rewritten.append(node)
+    }
+    if count > 0 {
+      graph.initializers.append(
+        Tensor(
+          name: constant, dims: [], dataType: DataType.float16,
+          raw: .owned(Elements.encode([1 / k], as: DataType.float16))))
+    }
+    graph.nodes = rewritten
+    return count
+  }
+
+}

JetlinkKit/Sources/JetlinkORT/Ort.swift

@@ -113,11 +113,45 @@ final class OrtSession: @unchecked Sendable {
   }
 }
 
+/// Host-accessible allocations tracked by QNN for RPC memory registration.
+/// Retaining the session keeps its allocator valid until every buffer is freed.
+final class OrtSharedAllocator: @unchecked Sendable {
+  let pointer: OpaquePointer
+  private let session: OrtSession
+
+  init(session: OrtSession) throws {
+    self.session = session
+    var allocator: OpaquePointer?
+    try OrtError.check(jl_shared_allocator_create(session.pointer, &allocator))
+    guard let allocator else { throw OrtError("QNN returned no shared allocator") }
+    pointer = allocator
+  }
+
+  deinit { jl_shared_allocator_release(pointer) }
+
+  var totalBytes: Int { jl_shared_allocator_total_bytes(pointer) }
+
+  var hostAllocator: HostAllocator {
+    HostAllocator(
+      allocate: { [self] bytes in
+        var memory: UnsafeMutableRawPointer?
+        try OrtError.check(jl_shared_allocator_alloc(pointer, bytes, &memory))
+        guard let memory else { throw OrtError("QNN shared allocation returned no pointer") }
+        return memory
+      },
+      free: { [self] memory in
+        jl_shared_allocator_free(pointer, memory)
+      })
+  }
+}
+
 /// Caller-owned buffers bound to one session's inputs and outputs.
 final class OrtBinding {
   private let pointer: OpaquePointer
+  private let sharedAllocator: OrtSharedAllocator?
 
-  init(session: OrtSession, inputs: [(TensorSpec, UnsafeMutableRawPointer)], outputs: [(TensorSpec, UnsafeMutableRawPointer)]) throws {
+  init(session: OrtSession, inputs: [(TensorSpec, UnsafeMutableRawPointer)], outputs: [(TensorSpec, UnsafeMutableRawPointer)], sharedAllocator: OrtSharedAllocator? = nil) throws {
+    self.sharedAllocator = sharedAllocator
     var names: [UnsafeMutablePointer<CChar>?] = []
     var dims: [UnsafeMutablePointer<Int64>] = []
     defer {
@@ -135,7 +169,7 @@ final class OrtBinding {
     let ins = inputs.map { tensor($0.0, $0.1) }
     let outs = outputs.map { tensor($0.0, $0.1) }
     var binding: OpaquePointer?
-    try OrtError.check(jl_binding_create(session.pointer, ins, ins.count, outs, outs.count, &binding))
+    try OrtError.check(jl_binding_create(session.pointer, ins, ins.count, outs, outs.count, sharedAllocator?.pointer, &binding))
     guard let binding else { throw OrtError("onnxruntime returned no binding") }
     pointer = binding
   }