| import CoreML |
| import Darwin |
| import Foundation |
|
|
| |
| |
| @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 |
|
|
