| |
| |
| import CoreML |
| import Foundation |
|
|
| let dir = URL(fileURLWithPath: CommandLine.arguments[1]) |
| let seq = 64 |
| let raw = try Data(contentsOf: dir.appendingPathComponent("bench_codes.f16")) |
| let nCites = raw.count / (seq * 2) |
|
|
| func load(_ b: Int, _ units: MLComputeUnits) throws -> MLModel { |
| let cfg = MLModelConfiguration() |
| cfg.computeUnits = units |
| return try MLModel(contentsOf: dir.appendingPathComponent("CiteNet_b\(b).mlmodelc"), configuration: cfg) |
| } |
|
|
| func batch(_ b: Int, offset: Int) throws -> MLDictionaryFeatureProvider { |
| let arr = try MLMultiArray(shape: [NSNumber(value: b), 1, 1, NSNumber(value: seq)], dataType: .float16) |
| raw.withUnsafeBytes { src in |
| let dst = arr.dataPointer |
| for i in 0..<b { |
| let c = (offset + i) % nCites |
| memcpy(dst + i * seq * 2, src.baseAddress! + c * seq * 2, seq * 2) |
| } |
| } |
| return try MLDictionaryFeatureProvider(dictionary: ["codes": arr]) |
| } |
|
|
| func now() -> UInt64 { DispatchTime.now().uptimeNanoseconds } |
|
|
| func pct(_ v: [Double], _ p: Double) -> Double { v.sorted()[min(v.count - 1, Int(Double(v.count) * p))] } |
|
|
| let unitsName: [(String, MLComputeUnits)] = [("cpu", .cpuOnly), ("ane", .cpuAndNeuralEngine), ("gpu", .cpuAndGPU)] |
| let only = CommandLine.arguments.count > 2 ? CommandLine.arguments[2] : "all" |
| for b in [1, 32, 512, 2048, 8192] where only == "all" { |
| for (uname, units) in unitsName { |
| if b == 1 && uname == "gpu" { continue } |
| if b >= 2048 && uname == "cpu" { continue } |
| let m = try load(b, units) |
| let inputs = try (0..<8).map { try batch(b, offset: $0 * b) } |
| for i in 0..<50 { _ = try m.prediction(from: inputs[i % 8]) } |
| let iters = b == 1 ? 5000 : (b == 32 ? 3000 : max(40, 300_000 / b)) |
| var lat = [Double]() |
| let t0 = now() |
| for i in 0..<iters { |
| let s = now() |
| _ = try m.prediction(from: inputs[i % 8]) |
| lat.append(Double(now() - s) / 1e3) |
| } |
| let total = Double(now() - t0) / 1e9 |
| print(String(format: "b%-4d %@ p50 %7.1f us p99 %7.1f us %8.0f cites/s", b, uname, |
| pct(lat, 0.5), pct(lat, 0.99), Double(iters * b) / total)) |
| } |
| } |
|
|
| |
| func concurrent(_ lanes: [(String, Int, MLComputeUnits)], seconds: Double) throws { |
| let models = try lanes.map { try load($0.1, $0.2) } |
| let inputs = try lanes.map { l in try (0..<4).map { try batch(l.1, offset: $0 * l.1) } } |
| for (i, m) in models.enumerated() { for j in 0..<20 { _ = try m.prediction(from: inputs[i][j % 4]) } } |
| var counts = [Int](repeating: 0, count: lanes.count) |
| let lock = NSLock() |
| let start = now(), deadline = start + UInt64(seconds * 1e9) |
| DispatchQueue.concurrentPerform(iterations: lanes.count) { i in |
| var n = 0, j = 0 |
| while now() < deadline { |
| _ = try! models[i].prediction(from: inputs[i][j % 4]); j += 1; n += lanes[i].1 |
| } |
| lock.lock(); counts[i] = n; lock.unlock() |
| } |
| let total = Double(now() - start) / 1e9 |
| let parts = zip(lanes, counts).map { String(format: "%@ b%d %.0fk/s", $0.0.0, $0.0.1, Double($0.1) / total / 1e3) } |
| print(String(format: "concurrent %@ -> %.0f cites/s", parts.joined(separator: " + "), Double(counts.reduce(0, +)) / total)) |
| } |
| try concurrent([("ane", 512, .cpuAndNeuralEngine)], seconds: 3) |
| try concurrent([("gpu", 8192, .cpuAndGPU)], seconds: 3) |
| try concurrent([("ane", 512, .cpuAndNeuralEngine), ("gpu", 8192, .cpuAndGPU)], seconds: 3) |
| try concurrent([("ane", 2048, .cpuAndNeuralEngine), ("gpu", 8192, .cpuAndGPU)], seconds: 3) |
| try concurrent([("ane", 512, .cpuAndNeuralEngine), ("ane", 512, .cpuAndNeuralEngine), ("gpu", 8192, .cpuAndGPU), ("gpu", 8192, .cpuAndGPU)], seconds: 3) |
|
|
