File size: 8,221 Bytes
606bc51 45847ba 606bc51 45847ba 606bc51 45847ba 606bc51 45847ba 606bc51 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | import CoreML
import Darwin
import Foundation
/// Load and prediction timing for a compiled Nemotron bundle on macOS or iOS.
/// Invoke one bundle and one compute-unit configuration per app launch.
@available(macOS 15.0, iOS 18.0, *)
public enum NemotronExportProbe {
public struct Report: Codable, Sendable {
public let bundle: String
public let os: String
public let machine: String
public let computeUnits: String
public let models: [ModelRecord]
}
public struct ModelRecord: Codable, Sendable {
public let component: String
public let firstObserved: Measurement
public let sameProcessReload: Measurement
public let plannedOperationsByPreferredDevice: [String: Int]
public let inputs: [FeatureSignature]
public let outputs: [FeatureSignature]
}
public struct Measurement: Codable, Sendable {
public let loadMs: Double
public let firstPredictionMs: Double
public let medianRepeatedPredictionMs: Double
}
public struct FeatureSignature: Codable, Sendable {
public let name: String
public let shape: [Int]
public let dtype: String
}
public static func run(bundleURL: URL, computeUnits: MLComputeUnits) async throws -> Report {
let configuration = MLModelConfiguration()
configuration.computeUnits = computeUnits
var records: [ModelRecord] = []
for name in ["encoder", "decoder", "joint"] {
let url = bundleURL.appendingPathComponent("\(name).mlmodelc", isDirectory: true)
print("MODEL_START \(name)")
let first = try measure(url: url, configuration: configuration)
print("MODEL_FIRST_COMPLETE \(name) loadMs=\(first.0.loadMs)")
let repeated = try measure(url: url, configuration: configuration)
print("MODEL_RELOAD_COMPLETE \(name) loadMs=\(repeated.0.loadMs)")
let plan = try await MLComputePlan.load(contentsOf: url, configuration: configuration)
print("MODEL_PLAN_COMPLETE \(name)")
var devices: [String: Int] = [:]
if case let .program(program) = plan.modelStructure,
let main = program.functions["main"] {
for operation in main.block.operations {
guard let usage = plan.deviceUsage(for: operation) else { continue }
let device: String
switch usage.preferred {
case .cpu: device = "cpu"
case .gpu: device = "gpu"
case .neuralEngine: device = "neuralEngine"
@unknown default: device = "unknown"
}
devices[device, default: 0] += 1
}
}
records.append(ModelRecord(
component: name, firstObserved: first.0, sameProcessReload: repeated.0,
plannedOperationsByPreferredDevice: devices, inputs: first.1, outputs: first.2))
}
var system = utsname()
uname(&system)
let machine = withUnsafePointer(to: &system.machine) {
String(cString: UnsafeRawPointer($0).assumingMemoryBound(to: CChar.self))
}
let units: String
switch computeUnits {
case .cpuOnly: units = "cpuOnly"
case .cpuAndGPU: units = "cpuAndGPU"
case .cpuAndNeuralEngine: units = "cpuAndNeuralEngine"
case .all: units = "all"
@unknown default: units = "unknown"
}
return Report(
bundle: bundleURL.lastPathComponent,
os: ProcessInfo.processInfo.operatingSystemVersionString,
machine: machine, computeUnits: units, models: records)
}
private static func signatures(_ descriptions: [String: MLFeatureDescription]) throws -> [FeatureSignature] {
try descriptions.keys.sorted().map { name in
guard let array = descriptions[name]?.multiArrayConstraint else {
throw ProbeError.unsupportedFeature(name)
}
return FeatureSignature(
name: name, shape: array.shape.map(\.intValue), dtype: String(describing: array.dataType))
}
}
private static func measure(
url: URL, configuration: MLModelConfiguration
) throws -> (Measurement, [FeatureSignature], [FeatureSignature]) {
try autoreleasepool {
let start = ProcessInfo.processInfo.systemUptime
let model = try MLModel(contentsOf: url, configuration: configuration)
let loadMs = (ProcessInfo.processInfo.systemUptime - start) * 1_000
let inputSignatures = try signatures(model.modelDescription.inputDescriptionsByName)
let outputSignatures = try signatures(model.modelDescription.outputDescriptionsByName)
var inputs: [String: Any] = [:]
for (name, description) in model.modelDescription.inputDescriptionsByName {
guard let constraint = description.multiArrayConstraint else {
throw ProbeError.unsupportedFeature(name)
}
let array = try MLMultiArray(shape: constraint.shape, dataType: constraint.dataType)
let itemBytes: Int
switch array.dataType {
case .float16: itemBytes = 2
case .float32, .int32: itemBytes = 4
case .double: itemBytes = 8
default: throw ProbeError.unsupportedFeature(name)
}
memset(array.dataPointer, 0, array.count * itemBytes)
if name == "audio_length" { array[0] = 32 }
if name == "language_mask" { array[0] = 1 }
if name == "token" { array[0] = 13087 }
inputs[name] = array
}
let features = try MLDictionaryFeatureProvider(dictionary: inputs)
var milliseconds: [Double] = []
for _ in 0..<4 {
let start = ProcessInfo.processInfo.systemUptime
let output = try model.prediction(from: features)
milliseconds.append((ProcessInfo.processInfo.systemUptime - start) * 1_000)
for name in output.featureNames {
guard let array = output.featureValue(for: name)?.multiArrayValue else {
throw ProbeError.unsupportedFeature(name)
}
for index in 0..<min(array.count, 16) where !array[index].doubleValue.isFinite {
throw ProbeError.nonFiniteOutput(name)
}
}
}
let repeated = milliseconds.dropFirst().sorted()
return (
Measurement(loadMs: loadMs, firstPredictionMs: milliseconds[0],
medianRepeatedPredictionMs: repeated[repeated.count / 2]),
inputSignatures, outputSignatures)
}
}
private enum ProbeError: Error {
case unsupportedFeature(String)
case nonFiniteOutput(String)
}
}
#if LOAD_PROBE_CLI
@main
struct LoadProbeCLI {
static func main() async throws {
guard CommandLine.arguments.count == 4 else {
throw NSError(domain: "NemotronExportProbe", code: 1, userInfo: [
NSLocalizedDescriptionKey: "usage: probe bundle-directory ane|cpu|gpu|all output.json",
])
}
let units: MLComputeUnits
switch CommandLine.arguments[2] {
case "ane": units = .cpuAndNeuralEngine
case "cpu": units = .cpuOnly
case "gpu": units = .cpuAndGPU
case "all": units = .all
default: throw NSError(domain: "NemotronExportProbe", code: 2)
}
let report = try await NemotronExportProbe.run(
bundleURL: URL(fileURLWithPath: CommandLine.arguments[1], isDirectory: true), computeUnits: units)
let encoder = JSONEncoder()
encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
let json = try encoder.encode(report)
try json.write(to: URL(fileURLWithPath: CommandLine.arguments[3]))
print(String(decoding: json, as: UTF8.self))
}
}
#endif
|