MLCompute Playground

MLCompute.swift

import Foundation
import MLCompute
import PlaygroundSupport

let page = PlaygroundPage.current
page.needsIndefiniteExecution = true

let tensor1 = MLCTensor(descriptor: MLCTensorDescriptor(shape: [6, 1], dataType: .float32)!)
let tensor2 = MLCTensor(descriptor: MLCTensorDescriptor(shape: [6, 1], dataType: .float32)!)
let tensor3 = MLCTensor(descriptor: MLCTensorDescriptor(shape: [6, 1], dataType: .float32)!)

let buffer1: [Float] = [1, 2, 3, 4, 5, 6]
let buffer2: [Float] = [7, 8, 9, 10, 11, 12]
let buffer3: [Float] = [6, 5, 4, 3, 2, 1]

let data1 = MLCTensorData(immutableBytesNoCopy: UnsafeRawPointer(buffer1), length: buffer1.count * MemoryLayout<Float>.size)
let data2 = MLCTensorData(immutableBytesNoCopy: UnsafeRawPointer(buffer2), length: buffer2.count * MemoryLayout<Float>.size)
let data3 = MLCTensorData(immutableBytesNoCopy: UnsafeRawPointer(buffer3), length: buffer3.count * MemoryLayout<Float>.size)

let g = MLCGraph()
let tensor1plus2 = g.node(with: MLCArithmeticLayer(operation: .add), sources: [tensor1, tensor2])
g.node(with: MLCArithmeticLayer(operation: .add), sources: [tensor1plus2!, tensor3])

let i = MLCInferenceGraph(graphObjects: [g])
i.addInputs(["data1" : tensor1, "data2" : tensor2, "data3" : tensor3])
i.compile(options: .debugLayers, device: MLCDevice())

i.execute(inputsData: ["data1" : data1, "data2" : data2, "data3" : data3],
          batchSize: 0,
          options: []) { (r, e, time) in
    print("Error: \(String(describing: e))")
    print("Result: \(String(describing: r))")

    let buffer3 = UnsafeMutableRawPointer.allocate(byteCount: 6 * MemoryLayout<Float>.size, alignment: MemoryLayout<Float>.alignment)
    
    r!.copyDataFromDeviceMemory(toBytes: buffer3, length: 6 * MemoryLayout<Float>.size, synchronizeWithDevice: false)
    
    let float4Ptr = buffer3.bindMemory(to: Float.self, capacity: 6)
    let float4Buffer = UnsafeBufferPointer(start: float4Ptr, count: 6)
    print(Array(float4Buffer))
    
    page.finishExecution()
}