Android: opt-in QNN shared I/O and NPU correctness fixes

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,8 @@
         }
         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))
       case let other:
         throw HostFailure("unknown backend \(other)")
       }

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
   }

JetlinkKit/Sources/JetlinkORT/OrtBackend.swift

@@ -138,25 +138,30 @@ public final class OrtBackend: EngineBackend {
   public let keepAlive: Bool
   /// Keep a CPU core warm while the Neural Engine or the NPU runs a frame.
   public let keepCPUWarm: Bool
+  public let qnn: QnnPolicy
   private let preparer: any ModelPreparer
   private let chip: String
   private let log = ServerLog(category: "ort")
 
   /// `chip` names what the artifacts are valid for: on Android the SoC's
   /// model, "SM8650" (Build.SOC_MODEL), which a QNN context is compiled for.
   /// Nil is the Apple chip on Apple platforms and "cpu" elsewhere.
-  public init(profile: OrtProfile, preparer: any ModelPreparer, keepAlive: Bool = true, keepCPUWarm: Bool = true, chip: String? = nil) {
+  public init(
+    profile: OrtProfile, preparer: any ModelPreparer, keepAlive: Bool = true, keepCPUWarm: Bool = true,
+    chip: String? = nil, qnn: QnnPolicy = .init()
+  ) {
     self.profile = profile
     self.preparer = preparer
     self.keepAlive = keepAlive
     self.keepCPUWarm = keepCPUWarm
     self.chip = HostChip.resolve(chip)
+    self.qnn = qnn
   }
 
   public var runtimeVersion: String { OrtRuntime.version }
 
   public func deviceTag() -> String {
-    sanitize("\(profile.rawValue)-\(chip)")
+    sanitize("\(profile.rawValue)-\(chip)\(profile.sessions.contains(where: { $0.unit == .htp }) ? qnn.cacheSuffix : "")")
   }
 
   public func deriveSpec(model: URL, sha256: String, nbytes: Int64, frameSkip: Int) throws -> ModelSpec {
@@ -167,18 +172,19 @@ public final class OrtBackend: EngineBackend {
   func providerOptions(_ unit: OrtUnit) -> [String: String] {
     switch unit {
     case .htp:
-      [
-        "backend_type": "htp",
+      var options = qnn.htpBackend.merging([
         // burst holds the clocks up between frames, the counterpart of the
         // Mac's GPU keep-alive; sustained trades speed for less heat.
         "htp_performance_mode": keepAlive ? "burst" : "sustained_high_performance",
         // the slowest finalization and the fastest graph: it runs once, at build
         "htp_graph_finalization_optimization_mode": "3",
         "enable_htp_fp16_precision": "1",
         "qnn_context_priority": "high",
-      ]
-    case .qnnGPU: ["backend_type": "gpu"]
-    case .coreML, .cpu: [:]
+      ]) { _, value in value }
+      if qnn.sharedIO { options["enable_htp_shared_memory_allocator"] = "1" }
+      return options
+    case .qnnGPU: return ["backend_type": "gpu"]
+    case .coreML, .cpu: return [:]
     }
   }
 
@@ -187,7 +193,9 @@ public final class OrtBackend: EngineBackend {
   func plan(_ unit: OrtUnit, model: URL, cache: URL?) -> SessionPlan {
     switch unit {
     case .coreML(let units): SessionPlan(model: model, computeUnits: units, cacheDirectory: cache)
-    case .htp: SessionPlan(model: model, provider: "QNN", options: providerOptions(unit), label: "QNN(htp)", usesNeuralEngine: true)
+    case .htp: SessionPlan(
+      model: model, provider: "QNN", options: providerOptions(unit), config: qnn.sessionConfig,
+      label: "QNN(htp)", usesNeuralEngine: true)
     case .qnnGPU: SessionPlan(model: model, provider: "QNN", options: providerOptions(unit), label: "QNN(gpu)", usesGPU: true)
     case .cpu: SessionPlan(model: model, provider: nil, threads: HostChip.cpuThreads, label: "CPU")
     }
@@ -204,6 +212,7 @@ public final class OrtBackend: EngineBackend {
   // MARK: build
 
   public func build(model: URL, artifact: URL, report: @escaping ProgressFn, metaExtra: [String: Any]) throws {
+    try qnn.validate()
     let started = Date()
     let expect = Artifact.sidecar(artifact)
     try Artifact.build(artifact, metaExtra: metaExtra, report: report) { staged in
@@ -267,12 +276,7 @@ public final class OrtBackend: EngineBackend {
     let context = staged.appending(path: "\(session)_ctx.onnx")
     report("compile", 0, "compiling \(session) for the NPU")
     try Ticker.during(interval: 2, Ticker.paced("compile", "compiling \(session) for the NPU", took: took, report: report)) {
-      let config = [
-        "ep.context_enable": "1",
-        "ep.context_file_path": context.path,
-        // the context binary in a file of its own beside the model, not base64 inside it
-        "ep.context_embed_mode": "0",
-      ]
+      let config = contextConfiguration(context: context)
       _ = try OrtSession(model: source, provider: "QNN", options: providerOptions(.htp), config: config)
     }
     guard FileManager.default.fileExists(atPath: context.path) else {
@@ -283,6 +287,16 @@ public final class OrtBackend: EngineBackend {
     return context.lastPathComponent
   }
 
+  /// The same strict provider contract applies during compilation and reload.
+  func contextConfiguration(context: URL) -> [String: String] {
+    qnn.sessionConfig.merging([
+        "ep.context_enable": "1",
+        "ep.context_file_path": context.path,
+        // the context binary in a file of its own beside the model, not base64 inside it
+        "ep.context_embed_mode": "0",
+    ]) { _, value in value }
+  }
+
   /// A CoreML build's first session creation is the conversion and CoreML's
   /// compile, so its progress is read off the cache directories as they fill.
   private func checkCoreML(
@@ -326,6 +340,7 @@ public final class OrtBackend: EngineBackend {
   // MARK: load
 
   public func load(artifact: URL, report: @escaping ProgressFn) throws -> any Engine {
+    try qnn.validate()
     let (manifest, meta) = try OrtArtifact.open(artifact, version: profile.prepareVersion) { entry in
       guard let unit = OrtUnit(entry: entry) else { return "a session names no unit" }
       // Without its compile, onnxruntime would recompile under a "loading"