| |
| |
| import CoreML |
| import Foundation |
|
|
| func deviceName(_ d: MLComputeDevice) -> String { |
| switch d { |
| case .cpu: return "CPU" |
| case .gpu: return "GPU" |
| case .neuralEngine: return "ANE" |
| @unknown default: return "?" |
| } |
| } |
|
|
| struct OpRecord: Codable { |
| let index: Int |
| let op: String |
| let outputs: [String] |
| let preferred: String |
| let supported: [String] |
| let cost: Double |
| } |
|
|
| func walk(_ block: MLModelStructure.Program.Block, plan: MLComputePlan, into records: inout [OpRecord]) { |
| for op in block.operations { |
| let usage = plan.deviceUsage(for: op) |
| let cost = plan.estimatedCost(of: op)?.weight ?? 0 |
| records.append(OpRecord( |
| index: records.count, |
| op: op.operatorName, |
| outputs: op.outputs.map { $0.name }, |
| preferred: usage.map { deviceName($0.preferred) } ?? "none", |
| supported: usage.map { $0.supported.map(deviceName) } ?? [], |
| cost: cost)) |
| for inner in op.blocks { walk(inner, plan: plan, into: &records) } |
| } |
| } |
|
|
| let args = CommandLine.arguments |
| let modelURL = URL(fileURLWithPath: args[1]) |
| let outURL = URL(fileURLWithPath: args[2]) |
|
|
| let compiledURL: URL |
| if modelURL.pathExtension == "mlmodelc" { |
| compiledURL = modelURL |
| } else { |
| let t0 = Date() |
| let tmp = try await MLModel.compileModel(at: modelURL) |
| compiledURL = modelURL.deletingPathExtension().appendingPathExtension("mlmodelc") |
| try? FileManager.default.removeItem(at: compiledURL) |
| try FileManager.default.moveItem(at: tmp, to: compiledURL) |
| print("compiled in \(String(format: "%.1f", Date().timeIntervalSince(t0)))s -> \(compiledURL.lastPathComponent)") |
| } |
|
|
| var result: [String: [OpRecord]] = [:] |
| for (label, units) in [("cpuAndNeuralEngine", MLComputeUnits.cpuAndNeuralEngine), ("all", MLComputeUnits.all)] { |
| let config = MLModelConfiguration() |
| config.computeUnits = units |
| let plan = try await MLComputePlan.load(contentsOf: compiledURL, configuration: config) |
| guard case let .program(program) = plan.modelStructure, let main = program.functions["main"] else { |
| fatalError("not an ML program") |
| } |
| var records: [OpRecord] = [] |
| walk(main.block, plan: plan, into: &records) |
| result[label] = records |
| print("\(label): \(records.count) ops") |
| } |
| let enc = JSONEncoder() |
| enc.outputFormatting = [.prettyPrinted, .sortedKeys] |
| try enc.encode(result).write(to: outURL) |
| print("wrote \(outURL.path)") |
|
|
