diff --git a/JetlinkKit/Package.swift b/JetlinkKit/Package.swift index 0e01a44c..7fc7cf7f 100644 --- a/JetlinkKit/Package.swift +++ b/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")]), diff --git a/JetlinkKit/Sources/COrt/include/jl_ort.h b/JetlinkKit/Sources/COrt/include/jl_ort.h index 2ab08937..21075924 100644 --- a/JetlinkKit/Sources/COrt/include/jl_ort.h +++ b/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); diff --git a/JetlinkKit/Sources/COrt/jl_ort.c b/JetlinkKit/Sources/COrt/jl_ort.c index 77422489..c48b6e65 100644 --- a/JetlinkKit/Sources/COrt/jl_ort.c +++ b/JetlinkKit/Sources/COrt/jl_ort.c @@ -2,6 +2,7 @@ #include #include +#include #include @@ -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); diff --git a/JetlinkKit/Sources/JetlinkAndroid/Bridge.swift b/JetlinkKit/Sources/JetlinkAndroid/Bridge.swift index 29e537bb..3bf6f190 100644 --- a/JetlinkKit/Sources/JetlinkAndroid/Bridge.swift +++ b/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)") } diff --git a/JetlinkKit/Sources/JetlinkONNX/CoreMLPreparation.swift b/JetlinkKit/Sources/JetlinkONNX/CoreMLPreparation.swift index dc79ff7f..0bd9483e 100644 --- a/JetlinkKit/Sources/JetlinkONNX/CoreMLPreparation.swift +++ b/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,6 +130,7 @@ 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)] @@ -135,7 +138,7 @@ public enum CoreMLPreparation { 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 { diff --git a/JetlinkKit/Sources/JetlinkONNX/QnnPreparation.swift b/JetlinkKit/Sources/JetlinkONNX/QnnPreparation.swift new file mode 100644 index 00000000..a0ed707e --- /dev/null +++ b/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 + } + +} diff --git a/JetlinkKit/Sources/JetlinkORT/Ort.swift b/JetlinkKit/Sources/JetlinkORT/Ort.swift index 3a6830ad..1021f6d5 100644 --- a/JetlinkKit/Sources/JetlinkORT/Ort.swift +++ b/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?] = [] var dims: [UnsafeMutablePointer] = [] 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 } diff --git a/JetlinkKit/Sources/JetlinkORT/OrtBackend.swift b/JetlinkKit/Sources/JetlinkORT/OrtBackend.swift index a789aba3..ea34377a 100644 --- a/JetlinkKit/Sources/JetlinkORT/OrtBackend.swift +++ b/JetlinkKit/Sources/JetlinkORT/OrtBackend.swift @@ -25,6 +25,8 @@ public enum OrtProfile: String, CaseIterable, Sendable { /// fp16 whatever the graph says, so its fp32 heads are fp16 there too, and /// only the parity check can say whether that is close enough. case htpWhole = "htp-whole" + /// Same-frame causal work: HTP vision and GPU history, then HTP decode. + case htpParallel = "htp-parallel" /// The whole graph on the Adreno GPU. case gpu /// onnxruntime's CPU provider alone: for tests, the emulator and Linux, @@ -36,7 +38,7 @@ public enum OrtProfile: String, CaseIterable, Sendable { #if canImport(Metal) [.ane, .aneWhole, .coreml, .cpu] #elseif os(Android) - [.htp, .htpWhole, .gpu, .cpu] + [.htp, .htpWhole, .htpParallel, .gpu, .cpu] #else [.cpu] #endif @@ -50,6 +52,7 @@ public enum OrtProfile: String, CaseIterable, Sendable { case .aneWhole: [("model", .coreML("ALL"))] case .htp: [("vision", .htp), ("policy", .qnnGPU)] case .htpWhole: [("model", .htp)] + case .htpParallel: [("vision", .htp), ("history", .qnnGPU), ("decode", .htp), ("join", .cpu)] case .gpu: [("model", .qnnGPU)] case .cpu: [("model", .cpu)] } @@ -64,6 +67,7 @@ public enum OrtProfile: String, CaseIterable, Sendable { switch self { case .ane, .htp: .split case .aneWhole, .htpWhole: .aneWhole + case .htpParallel: .qnnPrefillBase case .coreml, .gpu: .whole case .cpu: #if os(Android) @@ -77,7 +81,7 @@ public enum OrtProfile: String, CaseIterable, Sendable { /// What a build writes: 5 is every graph split on `ane` and Expand as Tile /// on both, 6 the plain graph. An artifact prepared under another rebuilds. public var prepareVersion: Int { - layout == .plain ? 6 : 5 + self == .htpParallel ? 11 : (layout == .plain ? 6 : 5) } var usesCoreML: Bool { @@ -122,8 +126,8 @@ enum OrtUnit: Equatable { /// and their manifest (OrtArtifact), CoreML's compiled models or QNN's EP /// context beside them. /// -/// Nothing QNN has run on a Snapdragon yet: its options follow onnxruntime -/// 1.29's QNN documentation and source. +/// The causal profile uses a model-specific contract and generic QNN provider +/// discovery. A phone/runtime still needs accuracy and sustained qualification. public final class OrtBackend: EngineBackend { /// A first NPU compile with no earlier build to go by. The QNN graph /// finalization of a big model is minutes on a phone (unmeasured). @@ -138,25 +142,35 @@ 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 causalBundle: URL? 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(), causalBundle: URL? = nil + ) { self.profile = profile self.preparer = preparer self.keepAlive = keepAlive self.keepCPUWarm = keepCPUWarm self.chip = HostChip.resolve(chip) + self.qnn = profile == .htpParallel ? QnnPolicy( + sharedIO: qnn.sharedIO, strict: true, disableAdvancedActivationFusion: qnn.disableAdvancedActivationFusion, + libraryDirectory: qnn.libraryDirectory, runtimeIdentity: qnn.runtimeIdentity) : qnn + self.causalBundle = causalBundle } public var runtimeVersion: String { OrtRuntime.version } public func deviceTag() -> String { - sanitize("\(profile.rawValue)-\(chip)") + let contract = profile == .htpParallel ? "-\(CausalPrefillBundle.identity)" : "" + return sanitize("\(profile.rawValue)-\(chip)\(contract)\(profile.sessions.contains(where: { $0.unit == .htp }) ? qnn.cacheSuffix : "")") } public func deriveSpec(model: URL, sha256: String, nbytes: Int64, frameSkip: Int) throws -> ModelSpec { @@ -167,8 +181,7 @@ 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", @@ -176,9 +189,11 @@ public final class OrtBackend: EngineBackend { "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,8 +202,12 @@ 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 .qnnGPU: SessionPlan(model: model, provider: "QNN", options: providerOptions(unit), label: "QNN(gpu)", usesGPU: 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), + config: profile == .htpParallel ? qnn.sessionConfig : [:], label: "QNN(gpu)", usesGPU: true) case .cpu: SessionPlan(model: model, provider: nil, threads: HostChip.cpuThreads, label: "CPU") } } @@ -201,16 +220,59 @@ public final class OrtBackend: EngineBackend { } } + /// Choose before preparation, not after an arbitrary inference error. + func effectiveProfile(for sha256: String) -> OrtProfile { + profile == .htpParallel && sha256 != CausalPrefillBundle.sourceSHA ? .htpWhole : profile + } + + private func wholeFallback() -> OrtBackend { + OrtBackend(profile: .htpWhole, preparer: preparer, keepAlive: keepAlive, keepCPUWarm: keepCPUWarm, chip: chip, qnn: qnn) + } + + /// The configured cache remains isolated; the sidecar and engine report the + /// route actually used for this model. Unknown/old route metadata rebuilds. + func artifactProfile(_ meta: [String: Any]) throws -> OrtProfile { + guard let source = meta["causal_prefill_source"] as? String, CacheLayout.isSHA256(source), + let route = (meta["causal_prefill_route"] as? String).flatMap(OrtProfile.init(rawValue:)), + route == effectiveProfile(for: source) + else { throw ArtifactInvalid("causal-prefill artifact has no matching source and execution route") } + return route + } + + private func engine(_ artifact: URL, _ manifest: [[String: Any]], keepAlive: Bool, keepCPUWarm: Bool) throws -> OrtEngine { + try OrtEngine( + plans: plans(artifact, manifest), device: deviceTag(), keepAlive: keepAlive, keepCPUWarm: keepCPUWarm, + dataflow: profile == .htpParallel, + accelerator: profile == .htpParallel ? "NPU + GPU" : (profile == .htpWhole ? "NPU" : "")) + } + // MARK: build public func build(model: URL, artifact: URL, report: @escaping ProgressFn, metaExtra: [String: Any]) throws { + try qnn.validate() + var extra = metaExtra + if profile == .htpParallel { + let source = try Registry.hashFile(model).0 + let route = effectiveProfile(for: source) + extra["causal_prefill_source"] = source + extra["causal_prefill_route"] = route.rawValue + if route == .htpWhole { + log.info("this model has no causal split; preparing the existing whole-NPU route") + return try wholeFallback().build(model: model, artifact: artifact, report: report, metaExtra: extra) + } + } let started = Date() let expect = Artifact.sidecar(artifact) - try Artifact.build(artifact, metaExtra: metaExtra, report: report) { staged in + try Artifact.build(artifact, metaExtra: extra, report: report) { staged in try FileManager.default.createDirectory(at: staged, withIntermediateDirectories: true) report("patch", 0, profile.usesCoreML ? "preparing the model for CoreML" : "preparing the model") - let prepared = try preparer.prepare(model: model, into: staged, layout: profile.layout) { - CoreMLPreparation.cacheKey(stem: artifact.deletingPathExtension().lastPathComponent, part: $0) + let prepared: PreparedModel + if profile == .htpParallel { + prepared = try CausalPrefillBundle.prepare(model: model, into: staged, bundle: causalBundle) + } else { + prepared = try preparer.prepare(model: model, into: staged, layout: profile.layout) { + CoreMLPreparation.cacheKey(stem: artifact.deletingPathExtension().lastPathComponent, part: $0) + } } log.info("prepared \(model.lastPathComponent): \(prepared.summary)") let sessions = profile.sessions @@ -247,7 +309,7 @@ public final class OrtBackend: EngineBackend { } // Prove it runs before calling it built: minutes on a CPU, so it ticks. let providers = try Ticker.during(interval: 1, Ticker.paced("load", "loading the model to check it runs", took: 0, report: report)) { - let engine = try OrtEngine(plans: plans(staged, manifest), device: deviceTag(), keepAlive: false, keepCPUWarm: false) + let engine = try self.engine(staged, manifest, keepAlive: false, keepCPUWarm: false) defer { engine.close() } try engine.run() return engine.providers @@ -267,12 +329,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 +340,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( @@ -297,7 +364,7 @@ public final class OrtBackend: EngineBackend { report(stage, frac, msg) } let engine = try Ticker.during(interval: 2, tick) { - try OrtEngine(plans: plans(staged, manifest), device: deviceTag(), keepAlive: false, keepCPUWarm: false) + try self.engine(staged, manifest, keepAlive: false, keepCPUWarm: false) } defer { engine.close() } let (converted, compiled) = progress.measure() @@ -326,6 +393,15 @@ public final class OrtBackend: EngineBackend { // MARK: load public func load(artifact: URL, report: @escaping ProgressFn) throws -> any Engine { + try qnn.validate() + if profile == .htpParallel, try artifactProfile(Artifact.sidecar(artifact)) == .htpWhole { + let (manifest, _) = try OrtArtifact.open(artifact, version: OrtProfile.htpWhole.prepareVersion) + guard manifest.count == 1, OrtUnit(entry: manifest[0]) == .htp, + ["model.onnx", "model_ctx.onnx"].contains(manifest[0]["model"] as? String ?? "") else { + throw ArtifactInvalid("causal-prefill fallback artifact is not the whole-NPU route") + } + return try wholeFallback().load(artifact: artifact, report: report) + } 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" @@ -336,9 +412,19 @@ public final class OrtBackend: EngineBackend { let contents = (try? FileManager.default.contentsOfDirectory(atPath: artifact.appending(path: cache, directoryHint: .isDirectory).path)) ?? [] return contents.isEmpty ? "the CoreML cache for \(entry["model"] as? String ?? cache) is empty" : nil } + if profile == .htpParallel { + guard manifest.count == 4, + manifest.compactMap(OrtUnit.init(entry:)) == [.htp, .qnnGPU, .htp, .cpu], + zip(manifest, ["vision", "history", "decode", "join"]).allSatisfy({ entry, name in + let file = entry["model"] as? String + return file == "\(name).onnx" || ((name == "vision" || name == "decode") && file == "\(name)_ctx.onnx") + }) + else { throw ArtifactInvalid("causal-prefill artifact has a different session contract") } + try CausalPrefillBundle.verifyRetainedParts(in: artifact) + } let what = profile.usesCoreML ? "the CoreML model" : "the model" let (engine, seconds) = try Artifact.load(artifact, meta: meta, what: what, report: report) { - try OrtEngine(plans: plans(artifact, manifest), device: deviceTag(), keepAlive: keepAlive, keepCPUWarm: keepCPUWarm) + try self.engine(artifact, manifest, keepAlive: keepAlive, keepCPUWarm: keepCPUWarm) } log.info("onnxruntime sessions on \(profile.rawValue) in \(String(format: "%.1f", seconds)) s: \(engine.providers.joined(separator: " then "))") return engine diff --git a/JetlinkKit/Sources/JetlinkORT/OrtEngine.swift b/JetlinkKit/Sources/JetlinkORT/OrtEngine.swift index f12f9e18..42ad55b0 100644 --- a/JetlinkKit/Sources/JetlinkORT/OrtEngine.swift +++ b/JetlinkKit/Sources/JetlinkORT/OrtEngine.swift @@ -9,6 +9,8 @@ public struct SessionPlan: Sendable, Equatable { /// provider alone. public let provider: String? public let options: [String: String] + /// Per-session configuration, including an explicit fallback contract. + public let config: [String: String] /// The CPU provider's intra-op pool: 1 where an accelerator does the work. public let threads: Int /// The session as the log and the hello name it: "CoreML(CPUAndGPU)", @@ -20,12 +22,13 @@ public struct SessionPlan: Sendable, Equatable { public let usesNeuralEngine: Bool public init( - model: URL, provider: String?, options: [String: String] = [:], threads: Int = 1, label: String, + model: URL, provider: String?, options: [String: String] = [:], config: [String: String] = [:], threads: Int = 1, label: String, usesGPU: Bool = false, usesNeuralEngine: Bool = false ) { self.model = model self.provider = provider self.options = options + self.config = config self.threads = threads self.label = label self.usesGPU = usesGPU @@ -56,7 +59,7 @@ public struct SessionPlan: Sendable, Equatable { } } -/// A loaded model: a chain of onnxruntime sessions run back to back over +/// A loaded model: a chain or dependency DAG of onnxruntime sessions over /// buffers the engine owns, the Swift form of the Python ort worker. /// /// The host writes a frame's inputs into `hostInput(_:)`, calls `run()`, and @@ -70,26 +73,50 @@ public final class OrtEngine: EngineCore, @unchecked Sendable { /// Whether the CPU is kept warm beside this engine: a busy thread, or /// Android's performance hints. public var keepsCPUWarm: Bool { keepWarm != nil || hint != nil } - public override var notes: String { "CPU keep-warm \(keepsCPUWarm ? "on" : "off")" } + public override var notes: String { + "CPU keep-warm \(keepsCPUWarm ? "on" : "off"); QNN shared I/O \(sharedAllocator == nil ? "off" : "on")" + + (schedule == nil ? "" : "; session DAG \(parallelExecution ? "parallel" : "serial")") + } + public override var accelerator: String { executedAccelerator } public override var coolsWhenIdle: Bool { usesNeuralEngine } private let chain: [OrtSession] + private let sharedAllocator: OrtSharedAllocator? + private let sharedSessions: [Bool] private var bindings: [[OrtBinding]] = [] // [parity][session] private let keepAlive: MetalKeepAlive? private let keepWarm: CPUKeepWarm? private let hint: PerformanceHint? private let usesNeuralEngine: Bool + private let schedule: SessionDAG? + private let parallelExecution: Bool + private let executedAccelerator: String /// `keepAlive` keeps the GPU clocked up between frames, `keepCPUWarm` the /// CPU; each only where a plan runs on that unit. - public init(plans: [SessionPlan], device: String, keepAlive: Bool = true, keepCPUWarm: Bool = true) throws { + public init( + plans: [SessionPlan], device: String, keepAlive: Bool = true, keepCPUWarm: Bool = true, + dataflow: Bool = false, parallel: Bool = true, accelerator: String = "" + ) throws { self.device = device + self.executedAccelerator = accelerator var chain: [OrtSession] = [] for plan in plans { - chain.append(try OrtSession(model: plan.model, provider: plan.provider, options: plan.options, threads: plan.threads)) + chain.append(try OrtSession(model: plan.model, provider: plan.provider, options: plan.options, config: plan.config, threads: plan.threads)) } self.chain = chain self.providers = plans.map(\.label) + self.schedule = dataflow ? try SessionDAG(inputs: chain.map { $0.inputs.map(\.name) }, outputs: chain.map { $0.outputs.map(\.name) }) : nil + self.parallelExecution = parallel + let sharedSessions = plans.map { $0.options["enable_htp_shared_memory_allocator"] == "1" } + for (index, enabled) in sharedSessions.enumerated() where enabled { + guard plans[index].provider == "QNN", plans[index].usesNeuralEngine, !plans[index].usesGPU else { + throw OrtError("only QNN HTP sessions may bind QNN shared memory") + } + } + let shared = try sharedSessions.firstIndex(of: true).map { try OrtSharedAllocator(session: chain[$0]) } + self.sharedAllocator = shared + self.sharedSessions = sharedSessions // What the host stages: every session's inputs no earlier session // produces. What it can read: every session's outputs no later one reads. @@ -108,10 +135,25 @@ public final class OrtEngine: EngineCore, @unchecked Sendable { outputs[spec.name] = spec } } + if dataflow { + let written = Set(chain.flatMap { $0.outputs.map(\.name) }) + let read = Set(chain.flatMap { $0.inputs.map(\.name) }) + inputs = [:] + outputs = [:] + for session in chain { + for spec in session.inputs where !written.contains(spec.name) { inputs[spec.name] = spec } + for spec in session.outputs where !read.contains(spec.name) { outputs[spec.name] = spec } + } + } // One buffer per tensor name: inputs, hand-offs and outputs alike. var sizes: [String: Int] = [:] + var tensorTypes: [String: TensorSpec] = [:] for session in chain { for spec in session.inputs + session.outputs { + if dataflow, let previous = tensorTypes[spec.name], previous != spec { + throw OrtError("incompatible session tensor \(spec.name)") + } + tensorTypes[spec.name] = spec sizes[spec.name] = max(sizes[spec.name] ?? 0, spec.byteCount) } } @@ -122,7 +164,7 @@ public final class OrtEngine: EngineCore, @unchecked Sendable { let hint = warmCPU ? PerformanceHint.make() : nil self.hint = hint self.keepWarm = warmCPU && hint == nil ? CPUKeepWarm() : nil - try super.init(inputs: inputs, outputs: outputs, sizes: sizes) + try super.init(inputs: inputs, outputs: outputs, sizes: sizes, allocator: shared?.hostAllocator ?? .heap) // A throw from here on closes through deinit. try rebind([]) } @@ -145,14 +187,18 @@ public final class OrtEngine: EngineCore, @unchecked Sendable { var sets: [[OrtBinding]] = [] for p in 0..<(pairs.isEmpty ? 1 : 2) { var set: [OrtBinding] = [] - for session in chain { + for (index, session) in chain.enumerated() { let ins = session.inputs.map { ($0, buffer($0.name, parity: p)!) } let outs = session.outputs.map { spec -> (TensorSpec, UnsafeMutableRawPointer) in guard let input = outputOf[spec.name] else { return (spec, buffer(spec.name, parity: p)!) } // writes the buffer the state_ input reads next run return (spec, buffer(input, parity: p ^ 1)!) } - set.append(try OrtBinding(session: session, inputs: ins, outputs: outs)) + // These allocations remain host-accessible. Only HTP sessions label + // their bindings QnnHtpShared; CPU/GPU sessions use CPU memory metadata. + set.append(try OrtBinding( + session: session, inputs: ins, outputs: outs, + sharedAllocator: sharedSessions[index] ? sharedAllocator : nil)) } sets.append(set) } @@ -172,6 +218,12 @@ public final class OrtEngine: EngineCore, @unchecked Sendable { } public override func execute() throws { + if let schedule { + try schedule.run(parallel: parallelExecution) { index in + try self.bindings[self.parity][index].run() + } + return + } for binding in bindings[parity] { try binding.run() } @@ -181,7 +233,8 @@ public final class OrtEngine: EngineCore, @unchecked Sendable { /// second is the steady state. public override func warm() throws -> String { _ = try super.warm() - return "onnxruntime \(OrtRuntime.version) on \(device) in process, sessions \(providers.joined(separator: " then "))" + let shared = sharedAllocator.map { "; QNN shared I/O: \($0.totalBytes) bytes allocated" } ?? "" + return "onnxruntime \(OrtRuntime.version) on \(device) in process, sessions \(providers.joined(separator: " then "))\(shared)" } public override func close() { diff --git a/JetlinkKit/Sources/JetlinkORT/QnnPolicy.swift b/JetlinkKit/Sources/JetlinkORT/QnnPolicy.swift new file mode 100644 index 00000000..5917e0f6 --- /dev/null +++ b/JetlinkKit/Sources/JetlinkORT/QnnPolicy.swift @@ -0,0 +1,66 @@ +import CoreFoundation +import Foundation +import JetlinkServer + +/// Optional QNN HTP behavior. Defaults preserve the existing provider and caches. +/// Shared I/O changes host allocation; strict mode changes fallback admission; +/// disabling activation fusion is a separate numerical-correctness workaround. +public struct QnnPolicy: Sendable, Equatable { + public let sharedIO: Bool + public let strict: Bool + public let disableAdvancedActivationFusion: Bool + public let libraryDirectory: String? + public let runtimeIdentity: String? + + public init( + sharedIO: Bool = false, strict: Bool = false, disableAdvancedActivationFusion: Bool = false, + libraryDirectory: String? = nil, runtimeIdentity: String? = nil + ) { + self.sharedIO = sharedIO + self.strict = strict + self.disableAdvancedActivationFusion = disableAdvancedActivationFusion + self.libraryDirectory = libraryDirectory + self.runtimeIdentity = runtimeIdentity + } + + /// App configuration accepts JSON Booleans, not strings or numeric lookalikes. + public init(configuration: [String: Any]) throws { + func flag(_ key: String) throws -> Bool { + guard let value = configuration[key] else { return false } + guard let number = value as? NSNumber, CFGetTypeID(number) == CFBooleanGetTypeID() else { + throw HostError.invalid("\(key) must be a Boolean") + } + return number.boolValue + } + self.init( + sharedIO: try flag("qnn_shared_io"), strict: try flag("qnn_strict"), + disableAdvancedActivationFusion: try flag("qnn_no_activation_fusion"), + libraryDirectory: configuration["native_library_dir"] as? String, + runtimeIdentity: configuration["qnn_runtime_identity"] as? String) + try validate() + } + + var sessionConfig: [String: String] { strict ? ["session.disable_cpu_ep_fallback": "1"] : [:] } + + var cacheSuffix: String { + guard sharedIO || strict || disableAdvancedActivationFusion else { return "" } + let sdk = disableAdvancedActivationFusion ? "-sdk\(runtimeIdentity ?? "missing")" : "" + return "-qnn-v1-shared\(sharedIO ? 1 : 0)-strict\(strict ? 1 : 0)-nofusion\(disableAdvancedActivationFusion ? 1 : 0)\(sdk)" + } + + func validate() throws { + guard disableAdvancedActivationFusion else { return } + guard let libraryDirectory, libraryDirectory.hasPrefix("/"), !libraryDirectory.contains("\0"), + let runtimeIdentity, runtimeIdentity.count == 64, + runtimeIdentity.utf8.allSatisfy({ (48...57).contains($0) || (97...102).contains($0) }) else { + throw HostError.invalid("the no-fusion adapter requires a matching SDK directory and runtime identity") + } + } + + var htpBackend: [String: String] { + guard disableAdvancedActivationFusion else { return ["backend_type": "htp"] } + let directory = libraryDirectory.map { URL(fileURLWithPath: $0, isDirectory: true) } + let path = directory?.appendingPathComponent("libJetlinkQnnHtpNoFusion.so").path ?? "libJetlinkQnnHtpNoFusion.so" + return ["backend_path": path] + } +} diff --git a/JetlinkKit/Sources/JetlinkORT/SessionDAG.swift b/JetlinkKit/Sources/JetlinkORT/SessionDAG.swift new file mode 100644 index 00000000..cdf347af --- /dev/null +++ b/JetlinkKit/Sources/JetlinkORT/SessionDAG.swift @@ -0,0 +1,60 @@ +import Foundation + +/// A dataflow schedule over session I/O. Every layer is a completion barrier; +/// sessions in one layer read completed tensors and have disjoint outputs. +struct SessionDAG: Sendable { + let layers: [[Int]] + + init(inputs: [[String]], outputs: [[String]]) throws { + guard inputs.count == outputs.count, !inputs.isEmpty else { throw OrtError("empty or inconsistent session DAG") } + var producers: [String: Int] = [:] + for (index, names) in outputs.enumerated() { + for name in names { + guard producers.updateValue(index, forKey: name) == nil else { throw OrtError("multiple sessions write \(name)") } + } + } + let dependencies = inputs.enumerated().map { index, names in + Set(names.compactMap { producers[$0] }.filter { $0 != index }) + } + var waiting = Set(inputs.indices), completed = Set(), result: [[Int]] = [] + while !waiting.isEmpty { + let ready = waiting.filter { dependencies[$0].isSubset(of: completed) }.sorted() + guard !ready.isEmpty else { throw OrtError("session dependency cycle") } + result.append(ready) + completed.formUnion(ready) + waiting.subtract(ready) + } + layers = result + } + + func run(parallel: Bool, stage: @Sendable (Int) throws -> Void) throws { + for layer in layers { + if parallel && layer.count > 1 { + let failure = FirstFailure() + // Join every branch, including after failure, before state can advance + // or any shared tensor can be freed or reused. + DispatchQueue.concurrentPerform(iterations: layer.count) { offset in + do { try stage(layer[offset]) } catch { failure.store(error) } + } + if let error = failure.error { throw error } + } else { + for index in layer { try stage(index) } + } + } + } +} + +private final class FirstFailure: @unchecked Sendable { + private let lock = NSLock() + private var value: (any Error)? + func store(_ error: any Error) { + lock.lock() + defer { lock.unlock() } + if value == nil { value = error } + } + var error: (any Error)? { + lock.lock() + defer { lock.unlock() } + return value + } +} diff --git a/JetlinkKit/Sources/JetlinkServer/CausalPrefillBundle.swift b/JetlinkKit/Sources/JetlinkServer/CausalPrefillBundle.swift new file mode 100644 index 00000000..5cd876a6 --- /dev/null +++ b/JetlinkKit/Sources/JetlinkServer/CausalPrefillBundle.swift @@ -0,0 +1,76 @@ +import Foundation +import JetlinkRegistry + +/// A source-only, hash-bound split of one model's same-frame causal work. +/// Another architecture needs its own independently validated contract. +package enum CausalPrefillBundle { + package static let identity = "causal-prefill-v1-5a159599" + package static let sourceSHA = "404a18cfd86d29637d20c697dfde245bb47c666ae016730ab674c65f4d1e1aa4" + private static let manifestSHA = "5a159599f57ce154026baf967825cc4890e2e93ebca00adda1819b1c5354b7a3" + + private struct Manifest: Decodable { + let format: Int + let strategy: String + let sourceModelSHA256: String + let parts: [Part] + private enum CodingKeys: String, CodingKey { + case format, strategy, parts + case sourceModelSHA256 = "source_model_sha256" + } + } + private struct Part: Decodable { + let name: String + let file: String + let unit: String + let weightBytes: Int64 + let sha256: String + private enum CodingKeys: String, CodingKey { + case name, file, unit, sha256 + case weightBytes = "weight_bytes" + } + } + + package static func prepare(model: URL, into directory: URL, bundle: URL?) throws -> PreparedModel { + guard let bundle else { throw HostError.invalid("the causal-prefill model bundle is missing") } + let manifestURL = bundle.appendingPathComponent("manifest.json") + guard try Registry.hashFile(manifestURL).0 == manifestSHA else { + throw HostError.invalid("causal-prefill manifest does not match this engine") + } + let manifest = try JSONDecoder().decode(Manifest.self, from: Data(contentsOf: manifestURL)) + guard manifest.format == 1, manifest.strategy == "causal-prefill-v1", + manifest.sourceModelSHA256 == sourceSHA, + manifest.parts.map(\.name) == ["vision", "history", "decode", "join"], + manifest.parts.map(\.unit) == ["htp", "gpu", "htp", "cpu"], + try Registry.hashFile(model).0 == sourceSHA + else { throw HostError.invalid("causal-prefill requires the exact Cinque Terre V3 export") } + var parts: [PreparedModel.Part] = [] + for part in manifest.parts { + guard part.file == "\(part.name).onnx", part.weightBytes >= 0, CacheLayout.isSHA256(part.sha256) else { + throw HostError.invalid("invalid causal-prefill part identity") + } + let source = bundle.appendingPathComponent(part.file) + guard try Registry.hashFile(source).0 == part.sha256 else { + throw HostError.invalid("causal-prefill part \(part.name) failed its hash check") + } + try FileManager.default.copyItem(at: source, to: directory.appendingPathComponent(part.file)) + parts.append(.init(name: part.name, file: part.file, weightBytes: part.weightBytes)) + } + try FileManager.default.copyItem(at: manifestURL, to: directory.appendingPathComponent("causal-prefill.json")) + return PreparedModel(parts: parts, summary: "hash-verified causal split; original weights and same-frame feedback") + } + + /// HTP parts become device-compiled contexts. The GPU graph and CPU assembly + /// remain source ONNX and must still match the pinned contract at reload. + package static func verifyRetainedParts(in directory: URL) throws { + let file = directory.appendingPathComponent("causal-prefill.json") + guard (try? Registry.hashFile(file).0) == manifestSHA else { + throw ArtifactInvalid("causal-prefill retained manifest differs") + } + let manifest = try JSONDecoder().decode(Manifest.self, from: Data(contentsOf: file)) + for part in manifest.parts where part.unit != "htp" { + guard (try? Registry.hashFile(directory.appendingPathComponent(part.file)).0) == part.sha256 else { + throw ArtifactInvalid("causal-prefill retained part \(part.name) differs") + } + } + } +} diff --git a/JetlinkKit/Sources/jetlink-prefill-base/main.swift b/JetlinkKit/Sources/jetlink-prefill-base/main.swift new file mode 100644 index 00000000..4905544f --- /dev/null +++ b/JetlinkKit/Sources/jetlink-prefill-base/main.swift @@ -0,0 +1,14 @@ +import Foundation +import JetlinkONNX + +// Called by scripts/android/compile_causal_prefill.py after its source hash check. +guard CommandLine.arguments.count == 3 else { + fatalError("usage: jetlink-prefill-base SOURCE.onnx OUTPUT_DIRECTORY") +} +let source = URL(fileURLWithPath: CommandLine.arguments[1]) +let directory = URL(fileURLWithPath: CommandLine.arguments[2], isDirectory: true) +// Fixed legacy metadata is part of the byte-verified experimental bundle. +let report = try CoreMLPreparation.prepare( + source: source, into: directory, layout: .qnnPrefillBase, + cacheKey: { _ in "xiaomiNativeStable1model" }) +print("prepared \(report.parts.count) base graph, \(report.norms) fp16 normalizations") diff --git a/JetlinkKit/Tests/JetlinkONNXTests/QnnPreparationTests.swift b/JetlinkKit/Tests/JetlinkONNXTests/QnnPreparationTests.swift new file mode 100644 index 00000000..bf2706c5 --- /dev/null +++ b/JetlinkKit/Tests/JetlinkONNXTests/QnnPreparationTests.swift @@ -0,0 +1,41 @@ +import Testing + +@testable import JetlinkONNX + +@Suite struct QnnPreparationTests { + @Test("Prescaling retains a native normalization and its original learned parameters") + func nativeNormalization() throws { + var builder = OnnxGraphBuilder() + builder.input("x", DataType.float16, [1, 2]) + builder.fp16("scale", [2], [1, 2]) + builder.fp16("bias", [2], [0.25, -0.5]) + builder.node("LayerNormalization", ["x", "scale", "bias"], ["y"], [.float("epsilon", 1e-4)]) + builder.output("y", DataType.float16, [1, 2]) + let weights = builder.graph.initializers + var graph = builder.graph + #expect(try QnnPreparation.fusedLayerNorms(&graph) == 1) + #expect(graph.nodes.map(\.op) == ["Mul", "LayerNormalization"]) + #expect(graph.nodes[1].inputs.dropFirst() == ["scale", "bias"]) + #expect(graph.nodes[1].outputs == ["y"]) + #expect(graph.nodes[1].attribute("epsilon")?.f == Float(1e-4) / 16) + #expect(graph.initializers.prefix(weights.count).map(\.key) == weights.map(\.key)) + let raw: (Tensor) -> [UInt8]? = { tensor in + guard case .owned(let bytes)? = tensor.raw else { return nil } + return bytes + } + #expect(graph.initializers.prefix(weights.count).map(raw) == weights.map(raw)) + #expect(graph.initializers.last?.key == "__qnn_fused_norm_scale4") + // A second transformation must not silently apply another factor of four. + #expect(throws: OnnxError.self) { _ = try QnnPreparation.fusedLayerNorms(&graph) } + } + + @Test("Optional normalization statistics are rejected rather than changing their meaning") + func optionalStatistics() { + var builder = OnnxGraphBuilder() + builder.input("x", DataType.float16, [1, 2]) + builder.node("LayerNormalization", ["x", "scale"], ["y", "mean"]) + var graph = builder.graph + #expect(throws: OnnxError.self) { _ = try QnnPreparation.fusedLayerNorms(&graph) } + #expect(graph.nodes.map(\.op) == ["LayerNormalization"]) + } +} diff --git a/JetlinkKit/Tests/JetlinkServerTests/CausalPrefillBundleTests.swift b/JetlinkKit/Tests/JetlinkServerTests/CausalPrefillBundleTests.swift new file mode 100644 index 00000000..9b2ec7e2 --- /dev/null +++ b/JetlinkKit/Tests/JetlinkServerTests/CausalPrefillBundleTests.swift @@ -0,0 +1,24 @@ +import Foundation +import JetlinkTestSupport +import Testing + +@testable import JetlinkServer + +@Suite("Hash-bound causal bundle", .serialized) +struct CausalPrefillBundleTests { + @Test("Missing bundles and substituted sources fail before any part is copied") + func sourceIdentity() throws { + let temporary = try TemporaryDirectory() + let bundle = temporary.url.appending(path: "bundle") + let staged = temporary.url.appending(path: "staged") + try FileManager.default.createDirectory(at: bundle, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: staged, withIntermediateDirectories: true) + let source = Fixture.url("tiny_stateful.onnx") + #expect(throws: HostError.self) { _ = try CausalPrefillBundle.prepare(model: source, into: staged, bundle: nil) } + try FileManager.default.copyItem(at: Fixture.url("causal_prefill_manifest.json"), to: bundle.appending(path: "manifest.json")) + #expect(throws: HostError.self) { _ = try CausalPrefillBundle.prepare(model: source, into: staged, bundle: bundle) } + #expect(try FileManager.default.contentsOfDirectory(atPath: staged.path).isEmpty) + try Data("{}".utf8).write(to: bundle.appending(path: "manifest.json")) + #expect(throws: HostError.self) { _ = try CausalPrefillBundle.prepare(model: source, into: staged, bundle: bundle) } + } +} diff --git a/JetlinkKit/Tests/JetlinkServerTests/CausalSessionEngineTests.swift b/JetlinkKit/Tests/JetlinkServerTests/CausalSessionEngineTests.swift new file mode 100644 index 00000000..539956bd --- /dev/null +++ b/JetlinkKit/Tests/JetlinkServerTests/CausalSessionEngineTests.swift @@ -0,0 +1,88 @@ +import Foundation +import JetlinkONNX +import JetlinkRegistry +import JetlinkTestSupport +import Testing + +@testable import JetlinkORT +@testable import JetlinkServer + +@Suite("Same-frame recurrent session DAG", .serialized) +struct CausalSessionEngineTests { + @Test("Real ORT bindings follow graph dependencies and preserve recurrent state", arguments: [false, true]) + func statefulBindings(parallel: Bool) throws { + // Deliberately not topological: graph I/O determines the host interface + // and barriers, independently of the array's order. + let plans = ["join", "decode", "history", "vision"].map { + SessionPlan(model: Fixture.url("parallel_\($0).onnx"), provider: nil, label: "CPU") + } + let engine = try OrtEngine( + plans: plans, device: "test", keepAlive: false, keepCPUWarm: false, + dataflow: true, parallel: parallel) + defer { engine.close() } + #expect(Set(engine.inputs.keys) == ["x", "state_h"]) + #expect(Set(engine.outputs.keys) == ["outputs", "next_state_h"]) + try engine.loopState([("state_h", "next_state_h")]) + let x = try #require(engine.hostInput("x")) + let output = try #require(engine.output("outputs")) + x.storeBytes(of: Float(2), as: Float.self) + for frame in 0..<20 { + try engine.run() + #expect(output.load(as: Float.self) == Float(5 + frame)) + } + engine.resetState() + try engine.run() + #expect(output.load(as: Float.self) == 5) + } + + @Test("Generic provider discovery keeps strict accelerator admission without a chip whitelist") + func portablePolicy() { + for chip in ["SM8450", "SM8650", "SM8750", "future-qualcomm"] { + let backend = OrtBackend(profile: .htpParallel, preparer: ONNXPreparer(), chip: chip) + let model = URL(fileURLWithPath: "/unused.onnx") + #expect(backend.qnn.strict) + #expect(backend.deviceTag().contains(CausalPrefillBundle.identity)) + #expect(backend.deviceTag() != OrtBackend(profile: .htpWhole, preparer: ONNXPreparer(), chip: chip).deviceTag()) + #expect(backend.providerOptions(.htp)["backend_type"] == "htp") + #expect(backend.providerOptions(.htp)["soc_model"] == nil) + #expect(backend.providerOptions(.htp)["htp_arch"] == nil) + for unit in [OrtUnit.htp, .qnnGPU] { + #expect(backend.plan(unit, model: model, cache: nil).config["session.disable_cpu_ep_fallback"] == "1") + } + #expect(backend.contextConfiguration(context: model)["session.disable_cpu_ep_fallback"] == "1") + } + } + + @Test("Unknown models choose whole NPU before bundle preparation, not after an inference error") + func unknownModelFallback() throws { + let backend = OrtBackend(profile: .htpParallel, preparer: FallbackPreparer()) + #expect(backend.effectiveProfile(for: CausalPrefillBundle.sourceSHA) == .htpParallel) + #expect(backend.effectiveProfile(for: String(repeating: "a", count: 64)) == .htpWhole) + let temporary = try TemporaryDirectory() + #expect(throws: Intercept.whole) { + try backend.build( + model: Fixture.url("tiny_stateful.onnx"), artifact: temporary.url.appending(path: "model.ortcache"), + report: { _, _, _ in }, metaExtra: [:]) + } + } + + @Test("A cached fallback cannot be relabelled as parallel, and old route metadata rebuilds") + func cachedRoute() throws { + let backend = OrtBackend(profile: .htpParallel, preparer: ONNXPreparer()) + let unknown = String(repeating: "a", count: 64) + #expect(try backend.artifactProfile(["causal_prefill_source": unknown, "causal_prefill_route": "htp-whole"]) == .htpWhole) + #expect(try backend.artifactProfile(["causal_prefill_source": CausalPrefillBundle.sourceSHA, "causal_prefill_route": "htp-parallel"]) == .htpParallel) + #expect(throws: ArtifactInvalid.self) { + _ = try backend.artifactProfile(["causal_prefill_source": unknown, "causal_prefill_route": "htp-parallel"]) + } + #expect(throws: ArtifactInvalid.self) { _ = try backend.artifactProfile([:]) } + } +} + +private enum Intercept: Error, Equatable { case whole, unexpected } +private struct FallbackPreparer: ModelPreparer { + func readSpec(model: URL, sha256: String, nbytes: Int64, frameSkip: Int) throws -> ModelSpec { throw Intercept.unexpected } + func prepare(model: URL, into directory: URL, layout: CoreMLPreparation.Layout, cacheKey: @escaping (String) -> String) throws -> PreparedModel { + throw layout == .aneWhole ? Intercept.whole : Intercept.unexpected + } +} diff --git a/JetlinkKit/Tests/JetlinkServerTests/Fixtures/causal_prefill_manifest.json b/JetlinkKit/Tests/JetlinkServerTests/Fixtures/causal_prefill_manifest.json new file mode 100644 index 00000000..b77f21e3 --- /dev/null +++ b/JetlinkKit/Tests/JetlinkServerTests/Fixtures/causal_prefill_manifest.json @@ -0,0 +1,136 @@ +{ + "format": 1, + "strategy": "causal-prefill-v1", + "source_model_sha256": "404a18cfd86d29637d20c697dfde245bb47c666ae016730ab674c65f4d1e1aa4", + "prepared_model_sha256": "4ef3b6f248bf8d3ae0e1f76cfd0d5fcc4f079b6d910f735bd4cc22172b0d0166", + "parts": [ + { + "name": "vision", + "file": "vision.onnx", + "unit": "htp", + "nodes": 500, + "weight_bytes": 705804828, + "inputs": [ + "new_img", + "state_img_q" + ], + "outputs": [ + "view_39", + "linear_76", + "linear_77", + "linear_78", + "linear_79", + "linear_80", + "mul_36", + "linear_82", + "linear_81", + "view_60", + "next_state_img_q" + ], + "sha256": "49a4c1c1778884f87f26360d4256e6238f059a8f840c136c614ebabd548d59bc" + }, + { + "name": "history", + "file": "history.onnx", + "unit": "gpu", + "nodes": 327, + "weight_bytes": 40836830, + "inputs": [ + "desire", + "traffic_convention", + "action_t", + "state_desire_q", + "state_feat_q" + ], + "outputs": [ + "prefill_0_key", + "prefill_0_value", + "prefill_1_key", + "prefill_1_value", + "prefill_2_key", + "prefill_2_value", + "prefill_3_key", + "prefill_3_value", + "prefill_4_key", + "prefill_4_value", + "prefill_5_key", + "prefill_5_value", + "prefill_6_key", + "prefill_6_value", + "prefill_7_key", + "prefill_7_value" + ], + "sha256": "1370a34bbdb94027f1a99e833f09c639a904f7b716ede3de51cfe04aef73f359" + }, + { + "name": "decode", + "file": "decode.onnx", + "unit": "htp", + "nodes": 422, + "weight_bytes": 57434348, + "inputs": [ + "view_39", + "desire", + "traffic_convention", + "action_t", + "state_desire_q", + "prefill_0_key", + "prefill_0_value", + "prefill_1_key", + "prefill_1_value", + "prefill_2_key", + "prefill_2_value", + "prefill_3_key", + "prefill_3_value", + "prefill_4_key", + "prefill_4_value", + "prefill_5_key", + "prefill_5_value", + "prefill_6_key", + "prefill_6_value", + "prefill_7_key", + "prefill_7_value" + ], + "outputs": [ + "mul_41", + "mul_42", + "linear_120", + "linear_126", + "mul_48" + ], + "sha256": "ce475ff23e7740281ae1a894c8eb332c0c2486e5701e38c0f55d7b21e0e56e52" + }, + { + "name": "join", + "file": "join.onnx", + "unit": "cpu", + "nodes": 9, + "weight_bytes": 68, + "inputs": [ + "linear_76", + "linear_77", + "linear_78", + "linear_79", + "linear_80", + "mul_36", + "linear_82", + "linear_81", + "view_60", + "mul_41", + "mul_42", + "linear_120", + "linear_126", + "mul_48", + "state_feat_q", + "desire", + "state_desire_q" + ], + "outputs": [ + "outputs", + "next_state_feat_q", + "next_state_desire_q" + ], + "sha256": "a4e72d70d8580283d1ba21af3b458fc4eca3dcf4e0ab3f01eda41a761a78873f" + } + ] +} \ No newline at end of file diff --git a/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_decode.onnx b/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_decode.onnx new file mode 100644 index 00000000..e466c503 Binary files /dev/null and b/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_decode.onnx differ diff --git a/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_history.onnx b/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_history.onnx new file mode 100644 index 00000000..517f0081 Binary files /dev/null and b/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_history.onnx differ diff --git a/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_join.onnx b/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_join.onnx new file mode 100644 index 00000000..b331ca09 Binary files /dev/null and b/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_join.onnx differ diff --git a/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_vision.onnx b/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_vision.onnx new file mode 100644 index 00000000..653d04c6 Binary files /dev/null and b/JetlinkKit/Tests/JetlinkServerTests/Fixtures/parallel_vision.onnx differ diff --git a/JetlinkKit/Tests/JetlinkServerTests/OrtBackendTests.swift b/JetlinkKit/Tests/JetlinkServerTests/OrtBackendTests.swift index c5faa603..9eee1e0f 100644 --- a/JetlinkKit/Tests/JetlinkServerTests/OrtBackendTests.swift +++ b/JetlinkKit/Tests/JetlinkServerTests/OrtBackendTests.swift @@ -24,6 +24,7 @@ struct OrtBackendTests { (.aneWhole, ["model"], [.coreML("ALL")], "aneWhole"), (.htp, ["vision", "policy"], [.htp, .qnnGPU], "split"), (.htpWhole, ["model"], [.htp], "aneWhole"), + (.htpParallel, ["vision", "history", "decode", "join"], [.htp, .qnnGPU, .htp, .cpu], "qnnPrefillBase"), (.gpu, ["model"], [.qnnGPU], "whole"), (.cpu, ["model"], [.cpu], cpuLayout), ] @@ -37,7 +38,7 @@ struct OrtBackendTests { #if canImport(Metal) #expect(OrtProfile.available == [.ane, .aneWhole, .coreml, .cpu]) #elseif os(Android) - #expect(OrtProfile.available == [.htp, .htpWhole, .gpu, .cpu]) + #expect(OrtProfile.available == [.htp, .htpWhole, .htpParallel, .gpu, .cpu]) #else #expect(OrtProfile.available == [.cpu]) #endif diff --git a/JetlinkKit/Tests/JetlinkServerTests/OrtSharedMemoryTests.swift b/JetlinkKit/Tests/JetlinkServerTests/OrtSharedMemoryTests.swift new file mode 100644 index 00000000..bc015571 --- /dev/null +++ b/JetlinkKit/Tests/JetlinkServerTests/OrtSharedMemoryTests.swift @@ -0,0 +1,22 @@ +import Foundation +import Testing + +@testable import JetlinkORT + +@Suite("QNN shared memory", .serialized) +struct OrtSharedMemoryTests { + @Test("A CPU session cannot silently substitute heap memory for QNN RPC memory") + func cpuDoesNotPretendToShare() throws { + let session = try OrtSession(model: Golden("tiny_stateful").model, provider: nil) + #expect(throws: OrtError.self) { _ = try OrtSharedAllocator(session: session) } + } + + @Test("CPU plans cannot request HTP shared bindings") + func wrongProvider() throws { + let plan = SessionPlan(model: try Golden("tiny_stateful").model, provider: nil, + options: ["enable_htp_shared_memory_allocator": "1"], label: "CPU") + #expect(throws: OrtError.self) { + _ = try OrtEngine(plans: [plan], device: "cpu", keepAlive: false, keepCPUWarm: false) + } + } +} diff --git a/JetlinkKit/Tests/JetlinkServerTests/QnnPolicyTests.swift b/JetlinkKit/Tests/JetlinkServerTests/QnnPolicyTests.swift new file mode 100644 index 00000000..e5d16634 --- /dev/null +++ b/JetlinkKit/Tests/JetlinkServerTests/QnnPolicyTests.swift @@ -0,0 +1,98 @@ +import Foundation +import JetlinkKit +import JetlinkServer +import Testing + +@testable import JetlinkORT + +@Suite("QNN policy", .serialized) +struct QnnPolicyTests { + private let model = URL(fileURLWithPath: "/tmp/model.onnx") + private let runtime = String(repeating: "a", count: 64) + + @Test("Default policy preserves existing providers, session configs and cache tags") + func defaults() throws { + let policy = try QnnPolicy(configuration: [:]) + #expect(policy == QnnPolicy()) + for profile in OrtProfile.allCases where profile != .htpParallel { + let backend = OrtBackend(profile: profile, preparer: ONNXPreparer(), chip: "example", qnn: policy) + #expect(backend.deviceTag() == "\(profile.rawValue)-example") + #expect(backend.providerOptions(.htp)["backend_type"] == "htp") + #expect(backend.providerOptions(.htp)["backend_path"] == nil) + #expect(backend.providerOptions(.htp)["enable_htp_shared_memory_allocator"] == nil) + #expect(backend.plan(.htp, model: model, cache: nil).config.isEmpty) + } + } + + @Test("Strict configuration reaches the compilation and reload session plans") + func strictBothPhases() { + let backend = OrtBackend(profile: .htpWhole, preparer: ONNXPreparer(), qnn: QnnPolicy(strict: true)) + let key = "session.disable_cpu_ep_fallback" + let context = URL(fileURLWithPath: "/tmp/model_ctx.onnx") + #expect(backend.contextConfiguration(context: context)[key] == "1") + #expect(backend.contextConfiguration(context: context)["ep.context_enable"] == "1") + #expect(backend.contextConfiguration(context: context)["ep.context_file_path"] == context.path) + #expect(backend.plan(.htp, model: model, cache: nil).config == [key: "1"]) + #expect(backend.plan(.qnnGPU, model: model, cache: nil).config.isEmpty) + #expect(backend.plan(.cpu, model: model, cache: nil).config.isEmpty) + } + + @Test("Shared I/O belongs only to HTP, independent of the numerical workaround") + func sharedOptions() { + let backend = OrtBackend(profile: .htpWhole, preparer: ONNXPreparer(), qnn: QnnPolicy(sharedIO: true)) + #expect(backend.providerOptions(.htp)["enable_htp_shared_memory_allocator"] == "1") + #expect(backend.providerOptions(.htp)["backend_type"] == "htp") + #expect(backend.providerOptions(.qnnGPU) == ["backend_type": "gpu"]) + #expect(backend.providerOptions(.cpu).isEmpty) + #expect(backend.providerOptions(.coreML("ALL")).isEmpty) + #expect(backend.plan(.htp, model: model, cache: nil).config.isEmpty) + } + + @Test("The adapter replaces backend discovery and requires matching-runtime metadata") + func adapter() throws { + let policy = QnnPolicy(disableAdvancedActivationFusion: true, libraryDirectory: "/app/native", runtimeIdentity: runtime) + try policy.validate() + let backend = OrtBackend(profile: .htpWhole, preparer: ONNXPreparer(), qnn: policy) + #expect(backend.providerOptions(.htp)["backend_type"] == nil) + #expect(backend.providerOptions(.htp)["backend_path"] == "/app/native/libJetlinkQnnHtpNoFusion.so") + #expect(throws: HostError.self) { try QnnPolicy(disableAdvancedActivationFusion: true).validate() } + #expect(throws: HostError.self) { + try QnnPolicy(disableAdvancedActivationFusion: true, libraryDirectory: "relative", runtimeIdentity: runtime).validate() + } + #expect(throws: HostError.self) { + try QnnPolicy(disableAdvancedActivationFusion: true, libraryDirectory: "/native", runtimeIdentity: "unknown").validate() + } + } + + @Test("All option combinations and SDK bytes have distinct HTP cache namespaces") + func caches() { + var tags = Set() + for shared in [false, true] { + for strict in [false, true] { + for fusion in [false, true] { + let policy = QnnPolicy(sharedIO: shared, strict: strict, disableAdvancedActivationFusion: fusion, + libraryDirectory: "/native", runtimeIdentity: runtime) + let backend = OrtBackend(profile: .htpWhole, preparer: ONNXPreparer(), chip: "example", qnn: policy) + tags.insert(backend.deviceTag()) + #expect(OrtBackend(profile: .cpu, preparer: ONNXPreparer(), chip: "example", qnn: policy).deviceTag() == "cpu-example") + } + } + } + #expect(tags.count == 8) + let first = QnnPolicy(disableAdvancedActivationFusion: true, libraryDirectory: "/native", runtimeIdentity: runtime) + let second = QnnPolicy(disableAdvancedActivationFusion: true, libraryDirectory: "/native", runtimeIdentity: String(repeating: "b", count: 64)) + #expect(first.cacheSuffix != second.cacheSuffix) + } + + @Test("JSON accepts actual Booleans and rejects strings and numeric lookalikes") + func typedConfiguration() throws { + let object = try JSONSerialization.jsonObject(with: Data("{\"qnn_shared_io\":true,\"qnn_strict\":false}".utf8)) as! [String: Any] + #expect(try QnnPolicy(configuration: object).sharedIO) + let invalid: [Any] = ["true", "false", 1, 0, NSNull()] + for key in ["qnn_shared_io", "qnn_strict", "qnn_no_activation_fusion"] { + for value in invalid { + #expect(throws: HostError.self) { _ = try QnnPolicy(configuration: [key: value]) } + } + } + } +} diff --git a/JetlinkKit/Tests/JetlinkServerTests/SessionDAGTests.swift b/JetlinkKit/Tests/JetlinkServerTests/SessionDAGTests.swift new file mode 100644 index 00000000..b95081eb --- /dev/null +++ b/JetlinkKit/Tests/JetlinkServerTests/SessionDAGTests.swift @@ -0,0 +1,88 @@ +import Foundation +import Testing + +@testable import JetlinkORT + +@Suite("Fork/join session scheduling", .serialized) +struct SessionDAGTests { + private func fork() throws -> SessionDAG { + try SessionDAG( + inputs: [["image"], ["vision"], ["vision"], ["off", "on"]], + outputs: [["vision"], ["off"], ["on"], ["output"]]) + } + + @Test("Independent policies overlap, and assembly sees both complete") + func forkJoin() throws { + let dag = try fork() + #expect(dag.layers == [[0], [1, 2], [3]]) + let progress = Progress() + let offEntered = DispatchSemaphore(value: 0), onEntered = DispatchSemaphore(value: 0) + try dag.run(parallel: true) { stage in + switch stage { + case 0: progress.finish(0) + case 1: + #expect(progress.contains(0)) + offEntered.signal() + #expect(onEntered.wait(timeout: .now() + 2) == .success) + progress.finish(1) + case 2: + #expect(progress.contains(0)) + onEntered.signal() + #expect(offEntered.wait(timeout: .now() + 2) == .success) + progress.finish(2) + default: + #expect(progress.contains(1) && progress.contains(2)) + progress.finish(3) + } + } + #expect(progress.contains(3)) + } + + @Test("A failing branch is joined before the error escapes; assembly is skipped") + func failureJoins() throws { + let dag = try fork() + let progress = Progress() + let entered = DispatchSemaphore(value: 0) + #expect(throws: Failure.self) { + try dag.run(parallel: true) { stage in + if stage == 1 { + entered.signal() + throw Failure.expected + } + if stage == 2 { + #expect(entered.wait(timeout: .now() + 2) == .success) + Thread.sleep(forTimeInterval: 0.02) + progress.finish(2) + } + if stage == 3 { progress.finish(3) } + } + } + #expect(progress.contains(2)) + #expect(!progress.contains(3)) + } + + @Test("The same DAG can run serially for a controlled timing comparison") + func serialControl() throws { + let dag = try fork() + let progress = Progress() + try dag.run(parallel: false) { stage in + if stage > 0 { #expect(progress.contains(stage - 1)) } + progress.finish(stage) + } + #expect(progress.contains(3)) + } + + @Test("Cycles and competing writers are rejected before sessions run") + func rejectsInvalid() { + #expect(throws: (any Error).self) { try SessionDAG(inputs: [["b"], ["a"]], outputs: [["a"], ["b"]]) } + #expect(throws: (any Error).self) { try SessionDAG(inputs: [["x"], ["y"]], outputs: [["a"], ["a"]]) } + } +} + +private enum Failure: Error { case expected } +private final class Progress: @unchecked Sendable { + private let lock = NSLock() + private var completed = Set() + func finish(_ stage: Int) { lock.lock(); completed.insert(stage); lock.unlock() } + func contains(_ stage: Int) -> Bool { lock.lock(); defer { lock.unlock() }; return completed.contains(stage) } +} diff --git a/android/app/build.gradle.kts b/android/app/build.gradle.kts index 058b1c99..e53dc47c 100644 --- a/android/app/build.gradle.kts +++ b/android/app/build.gradle.kts @@ -22,6 +22,48 @@ val jetlinkVersionCode: Int = jetlinkVersion.split(".").map { it.takeWhile(Char::isDigit).toIntOrNull() ?: 0 } .let { (it.getOrElse(0) { 0 } * 10000) + (it.getOrElse(1) { 0 } * 100) + it.getOrElse(2) { 0 } } +// All options are build-time opt-ins; ordinary releases retain the Maven runtime. +fun qnnFlag(name: String): Boolean = providers.gradleProperty(name).orNull?.let { + it.toBooleanStrictOrNull() ?: error("$name must be true or false") +} ?: false +val qnnSharedIO = qnnFlag("jetlink.qnnSharedIO") +val causalPrefill = qnnFlag("jetlink.causalPrefill") +val qnnStrict = qnnFlag("jetlink.qnnStrict") +val qnnNoActivationFusion = qnnFlag("jetlink.qnnNoActivationFusion") +val qnnJniDir = providers.gradleProperty("jetlink.qnnJniDir").orNull +val qnnJniRoot = qnnJniDir?.let { file(it) } +val qnnRuntimeIdentity = if (qnnNoActivationFusion) { + check(qnnJniDir != null) { "no-fusion requires one matching SDK; see docs/android-qnn-npu.md" } + val directory = File(qnnJniRoot!!, "arm64-v8a") + val libraries = directory.listFiles()?.filter { it.isFile && it.extension == "so" }?.sortedBy { it.name }.orEmpty() + val names = libraries.map { it.name }.toSet() + check(names.containsAll(listOf("libJetlinkQnnHtpNoFusion.so", "libQnnHtp.so", "libQnnGpu.so", "libQnnSystem.so", + "libQnnHtpPrepare.so", "libQnnIr.so", "libQnnSaver.so"))) { "incomplete SDK jniLibs directory: $directory" } + val stub = names.filter { it.matches(Regex("libQnnHtpV\\d+Stub\\.so")) }.map { it.removeSuffix("Stub.so") }.toSet() + val skel = names.filter { it.matches(Regex("libQnnHtpV\\d+Skel\\.so")) }.map { it.removeSuffix("Skel.so") }.toSet() + check(stub.isNotEmpty() && stub == skel) { "matching HTP stub/skel pairs are required" } + // Actual runtime bytes form a cache namespace, not a preselected SDK/SoC pin. + val identity = MessageDigest.getInstance("SHA-256") + for (library in libraries) { + val digest = MessageDigest.getInstance("SHA-256") + library.inputStream().use { stream -> + val buffer = ByteArray(1024 * 1024) + while (true) { + val count = stream.read(buffer) + if (count < 0) break + digest.update(buffer, 0, count) + } + } + identity.update(library.name.toByteArray(Charsets.UTF_8)) + identity.update(0.toByte()) + identity.update(digest.digest()) + } + identity.digest().joinToString("") { "%02x".format(it) } +} else { + check(qnnJniDir == null) { "jetlink.qnnJniDir requires jetlink.qnnNoActivationFusion=true" } + "" +} + android { namespace = "io.zoompilot.jetlink" compileSdk = 37 @@ -36,6 +78,11 @@ android { targetSdk = 36 versionCode = jetlinkVersionCode versionName = jetlinkVersion + buildConfigField("boolean", "EXPERIMENTAL_CAUSAL_PREFILL", causalPrefill.toString()) + buildConfigField("boolean", "QNN_SHARED_IO", qnnSharedIO.toString()) + buildConfigField("boolean", "QNN_STRICT", qnnStrict.toString()) + buildConfigField("boolean", "QNN_NO_ACTIVATION_FUSION", qnnNoActivationFusion.toString()) + buildConfigField("String", "QNN_RUNTIME_IDENTITY", "\"$qnnRuntimeIdentity\"") // The QNN runtime is arm64 only, and so is every Snapdragon with an NPU // worth driving on. The emulator on an Apple silicon Mac is arm64 too. ndk { abiFilters += "arm64-v8a" } @@ -84,6 +131,8 @@ android { } } + if (qnnNoActivationFusion) sourceSets.getByName("main").jniLibs.srcDir(qnnJniRoot!!) + compileOptions { sourceCompatibility = JavaVersion.VERSION_17 targetCompatibility = JavaVersion.VERSION_17 @@ -132,7 +181,9 @@ dependencies { implementation(libs.kotlinx.coroutines.android) implementation(libs.kotlinx.serialization.json) // libonnxruntime.so with the QNN provider, and the QNN runtime's libraries - implementation(libs.onnxruntime.qnn) + implementation(libs.onnxruntime.qnn) { + if (qnnNoActivationFusion) exclude(group = "com.qualcomm.qti", module = "qnn-runtime") + } debugImplementation(libs.compose.ui.tooling) onnxruntimeAar(variantOf(libs.onnxruntime.qnn) { artifactType("aar") }) diff --git a/android/app/src/main/java/io/zoompilot/jetlink/server/ServerController.kt b/android/app/src/main/java/io/zoompilot/jetlink/server/ServerController.kt index c14ae1cf..1ca0d3f5 100644 --- a/android/app/src/main/java/io/zoompilot/jetlink/server/ServerController.kt +++ b/android/app/src/main/java/io/zoompilot/jetlink/server/ServerController.kt @@ -4,7 +4,9 @@ import android.content.Context import android.net.Uri import android.os.Build import android.util.Log +import io.zoompilot.jetlink.BuildConfig import io.zoompilot.jetlink.settings.Chip +import io.zoompilot.jetlink.settings.Processor import io.zoompilot.jetlink.settings.SettingsValues import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers @@ -63,6 +65,11 @@ class ServerController(private val context: Context, private val scope: Coroutin /** Starts the server with [settings], restarting one that runs. */ suspend fun start(settings: SettingsValues) = lifecycle.withLock { + if (settings.processor.resolved() == Processor.NpuGpuParallel && !BuildConfig.EXPERIMENTAL_CAUSAL_PREFILL) { + withContext(Dispatchers.IO) { Native.stop() } + run.value = RunState.Failed("This build does not enable the causal-prefill experiment.") + return@withLock + } run.value = RunState.Starting // the server stops one that runs first val error = withContext(Dispatchers.IO) { Native.start(config(settings).toString()) } @@ -99,6 +106,15 @@ class ServerController(private val context: Context, private val scope: Coroutin // a model compiled for the NPU is good for this system build alone put("firmware", Build.FINGERPRINT) put("native_library_dir", context.applicationInfo.nativeLibraryDir) + put("qnn_shared_io", BuildConfig.QNN_SHARED_IO) + put("qnn_strict", BuildConfig.QNN_STRICT) + put("qnn_no_activation_fusion", BuildConfig.QNN_NO_ACTIVATION_FUSION) + put("qnn_runtime_identity", BuildConfig.QNN_RUNTIME_IDENTITY) + if (processor == Processor.NpuGpuParallel) { + context.getExternalFilesDir(null)?.let { + put("parallel_bundle", File(it, "causal-prefill-v1").absolutePath) + } + } } /** Follows the snapshot and the log for as long as the app runs. */ diff --git a/android/app/src/main/java/io/zoompilot/jetlink/settings/Chip.kt b/android/app/src/main/java/io/zoompilot/jetlink/settings/Chip.kt index 1addc86c..2cc7319e 100644 --- a/android/app/src/main/java/io/zoompilot/jetlink/settings/Chip.kt +++ b/android/app/src/main/java/io/zoompilot/jetlink/settings/Chip.kt @@ -1,6 +1,7 @@ package io.zoompilot.jetlink.settings import android.os.Build +import io.zoompilot.jetlink.BuildConfig /** * The phone's SoC, and whether its NPU can run the model. QNN runs the @@ -28,7 +29,8 @@ object Chip { fun tensorGeneration(model: String): Int? = tensorModel.find(model.trim())?.groupValues?.get(1)?.toIntOrNull() /** What this phone's Settings offer. */ - fun processors(current: Processor): List = Processor.choices(isQualcomm, hasTensorNpu, isEmulator, current) + fun processors(current: Processor): List = Processor.choices( + isQualcomm, hasTensorNpu, isEmulator, current, parallel = BuildConfig.EXPERIMENTAL_CAUSAL_PREFILL) /** What Automatic runs on here. */ val automatic: Processor by lazy { Processor.automatic(hasTensorNpu, isEmulator) } diff --git a/android/app/src/main/java/io/zoompilot/jetlink/settings/Settings.kt b/android/app/src/main/java/io/zoompilot/jetlink/settings/Settings.kt index c4845926..3e38e0c1 100644 --- a/android/app/src/main/java/io/zoompilot/jetlink/settings/Settings.kt +++ b/android/app/src/main/java/io/zoompilot/jetlink/settings/Settings.kt @@ -44,6 +44,9 @@ enum class Processor(val backend: Backend, val device: String, val title: String */ NpuGpu(Backend.Ort, "htp", "NPU + GPU"), + /** Same-frame causal work in parallel; exact model contracts only, otherwise whole NPU. */ + NpuGpuParallel(Backend.Ort, "htp-parallel", "NPU + GPU (experimental)"), + /** The CPU: the emulator and tests. Seconds a frame with a real model. */ Cpu(Backend.Ort, "cpu", "CPU"); @@ -58,7 +61,11 @@ enum class Processor(val backend: Backend, val device: String, val title: String * NPU: the NPU took no model. [accelerator] is the ready engine's. */ fun fellBack(accelerator: String?, automatic: Processor = Chip.automatic): Boolean = - resolved(automatic) == TensorNpu && hardwareOf(accelerator).let { it != null && it != "NPU" } + when (resolved(automatic)) { + TensorNpu -> hardwareOf(accelerator).let { it != null && it != "NPU" } + NpuGpuParallel -> hardwareOf(accelerator) == "NPU" + else -> false + } /** What the Processor row says: the choice, and with Automatic what it runs on. */ fun summary(accelerator: String?, automatic: Processor = Chip.automatic): String { @@ -88,9 +95,9 @@ enum class Processor(val backend: Backend, val device: String, val title: String * CPU is offered on a Snapdragon only once chosen; the emulator offers * everything, for testing. */ - fun choices(qualcomm: Boolean, tensorNpu: Boolean, emulator: Boolean, current: Processor): List = when { - emulator -> entries - qualcomm -> entries.filter { it != TensorNpu && (it != Cpu || current == Cpu) } + fun choices(qualcomm: Boolean, tensorNpu: Boolean, emulator: Boolean, current: Processor, parallel: Boolean = false): List = when { + emulator -> entries.filter { it != NpuGpuParallel || parallel } + qualcomm -> entries.filter { it != TensorNpu && (it != Cpu || current == Cpu) && (it != NpuGpuParallel || parallel) } tensorNpu -> listOf(Auto, TensorNpu, Gpu, Cpu) else -> listOf(Auto, Gpu, Cpu) } diff --git a/android/app/src/main/java/io/zoompilot/jetlink/ui/settings/SettingsScreen.kt b/android/app/src/main/java/io/zoompilot/jetlink/ui/settings/SettingsScreen.kt index e92e7182..3470ca07 100644 --- a/android/app/src/main/java/io/zoompilot/jetlink/ui/settings/SettingsScreen.kt +++ b/android/app/src/main/java/io/zoompilot/jetlink/ui/settings/SettingsScreen.kt @@ -263,6 +263,10 @@ private fun Performance(values: SettingsValues, info: SettingsInfo, actions: Set var choosing by remember { mutableStateOf(false) } val choices = info.processors val footer = when { + values.processor.resolved(info.automatic) == Processor.NpuGpuParallel && + values.processor.fellBack(info.accelerator, info.automatic) -> + "This model has no validated causal split, so it runs on the whole NPU. " + + "Run Benchmark and check accuracy on this phone before use." values.processor.fellBack(info.accelerator, info.automatic) -> "The NPU did not take this model, so it runs on the GPU. Logs (in Help) say why. " + "Changing the processor prepares models again." diff --git a/android/app/src/test/java/io/zoompilot/jetlink/settings/ProcessorTest.kt b/android/app/src/test/java/io/zoompilot/jetlink/settings/ProcessorTest.kt index 78a2fa12..4480f01c 100644 --- a/android/app/src/test/java/io/zoompilot/jetlink/settings/ProcessorTest.kt +++ b/android/app/src/test/java/io/zoompilot/jetlink/settings/ProcessorTest.kt @@ -69,16 +69,17 @@ class ProcessorTest { @Test fun onlyOnnxruntimesNpuAndGpuChoicesNeedQnn() { - assertEquals(listOf(Processor.Npu, Processor.NpuGpu), Processor.entries.filter { it.usesQnn }) + assertEquals(listOf(Processor.Npu, Processor.NpuGpu, Processor.NpuGpuParallel), Processor.entries.filter { it.usesQnn }) assertEquals( - listOf(Backend.Ort, Backend.LiteRt, Backend.LiteRt, Backend.Ort, Backend.Ort), + listOf(Backend.Ort, Backend.LiteRt, Backend.LiteRt, Backend.Ort, Backend.Ort, Backend.Ort), Processor.entries.filter { it != Processor.Auto }.map { it.backend }, ) } @Test fun theEmulatorOffersEverything() { - assertEquals(Processor.entries, Processor.choices(qualcomm = false, tensorNpu = false, emulator = true, current = Processor.Cpu)) + assertEquals(Processor.entries.filter { it != Processor.NpuGpuParallel }, Processor.choices(false, false, true, Processor.Cpu)) + assertEquals(Processor.entries, Processor.choices(false, false, true, Processor.Cpu, parallel = true)) } @Test @@ -110,4 +111,15 @@ class ProcessorTest { assertEquals(null, Chip.tensorGeneration("SM8650")) assertEquals(null, Chip.tensorGeneration("Tensor")) } + + @Test + fun theParallelExperimentIsBuildOptInForQualcommAndReportsItsFallback() { + assertFalse(Processor.choices(true, false, false, Processor.Npu).contains(Processor.NpuGpuParallel)) + assertTrue(Processor.choices(true, false, false, Processor.Npu, parallel = true).contains(Processor.NpuGpuParallel)) + assertFalse(Processor.choices(false, true, false, Processor.TensorNpu, parallel = true).contains(Processor.NpuGpuParallel)) + assertEquals(Processor.NpuGpuParallel, Processor.of("htp-parallel")) + assertTrue(Processor.NpuGpuParallel.fellBack("NPU", automatic = Processor.Gpu)) + assertFalse(Processor.NpuGpuParallel.fellBack("NPU + GPU", automatic = Processor.Gpu)) + assertEquals("NPU + GPU (experimental) (NPU)", Processor.NpuGpuParallel.summary("NPU", automatic = Processor.Gpu)) + } } diff --git a/android/qnn/build.sh b/android/qnn/build.sh new file mode 100644 index 00000000..62aa1572 --- /dev/null +++ b/android/qnn/build.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# Use one locally obtained SDK for headers AND every packaged QNN library. +# No vendor code or model is downloaded or stored in this repository. +set -euo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +: "${QNN_SDK_ROOT:?set QNN_SDK_ROOT to the matching SDK}" +: "${ANDROID_NDK_HOME:?set ANDROID_NDK_HOME to the Android NDK}" +OUT=${1:?usage: android/qnn/build.sh NEW_OUTPUT_JNILIBS_DIRECTORY} +[[ ! -e "$OUT" ]] || { echo "Output must not exist: $OUT" >&2; exit 1; } +bash "$HERE/test.sh" +INCLUDE="$QNN_SDK_ROOT/include/QNN" +[[ -d "$INCLUDE" ]] || INCLUDE="$QNN_SDK_ROOT/include" +TOOLCHAIN="$ANDROID_NDK_HOME/toolchains/llvm/prebuilt/linux-x86_64/bin" +mkdir -p "$OUT/arm64-v8a" +"$TOOLCHAIN/aarch64-linux-android31-clang" -std=gnu11 -Wall -Wextra -Werror \ + -I "$INCLUDE" -fPIC -shared "$HERE/htp_no_activation_fusion.c" \ + -ldl -pthread -Wl,-z,relro,-z,now -o "$OUT/arm64-v8a/libJetlinkQnnHtpNoFusion.so" +python3 - "$QNN_SDK_ROOT" "$OUT/arm64-v8a" <<'PY' +import hashlib +from pathlib import Path +import re +import shutil +import sys +sdk, out = map(Path, sys.argv[1:]) +core = {'libQnnHtp.so', 'libQnnGpu.so', 'libQnnSystem.so', 'libQnnHtpPrepare.so', 'libQnnIr.so', 'libQnnSaver.so'} +found = {} +for path in sdk.rglob('libQnn*.so'): + name = path.name + skel = re.fullmatch(r'libQnnHtpV\d+Skel\.so', name) + stub = re.fullmatch(r'libQnnHtpV\d+Stub\.so', name) + if name not in core and not skel and not stub: + continue + if skel: + if 'unsigned' not in path.parts: + continue + elif 'aarch64-android' not in path.parts: + continue + with path.open('rb') as stream: + header = stream.read(20) + machine = 164 if skel else 183 + if header[:4] != b'\x7fELF' or header[5] != 1 or int.from_bytes(header[18:20], 'little') != machine: + continue + with path.open('rb') as stream: + digest = hashlib.file_digest(stream, 'sha256').hexdigest() + if name in found and found[name][0] != digest: + raise SystemExit(f'Ambiguous SDK library: {name}; choose one SDK root') + found[name] = (digest, path) +missing = core - found.keys() +stubs = {name.removesuffix('Stub.so') for name in found if name.endswith('Stub.so')} +skels = {name.removesuffix('Skel.so') for name in found if name.endswith('Skel.so')} +if missing or not stubs or stubs != skels: + raise SystemExit(f'Incomplete matching SDK: missing={sorted(missing)}, stubs={sorted(stubs)}, skels={sorted(skels)}') +for _, path in found.values(): + shutil.copyfile(path, out / path.name) +print(f'Staged adapter and {len(found)} libraries from one SDK; Gradle fingerprints their actual bytes for cache identity.') +PY diff --git a/android/qnn/htp_no_activation_fusion.c b/android/qnn/htp_no_activation_fusion.c new file mode 100644 index 00000000..fb0a5cfb --- /dev/null +++ b/android/qnn/htp_no_activation_fusion.c @@ -0,0 +1,124 @@ +// Public QNN API adapter; build against the same SDK as the packaged backend. +// This is a numerical workaround, independent of shared I/O and scheduling. +#define _GNU_SOURCE +#include +#include +#include +#include +#include +#include +#include +#include "QnnInterface.h" +#include "HTP/QnnHtpCommon.h" +#include "HTP/QnnHtpGraph.h" + +static pthread_once_t initialized = PTHREAD_ONCE_INIT; +static void *backend; +static QnnInterface_t wrapped; +static const QnnInterface_t *providers[] = {&wrapped}; +static QnnGraph_CreateFn_t real_create; +static QnnContext_FreeFn_t real_context_free; +static Qnn_ErrorHandle_t init_status = QNN_COMMON_ERROR_GENERAL; + +// Some backends consume graph configuration during finalization. Keep both +// the list and the custom configuration alive until contextFree succeeds. +typedef struct ConfigStorage { + struct ConfigStorage *next; + Qnn_ContextHandle_t context; + QnnHtpGraph_CustomConfig_t custom; + QnnGraph_Config_t config; + const QnnGraph_Config_t *all[]; +} ConfigStorage; +static ConfigStorage *retained; +static pthread_mutex_t retained_mutex = PTHREAD_MUTEX_INITIALIZER; + +static Qnn_ErrorHandle_t context_free(Qnn_ContextHandle_t context, Qnn_ProfileHandle_t profile) { + Qnn_ErrorHandle_t status = real_context_free(context, profile); + if (status != QNN_SUCCESS) return status; + pthread_mutex_lock(&retained_mutex); + ConfigStorage **link = &retained; + while (*link != NULL) { + ConfigStorage *entry = *link; + if (entry->context == context) { + *link = entry->next; + free(entry); + } else { + link = &entry->next; + } + } + pthread_mutex_unlock(&retained_mutex); + return status; +} + +static Qnn_ErrorHandle_t graph_create(Qnn_ContextHandle_t context, const char *name, + const QnnGraph_Config_t **original, Qnn_GraphHandle_t *graph) { + size_t count = 0; + if (original != NULL) { + while (count < 128 && original[count] != NULL) ++count; + if (count == 128) return QNN_COMMON_ERROR_INVALID_ARGUMENT; + } + ConfigStorage *storage = calloc(1, sizeof(*storage) + (count + 2) * sizeof(*storage->all)); + if (storage == NULL) return QNN_COMMON_ERROR_MEM_ALLOC; + storage->context = context; + storage->custom.option = QNN_HTP_GRAPH_CONFIG_OPTION_ADVANCED_ACTIVATION_FUSION; + storage->custom.advancedActivationFusion = false; + storage->config.option = QNN_GRAPH_CONFIG_OPTION_CUSTOM; + storage->config.customConfig = &storage->custom; + for (size_t i = 0; i < count; ++i) storage->all[i] = original[i]; + storage->all[count] = &storage->config; + Qnn_ErrorHandle_t status = real_create(context, name, storage->all, graph); + if (status == QNN_SUCCESS) { + pthread_mutex_lock(&retained_mutex); + storage->next = retained; + retained = storage; + pthread_mutex_unlock(&retained_mutex); + } else { + free(storage); + } + return status; +} + +static void initialize(void) { + Dl_info location; + if (dladdr((void *)&initialize, &location) == 0 || location.dli_fname == NULL) return; + const char *slash = strrchr(location.dli_fname, '/'); + if (slash == NULL) return; + char path[4096]; + if (snprintf(path, sizeof(path), "%.*s/libQnnHtp.so", (int)(slash - location.dli_fname), location.dli_fname) >= (int)sizeof(path)) return; + backend = dlopen(path, RTLD_NOW | RTLD_LOCAL); + if (backend == NULL) { + fprintf(stderr, "QNN no-fusion adapter: cannot open the sibling HTP backend: %s\n", dlerror()); + return; + } + typedef Qnn_ErrorHandle_t (*GetProviders)(const QnnInterface_t ***, uint32_t *); + GetProviders get = (GetProviders)dlsym(backend, "QnnInterface_getProviders"); + if (get == NULL) return; + const QnnInterface_t **list = NULL; + uint32_t count = 0; + if (get(&list, &count) != QNN_SUCCESS || list == NULL) return; + for (uint32_t i = 0; i < count; ++i) { + const QnnInterface_t *p = list[i]; + if (p == NULL || p->backendId != QNN_BACKEND_ID_HTP || + p->apiVersion.coreApiVersion.major != QNN_API_VERSION_MAJOR || + p->apiVersion.coreApiVersion.minor < QNN_API_VERSION_MINOR) continue; + wrapped = *p; + real_create = p->QNN_INTERFACE_VER_NAME.graphCreate; + real_context_free = p->QNN_INTERFACE_VER_NAME.contextFree; + if (real_create == NULL || real_context_free == NULL) return; + wrapped.QNN_INTERFACE_VER_NAME.graphCreate = graph_create; + wrapped.QNN_INTERFACE_VER_NAME.contextFree = context_free; + init_status = QNN_SUCCESS; + return; + } +} + +Qnn_ErrorHandle_t QnnInterface_getProviders(const QnnInterface_t ***out, uint32_t *count) { + if (out == NULL || count == NULL) return QNN_COMMON_ERROR_INVALID_ARGUMENT; + *out = NULL; + *count = 0; + pthread_once(&initialized, initialize); + if (init_status != QNN_SUCCESS) return init_status; + *out = providers; + *count = 1; + return QNN_SUCCESS; +} diff --git a/android/qnn/test.sh b/android/qnn/test.sh new file mode 100644 index 00000000..6beb3d37 --- /dev/null +++ b/android/qnn/test.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +# Host-only public-API forwarding/lifetime test; no accelerator execution. +set -euo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +: "${QNN_SDK_ROOT:?set QNN_SDK_ROOT to the matching SDK}" +INCLUDE="$QNN_SDK_ROOT/include/QNN" +[[ -d "$INCLUDE" ]] || INCLUDE="$QNN_SDK_ROOT/include" +BUILD=$(mktemp -d) +trap 'rm -rf "$BUILD"' EXIT +FLAGS=(-std=gnu11 -Wall -Wextra -Werror -I "$INCLUDE") +cc "${FLAGS[@]}" -fPIC -shared "$HERE/test_htp_no_activation_fusion.c" -DFAKE_BACKEND -o "$BUILD/libQnnHtp.so" +cc "${FLAGS[@]}" -fPIC -shared "$HERE/htp_no_activation_fusion.c" -ldl -pthread -o "$BUILD/libJetlinkQnnHtpNoFusion.so" +cc "${FLAGS[@]}" "$HERE/test_htp_no_activation_fusion.c" -ldl -o "$BUILD/test-no-fusion" +"$BUILD/test-no-fusion" "$BUILD/libJetlinkQnnHtpNoFusion.so" diff --git a/android/qnn/test_htp_no_activation_fusion.c b/android/qnn/test_htp_no_activation_fusion.c new file mode 100644 index 00000000..9a7cc672 --- /dev/null +++ b/android/qnn/test_htp_no_activation_fusion.c @@ -0,0 +1,79 @@ +#include +#include +#include +#include +#include +#include "QnnInterface.h" +#include "HTP/QnnHtpCommon.h" +#include "HTP/QnnHtpGraph.h" + +#ifdef FAKE_BACKEND +static const QnnGraph_Config_t **saved[4]; +static Qnn_ContextHandle_t contexts[4]; +static size_t saved_count; +static bool first_free = true; +static void check(const QnnGraph_Config_t **configs) { + assert(configs && configs[0] && configs[1] && !configs[2]); + assert(configs[0]->option == QNN_GRAPH_CONFIG_OPTION_PRIORITY); + assert(configs[1]->option == QNN_GRAPH_CONFIG_OPTION_CUSTOM); + const QnnHtpGraph_CustomConfig_t *c = configs[1]->customConfig; + assert(c->option == QNN_HTP_GRAPH_CONFIG_OPTION_ADVANCED_ACTIVATION_FUSION); + assert(!c->advancedActivationFusion); +} +static Qnn_ErrorHandle_t create(Qnn_ContextHandle_t context, const char *name, + const QnnGraph_Config_t **configs, Qnn_GraphHandle_t *graph) { + check(configs); + if (strcmp(name, "fail") == 0) return QNN_COMMON_ERROR_GENERAL; + assert(saved_count < 4); + saved[saved_count] = configs; + contexts[saved_count++] = context; + *graph = (void *)0x1; + return QNN_SUCCESS; +} +static Qnn_ErrorHandle_t free_context(Qnn_ContextHandle_t context, Qnn_ProfileHandle_t profile) { + (void)profile; + // Inspect configuration after graphCreate has returned, across contexts and + // after a failed contextFree. The wrapper may release only after success. + for (size_t i = 0; i < saved_count; ++i) if (saved[i]) check(saved[i]); + if (first_free) { first_free = false; return QNN_COMMON_ERROR_GENERAL; } + for (size_t i = 0; i < saved_count; ++i) if (contexts[i] == context) saved[i] = NULL; + return QNN_SUCCESS; +} +Qnn_ErrorHandle_t QnnInterface_getProviders(const QnnInterface_t ***list, uint32_t *count) { + static QnnInterface_t p = QNN_INTERFACE_INIT; + static const QnnInterface_t *ps[] = {&p}; + p.backendId = QNN_BACKEND_ID_HTP; + p.apiVersion.coreApiVersion.major = QNN_API_VERSION_MAJOR; + p.apiVersion.coreApiVersion.minor = QNN_API_VERSION_MINOR; + p.QNN_INTERFACE_VER_NAME.graphCreate = create; + p.QNN_INTERFACE_VER_NAME.contextFree = free_context; + *list = ps; *count = 1; + return QNN_SUCCESS; +} +#else +#include +int main(int argc, char **argv) { + assert(argc == 2); + void *lib = dlopen(argv[1], RTLD_NOW | RTLD_LOCAL); + if (!lib) { puts(dlerror()); return 1; } + typedef Qnn_ErrorHandle_t (*Get)(const QnnInterface_t ***, uint32_t *); + Get get = (Get)dlsym(lib, "QnnInterface_getProviders"); + assert(get && get(NULL, NULL) == QNN_COMMON_ERROR_INVALID_ARGUMENT); + const QnnInterface_t **p; uint32_t n; + assert(get(&p, &n) == QNN_SUCCESS && n == 1); + QnnGraph_Config_t priority = QNN_GRAPH_CONFIG_INIT; + priority.option = QNN_GRAPH_CONFIG_OPTION_PRIORITY; + const QnnGraph_Config_t *configs[] = {&priority, NULL}; + Qnn_GraphHandle_t graph = NULL; + Qnn_ContextHandle_t a = (void *)0x10, b = (void *)0x20; + assert(p[0]->QNN_INTERFACE_VER_NAME.graphCreate(a, "first", configs, &graph) == QNN_SUCCESS); + assert(p[0]->QNN_INTERFACE_VER_NAME.graphCreate(a, "second", configs, &graph) == QNN_SUCCESS); + assert(p[0]->QNN_INTERFACE_VER_NAME.graphCreate(b, "third", configs, &graph) == QNN_SUCCESS); + assert(p[0]->QNN_INTERFACE_VER_NAME.graphCreate(a, "fail", configs, &graph) == QNN_COMMON_ERROR_GENERAL); + assert(p[0]->QNN_INTERFACE_VER_NAME.contextFree(a, NULL) == QNN_COMMON_ERROR_GENERAL); + assert(p[0]->QNN_INTERFACE_VER_NAME.contextFree(a, NULL) == QNN_SUCCESS); + assert(p[0]->QNN_INTERFACE_VER_NAME.contextFree(b, NULL) == QNN_SUCCESS); + puts("configuration forwarding, graph failure and context lifetime passed"); + return 0; +} +#endif diff --git a/docs/android-causal-prefill.md b/docs/android-causal-prefill.md new file mode 100644 index 00000000..02839f06 --- /dev/null +++ b/docs/android-causal-prefill.md @@ -0,0 +1,110 @@ +# Experimental Android causal prefill + +This build-time opt-in overlaps Cinque Terre V3's current-image vision work on +HTP with its historical-token work on the QNN GPU. HTP decode waits for both, +then a small CPU graph assembles the original outputs and next queues. Current +queries still attend to all 288 original keys and values; no previous-frame +answer replaces current-frame computation. + +The scheduler and provider discovery are generic. The graph split is deliberately +limited to the exact CTV3 export below. Other models select the existing whole-NPU +route **before preparation**; the sidecar records that route and Settings reports +NPU. A missing or altered CTV3 bundle, unsupported operator/provider, allocation +failure or inference error fails explicitly. Such errors do not trigger a silent +CPU fallback or a different graph. + +There is no SM8750, manufacturer, HTP architecture or ORT-version whitelist. +The normal QNN backend discovers hardware using the runtime/DSP libraries the +build ships. This does not establish numerical correctness or speed on every +Snapdragon. Run original-model accuracy checks and sustained tests on each +device/runtime/firmware before relying on it. Existing Automatic and processor +defaults remain unchanged. + +## Generate the source-only bundle + +Use the normal Android/Swift build prerequisites and the original catalog ONNX: + +``` +SHA-256: 404a18cfd86d29637d20c697dfde245bb47c666ae016730ab674c65f4d1e1aa4 +Size: 766354845 bytes +``` + +```bash +python3 -m venv .venv-prefill +.venv-prefill/bin/python -m pip install -r scripts/android/requirements-causal-prefill.txt +swift build --package-path JetlinkKit --product jetlink-prefill-base +BASE_TOOL="$(swift build --package-path JetlinkKit --show-bin-path)/jetlink-prefill-base" +.venv-prefill/bin/python scripts/android/compile_causal_prefill.py \ + --model /path/to/original.onnx --base-tool "$BASE_TOOL" \ + --out /path/to/NEW-causal-prefill-v1 +``` + +The compiler checks the full source, prepared-base, four-part and manifest +identities, and rejects Python `-O`. It preserves original learned weights, +precision, output slices and recurrent contract. The legacy metadata cache-key +string in the base helper is retained solely to reproduce the historical graph +bytes; it is not a hardware selector. Model weights and compiled contexts are +not distributed in the source tree. + +## Build and stage + +```bash +cd android +./gradlew :app:assembleRelease :app:testDebugUnitTest \ + -Pjetlink.causalPrefill=true -Pjetlink.qnnSharedIO=true +``` + +Ordinary builds hide and reject the experimental choice. The experimental +profile always disables CPU EP fallback in its HTP and GPU sessions, including +HTP context compilation and reload. Shared I/O is optional: only actual HTP +allocators label tensors `QnnHtpShared`; GPU and CPU bindings use host metadata. +The engine retains allocations, sessions and bindings in the required order. + +The historical Xiaomi experiment required disabling advanced HTP activation +fusion to pass its original accuracy gates. To reproduce that numerical path, +build the optional adapter and **one complete matching SDK runtime** as described +in [Android QNN policy](android-qnn-npu.md), adding +`-Pjetlink.qnnNoActivationFusion=true -Pjetlink.qnnJniDir=/path/to/NEW-qnn-jni` +to the command above. No fixed SDK binary hashes or SoC settings are imposed. +Actual SDK library bytes identify its separate cache namespace. The generic +stock-runtime path still needs independent numerical qualification. + +Stage only generated `vision.onnx`, `history.onnx`, `decode.onnx`, `join.onnx` +and `manifest.json` into the app's external files directory, using the project's +normal signing/install workflow. For an unlocked development device: + +```bash +BUNDLE=/sdcard/Android/data/io.zoompilot.jetlink.android/files/causal-prefill-v1 +adb shell mkdir -p "$BUNDLE" +adb push /path/to/NEW-causal-prefill-v1/. "$BUNDLE/" +adb shell chmod 755 "$BUNDLE" +adb shell chmod 644 "$BUNDLE"/*.onnx "$BUNDLE/manifest.json" +``` + +Keep the original CTV3 model available in Jetlink and select **NPU + GPU +(experimental)**. The app verifies source/manifest/part hashes before copying; +reload verifies the exact session contract and retained GPU/assembly graphs. +Contexts compile locally for the device. Do not transfer another phone's +compiled contexts; rebuild and recheck accuracy after firmware/runtime changes. + +## Validation and evidence + +`SessionDAG` derives barriers from typed graph I/O and rejects cycles/multiple +writers. All branches join before an error escapes. Recurrent parity advances +only after a successful frame. Focused tests exercise reordered sessions, +serial/parallel recurrent equivalence, reset, worker failure, generic provider +options, model routing and invalid bundle/cache metadata. + +A historical Xiaomi 15 prototype measured **27.2 ms mean / 29.5 ms P99** across +12,001 synthetic frames at 20 Hz, with original 128-frame accuracy checks before +and after. Its exact four graphs measured **45.2 / 47.2 ms serial** versus +**26.8 / 29.0 ms parallel** in shorter sequential runs. These are historical +prototype measurements, not measurements of this portable port or another phone. +[Immutable prototype evidence](https://github.com/joe-sama/jetlink/blob/e681e507b5e83384f3316da1d25b4cce823e5ea9/docs/evidence/sm8750-causal-prefill/README.md) +retains the logs, screenshot, identities and thermal limitations. + +Use the existing `scripts/verify_parity.py` capture/reference/compare flow with +the original model, 128 identical inputs and unchanged gates, including recurrent +resets. Compare the same graphs serially and in parallel, then run a sustained +benchmark and cold reload. Hardware occupancy, a universal speedup and vehicle +readiness are not established by the host tests or provider-call concurrency. diff --git a/docs/android-qnn-npu.md b/docs/android-qnn-npu.md new file mode 100644 index 00000000..ae0e92c2 --- /dev/null +++ b/docs/android-qnn-npu.md @@ -0,0 +1,77 @@ +# Optional Android QNN HTP policy + +Ordinary builds retain the existing profiles, QNN runtime, provider options and +cache names. Three independent build-time options are available: + +| Gradle property | Effect | +| --- | --- | +| `jetlink.qnnSharedIO=true` | Allocate host tensors through ORT's actual QNN HTP shared allocator and bind them with the matching memory information. | +| `jetlink.qnnStrict=true` | Disable CPU execution-provider fallback in HTP sessions during both compilation and reload. Unsupported graphs fail instead of silently leaving compute on CPU. | +| `jetlink.qnnNoActivationFusion=true` | Use the optional matching-SDK adapter to disable advanced HTP activation fusion, a numerical workaround. | + +All default to false. The policy is typed; JSON configuration accepts Booleans, +not strings or numeric substitutes. GPU and CPU sessions do not register their +bindings as HTP shared memory. These policy options leave original models and +preparation unchanged; the optional causal profile is documented separately in +[android-causal-prefill.md](android-causal-prefill.md). + +For shared I/O and strict HTP admission using the normal runtime: + +```bash +cd android +./gradlew :app:assembleRelease :app:testDebugUnitTest \ + -Pjetlink.qnnSharedIO=true -Pjetlink.qnnStrict=true +``` + +The optional adapter requires locally obtained headers and libraries from one +matching Qualcomm SDK. Vendor headers, libraries and models are not included or +downloaded. Use the project's normal Android/Swift build prerequisites, then: + +```bash +export QNN_SDK_ROOT=/path/to/one/qairt-sdk +export ANDROID_NDK_HOME=/path/to/android-ndk +bash android/qnn/build.sh /path/to/NEW-qnn-jni +cd android +./gradlew :app:assembleRelease :app:testDebugUnitTest \ + -Pjetlink.qnnSharedIO=true -Pjetlink.qnnStrict=true \ + -Pjetlink.qnnNoActivationFusion=true -Pjetlink.qnnJniDir=/path/to/NEW-qnn-jni +``` + +The helper runs a host forwarding/lifetime test, builds the adapter and stages +Android libraries plus available matched HTP stub/skel pairs from that same SDK. +The SDK must expose `QNN_HTP_GRAPH_CONFIG_OPTION_ADVANCED_ACTIVATION_FUSION`. +The adapter forwards the public backend interface and disables only this option. +It retains custom configuration through successful `contextFree`, including when +the backend consumes configuration after graph creation. API incompatibility, +missing libraries, missing allocators or unsupported graphs produce errors. + +Gradle replaces the transitive Maven QNN runtime only in this explicit adapter +build, avoiding a mixture of runtimes. It fingerprints the actual staged library +bytes for the cache namespace; these are not fixed vendor-binary or SoC pins. +The runtime directory and identity flow through `QnnPolicy`, without global +tuning variables. Each changed policy has a different HTP cache namespace. + +## Evidence and limits + +On one Xiaomi 15 / Snapdragon 8 Elite, a historical whole-NPU CTV3 test using +ORT 1.29 and QAIRT 2.47 compared the same range-safe prepared graph with stock +advanced activation fusion versus fusion disabled. The stock capture failed +**six output slices** despite finite values; the disabled-fusion capture passed +all original gates across **32 identical input frames**. This establishes a +numerical workaround for that tested stack, not a universal cause or speedup. +The preparation used in that experiment is not added or changed here. + +A separate sequential one-minute comparison of the same whole-NPU engine +reported ordinary memory **42.9 ms mean / 45.2 ms P99**, versus shared I/O +**40.5 ms mean / 42.2 ms P99**. The mean reduction was about 5.6%. Order, +temperature and clocks were not fully controlled, so this is indicative evidence, +not a general performance guarantee. Strict fallback prevention is a correctness +contract rather than a speed optimization. None of these measurements establishes +50% improvement, 100% NPU occupancy or vehicle readiness. + +Before promoting an option on another graph/runtime/device, retain original +precision and use the unchanged original-model parity gates on identical inputs +with independent recurrent state, including reset tests. Compare alternating +heap/shared runs under similar temperatures, then perform sustained and cold-load +tests. Record exact runtime/options, package/native identities and mean/P99/max. +The current patch's host tests do not qualify an HTP device or a car connection. diff --git a/scripts/android/causal-prefill-cut-points.json b/scripts/android/causal-prefill-cut-points.json new file mode 100644 index 00000000..c16a3fba --- /dev/null +++ b/scripts/android/causal-prefill-cut-points.json @@ -0,0 +1,110 @@ +[ + { + "name": "vision", + "file": "vision.onnx", + "unit": "htp", + "after": [], + "inputs": [ + "new_img", + "state_img_q" + ], + "outputs": [ + "view_39", + "linear_76", + "linear_77", + "linear_78", + "linear_79", + "linear_80", + "mul_36", + "linear_82", + "linear_81", + "view_60", + "next_state_img_q" + ], + "nodes": 500, + "weight_bytes": 705804828, + "sha256": "56c0b305332e0de7ba2388c530287d3169c552a98ddd752520f4ef3766c1035f" + }, + { + "name": "off", + "file": "off.onnx", + "unit": "htp", + "after": [ + "vision" + ], + "inputs": [ + "view_39", + "desire", + "traffic_convention", + "action_t", + "state_desire_q", + "state_feat_q" + ], + "outputs": [ + "mul_41", + "mul_42", + "linear_120", + "linear_126", + "next_state_desire_q" + ], + "nodes": 230, + "weight_bytes": 32757604, + "sha256": "14304ea9fe4ca891ae7fb9a9c0e6f20c479837653944c4ff72777373adc52215" + }, + { + "name": "on", + "file": "on.onnx", + "unit": "gpu", + "after": [ + "vision" + ], + "inputs": [ + "view_39", + "desire", + "traffic_convention", + "action_t", + "state_desire_q", + "state_feat_q" + ], + "outputs": [ + "mul_48" + ], + "nodes": 208, + "weight_bytes": 25287078, + "sha256": "85d7c83ec14f8bafd7ae899ffb3f905840edf8c85de0b1a6ba1de41c315d903e" + }, + { + "name": "join", + "file": "join.onnx", + "unit": "cpu", + "after": [ + "vision", + "off", + "on" + ], + "inputs": [ + "linear_76", + "linear_77", + "linear_78", + "linear_79", + "linear_80", + "mul_36", + "linear_82", + "linear_81", + "view_60", + "mul_41", + "mul_42", + "linear_120", + "linear_126", + "mul_48", + "state_feat_q" + ], + "outputs": [ + "outputs", + "next_state_feat_q" + ], + "nodes": 6, + "weight_bytes": 44, + "sha256": "014a7415f2a15b6278e2217957ac8e5e65beeac55a2033fc123bfd3284cfb2e3" + } +] \ No newline at end of file diff --git a/scripts/android/causal-prefill-manifest.json b/scripts/android/causal-prefill-manifest.json new file mode 100644 index 00000000..b77f21e3 --- /dev/null +++ b/scripts/android/causal-prefill-manifest.json @@ -0,0 +1,136 @@ +{ + "format": 1, + "strategy": "causal-prefill-v1", + "source_model_sha256": "404a18cfd86d29637d20c697dfde245bb47c666ae016730ab674c65f4d1e1aa4", + "prepared_model_sha256": "4ef3b6f248bf8d3ae0e1f76cfd0d5fcc4f079b6d910f735bd4cc22172b0d0166", + "parts": [ + { + "name": "vision", + "file": "vision.onnx", + "unit": "htp", + "nodes": 500, + "weight_bytes": 705804828, + "inputs": [ + "new_img", + "state_img_q" + ], + "outputs": [ + "view_39", + "linear_76", + "linear_77", + "linear_78", + "linear_79", + "linear_80", + "mul_36", + "linear_82", + "linear_81", + "view_60", + "next_state_img_q" + ], + "sha256": "49a4c1c1778884f87f26360d4256e6238f059a8f840c136c614ebabd548d59bc" + }, + { + "name": "history", + "file": "history.onnx", + "unit": "gpu", + "nodes": 327, + "weight_bytes": 40836830, + "inputs": [ + "desire", + "traffic_convention", + "action_t", + "state_desire_q", + "state_feat_q" + ], + "outputs": [ + "prefill_0_key", + "prefill_0_value", + "prefill_1_key", + "prefill_1_value", + "prefill_2_key", + "prefill_2_value", + "prefill_3_key", + "prefill_3_value", + "prefill_4_key", + "prefill_4_value", + "prefill_5_key", + "prefill_5_value", + "prefill_6_key", + "prefill_6_value", + "prefill_7_key", + "prefill_7_value" + ], + "sha256": "1370a34bbdb94027f1a99e833f09c639a904f7b716ede3de51cfe04aef73f359" + }, + { + "name": "decode", + "file": "decode.onnx", + "unit": "htp", + "nodes": 422, + "weight_bytes": 57434348, + "inputs": [ + "view_39", + "desire", + "traffic_convention", + "action_t", + "state_desire_q", + "prefill_0_key", + "prefill_0_value", + "prefill_1_key", + "prefill_1_value", + "prefill_2_key", + "prefill_2_value", + "prefill_3_key", + "prefill_3_value", + "prefill_4_key", + "prefill_4_value", + "prefill_5_key", + "prefill_5_value", + "prefill_6_key", + "prefill_6_value", + "prefill_7_key", + "prefill_7_value" + ], + "outputs": [ + "mul_41", + "mul_42", + "linear_120", + "linear_126", + "mul_48" + ], + "sha256": "ce475ff23e7740281ae1a894c8eb332c0c2486e5701e38c0f55d7b21e0e56e52" + }, + { + "name": "join", + "file": "join.onnx", + "unit": "cpu", + "nodes": 9, + "weight_bytes": 68, + "inputs": [ + "linear_76", + "linear_77", + "linear_78", + "linear_79", + "linear_80", + "mul_36", + "linear_82", + "linear_81", + "view_60", + "mul_41", + "mul_42", + "linear_120", + "linear_126", + "mul_48", + "state_feat_q", + "desire", + "state_desire_q" + ], + "outputs": [ + "outputs", + "next_state_feat_q", + "next_state_desire_q" + ], + "sha256": "a4e72d70d8580283d1ba21af3b458fc4eca3dcf4e0ab3f01eda41a761a78873f" + } + ] +} \ No newline at end of file diff --git a/scripts/android/compile_causal_prefill.py b/scripts/android/compile_causal_prefill.py new file mode 100644 index 00000000..7ca5ac90 --- /dev/null +++ b/scripts/android/compile_causal_prefill.py @@ -0,0 +1,288 @@ +"""Exact causal past/current partition of the audited Cinque V3 fp16 graph. + +Vision and past-token computation are independent in the same frame. The decoder +waits for both and uses every original past/current attention key and value. +No model parameters, recurrent feedback, frame IDs, or accuracy gates change. +""" + +from pathlib import Path +import copy +import hashlib +import json +import numpy as np +import onnx +from onnx import helper as h, numpy_helper as nh +from onnx.utils import Extractor + +import argparse +import subprocess +import tempfile + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("--model", type=Path, required=True, help="original Cinque V3 model") +parser.add_argument("--out", type=Path, required=True) +parser.add_argument("--base-tool", type=Path, required=True, help="built jetlink-prefill-base executable") +args = parser.parse_args() +if not __debug__: + parser.error("Python -O disables the model-contract assertions; run without -O") + + +def sha256(path): + digest = hashlib.sha256() + with path.open('rb') as stream: + for block in iter(lambda: stream.read(8 * 1024 * 1024), b''): + digest.update(block) + return digest.hexdigest() + + +assert onnx.__version__ == "1.23.2", "reproduction is pinned to onnx==1.23.2" +assert np.__version__ == "2.5.4", "reproduction is pinned to numpy==2.5.4" +assert sha256(args.model) == "404a18cfd86d29637d20c697dfde245bb47c666ae016730ab674c65f4d1e1aa4", ( + "only the validated Cinque V3 export is supported" +) +OUT = args.out.resolve() +OUT.mkdir(parents=True, exist_ok=True) +assert not any(OUT.iterdir()), "choose an empty output directory" +temporary = tempfile.TemporaryDirectory(prefix="jetlink-prefill-base-") +source = Path(temporary.name) / "model.onnx" +subprocess.run([str(args.base_tool.resolve()), str(args.model.resolve()), temporary.name], check=True) +assert sha256(source) == "4ef3b6f248bf8d3ae0e1f76cfd0d5fcc4f079b6d910f735bd4cc22172b0d0166", ( + "base graph differs from the phone-validated graph" +) +m = onnx.load(source) +ex = Extractor(m) +plan = json.loads(Path(__file__).with_name("causal-prefill-cut-points.json").read_text()) +vision_terms = plan[0]['outputs'][1:-1] +policy_terms = [*plan[1]['outputs'][:-1], *plan[2]['outputs']] +original_inputs = [v.name for v in m.graph.input] +other_inputs = [v for v in original_inputs if v not in ['new_img', 'state_img_q']] +policy = ex.extract_model(['view_39', *other_inputs], policy_terms) +shapes = {v.name: copy.deepcopy(v) for v in [*policy.graph.input, *policy.graph.output, *policy.graph.value_info]} +producer = {s: n for n in policy.graph.node for s in n.output} +users = {} +for n in policy.graph.node: + for s in n.input: + users.setdefault(s, []).append(n) +attention = [] +for n in policy.graph.node: + if n.op_type != 'Softmax': + continue + where = producer[n.input[0]] + assert where.op_type == 'Where' + negative_mask = producer[where.input[0]] + assert negative_mask.op_type == 'Not' + score = producer[where.input[2]] + assert score.op_type == 'Mul' + qk = producer[score.input[0]] + assert qk.op_type == 'MatMul' + kt = producer[qk.input[1]] + assert kt.op_type == 'Transpose' + av = [x for x in users[n.output[0]] if x.op_type == 'MatMul'] + assert len(av) == 1 and av[0].input[0] == n.output[0] + attention.append({'mask': negative_mask.input[0], 'key': kt.input[0], 'value': av[0].input[1], 'key_transpose': kt.name, 'value_product': av[0].name}) +assert len(attention) == 8 + + +def prune(model): + g = model.graph + producers = {s: n for n in g.node for s in n.output} + needed = {v.name for v in g.output} + pending = list(needed) + nodes = set() + while pending: + s = pending.pop() + if s not in producers: + continue + node = producers[s] + if node.name in nodes: + continue + nodes.add(node.name) + for i in node.input: + if i and i not in needed: + needed.add(i) + pending.append(i) + kept = [n for n in g.node if n.name in nodes] + del g.node[:] + g.node.extend(kept) + for field in ['initializer', 'input']: + seq = getattr(g, field) + retained = [v for v in seq if v.name in needed] + del seq[:] + seq.extend(retained) + del g.value_info[:] + model = onnx.shape_inference.infer_shapes(model, strict_mode=True, data_prop=True) + onnx.checker.check_model(model, full_check=True) + return model + + +def make(role): + part = copy.deepcopy(policy) + g = part.graph + init = {v.name: v for v in g.initializer} + prod = {s: n for n in g.node for s in n.output} + length = 256 if role == 'history' else 32 + blocks = 8 if role == 'history' else 1 + + def constant(name, arr): + t = nh.from_array(np.asarray(arr), name) + if name in init: + init[name].CopyFrom(t) + else: + g.initializer.append(t) + init[name] = g.initializer[-1] + return name + + # The eight historical feature blocks are 24...31 of the 32-block + # historical window. The ninth block is the current view_39 image feature. + old_idx = nh.to_array(init['val_551']) + assert old_idx.shape == (9, 1) and np.array_equal(old_idx.reshape(-1), np.arange(-9, 0)) + if role == 'history': + cat = prod['cat_3'] + assert cat.op_type == 'Concat' and list(cat.input) == ['_to_copy_1', 'unsqueeze_1'] + cat.CopyFrom(h.make_node('Identity', ['_to_copy_1'], ['cat_3'], name='prefill_history_only')) + constant('val_551', (old_idx + 33)[:8].astype(np.int64)) + else: + index = prod['index'] + index.CopyFrom(h.make_node('Unsqueeze', ['view_39', 'val_3'], ['index'], name='prefill_current_only')) + # Desire windows and temporal positions are token-local embedding inputs. + desire_indices = nh.to_array(init['val_569']).reshape(9, 25) + constant('val_569', (desire_indices[:8] if role == 'history' else desire_indices[8:]).reshape(-1)) + # Temporal position broadcasts must select the same absolute positions. + temporal = ['unsqueeze_4', 'unsqueeze_8'] + for name in temporal: + if name in init: + a = nh.to_array(init[name]) + assert list(a.shape) == [9, 1, 512], (name, a.shape) + constant(name + '__' + role, a[:8] if role == 'history' else a[8:]) + nodes = [] + for n in g.node: + nodes.append(n) + for name in temporal: + if name not in n.output: + continue + output = name + '__' + role + start = 0 if role == 'history' else 8 + end = 8 if role == 'history' else 9 + args = [ + constant(output + '__start', np.array([start], np.int64)), + constant(output + '__end', np.array([end], np.int64)), + constant(output + '__axis', np.array([0], np.int64)), + ] + nodes.append(h.make_node('Slice', [name, *args], [output], name=output)) + # Retarget only consumers of the positions, leaving the Slice source intact. + for n in nodes: + if n.op_type == 'Slice' and n.name.endswith('__' + role): + continue + for i, s in enumerate(n.input): + if s in temporal: + n.input[i] = s + '__' + role + del g.node[:] + g.node.extend(nodes) + # Clone shape constants for each Reshape so no unrelated shape changes. + for n in g.node: + if n.op_type != 'Reshape' or n.input[1] not in init: + continue + a = nh.to_array(init[n.input[1]]) + if not np.any((a == 288) | (a == 9)): + continue + b = np.where(a == 288, length, np.where(a == 9, blocks, a)) + n.input[1] = constant(n.input[1] + '__' + role, b) + # Causal attention allows past queries to be evaluated without current data. + for info in attention: + mask = nh.to_array(init[info['mask']]) + assert mask.shape == (1, 1, 288, 288) and not mask[:, :, :256, 256:].any() + selected = mask[:, :, :256, :256] if role == 'history' else mask[:, :, 256:, :] + constant(info['mask'], selected) + past = [] + insert_before = {} + renamed = {} + for i, info in enumerate(attention): + for kind in ['key', 'value']: + current = info[kind] + cached = f'prefill_{i}_{kind}' + past.append(h.make_tensor_value_info(cached, onnx.TensorProto.FLOAT16, [1, 8, 256, 64])) + if role == 'history': + # A named graph output may still feed later layers. QNN GPU + # rejects output Identity aliases in that case; rename the + # producer and all consumers instead, with no arithmetic op. + renamed[current] = cached + else: + target = info['key_transpose'] if kind == 'key' else info['value_product'] + joined = f'{cached}__with_current' + insert_before.setdefault(target, []).append(h.make_node('Concat', [cached, current], [joined], name=joined, axis=2)) + node = next(n for n in g.node if n.name == target) + node.input[0 if kind == 'key' else 1] = joined + if role == 'history': + for n in g.node: + for field in [n.input, n.output]: + for at, value in enumerate(field): + if value in renamed: + field[at] = renamed[value] + del g.output[:] + g.output.extend(past) + else: + g.input.extend(past) + ordered = [] + for n in g.node: + ordered.extend(insert_before.get(n.name, [])) + ordered.append(n) + if n.op_type == 'Gather' and n.input[1] in init: + a = nh.to_array(init[n.input[1]]) + if a.size == 1 and int(a.reshape(-1)[0]) == 287: + n.input[1] = constant(n.input[1] + '__' + role, np.array(31, dtype=a.dtype)) + del g.node[:] + g.node.extend(ordered) + part = prune(part) + assert ('view_39' not in [v.name for v in part.graph.input]) == (role == 'history') + if role == 'decode': + assert 'state_feat_q' not in [v.name for v in part.graph.input] + return part + + +records = [] +models = [ + ('vision', ex.extract_model(['new_img', 'state_img_q'], plan[0]['outputs']), 'htp'), + ('history', make('history'), 'gpu'), + ('decode', make('decode'), 'htp'), + ( + 'join', + ex.extract_model([*vision_terms, *policy_terms, 'state_feat_q', 'desire', 'state_desire_q'], ['outputs', 'next_state_feat_q', 'next_state_desire_q']), + 'cpu', + ), +] +for name, part, unit in models: + part = prune(part) + path = OUT / (name + '.onnx') + onnx.save(part, path) + onnx.checker.check_model(str(path), full_check=True) + record = { + 'name': name, + 'file': path.name, + 'unit': unit, + 'nodes': len(part.graph.node), + 'weight_bytes': sum(len(t.raw_data) for t in part.graph.initializer), + 'inputs': [v.name for v in part.graph.input], + 'outputs': [v.name for v in part.graph.output], + 'sha256': sha256(path), + } + records.append(record) + print(json.dumps(record), flush=True) +(OUT / 'manifest.json').write_text( + json.dumps( + { + 'format': 1, + 'strategy': 'causal-prefill-v1', + 'source_model_sha256': '404a18cfd86d29637d20c697dfde245bb47c666ae016730ab674c65f4d1e1aa4', + 'prepared_model_sha256': '4ef3b6f248bf8d3ae0e1f76cfd0d5fcc4f079b6d910f735bd4cc22172b0d0166', + 'parts': records, + }, + indent=2, + ) +) +print('Saved exact same-frame (vision || history) -> decode -> join.', flush=True) + +assert (OUT / "manifest.json").read_bytes() == Path(__file__).with_name("causal-prefill-manifest.json").read_bytes(), ( + "generated bundle differs from the validated bundle" +) +print("All four generated graph hashes and the manifest exactly match the phone-tested bundle.") +temporary.cleanup() diff --git a/scripts/android/requirements-causal-prefill.txt b/scripts/android/requirements-causal-prefill.txt new file mode 100644 index 00000000..7a49b3a4 --- /dev/null +++ b/scripts/android/requirements-causal-prefill.txt @@ -0,0 +1,3 @@ +onnx==1.23.2 +numpy==2.5.4 +onnxruntime==1.29.0