Implement CoreML Custom Layer With Two Inputs - metal

Question

I have a tensorflow graph that I want to convert to CoreML, but it uses some operations that are missing, which I will have to implement as Custom Layers.
The two operations I'm focussing on now are Sin and FloorDiv.
Sin was quite straightforward, I could follow this tutorial, and I have a working Swift class and Metal kernel that does the job, which I tested with a toy coreml file:
import Foundation
import CoreML
import Accelerate
#objc(Sin) class Sin: NSObject, MLCustomLayer {
let sinPipeline: MTLComputePipelineState
required init(parameters: [String : Any]) throws {
print(#function, parameters)
let sinFunction = GPUDispatch.sharedInstance.library.makeFunction(name: "sin")!
sinPipeline = try! GPUDispatch.sharedInstance.device.makeComputePipelineState(
function: sinFunction)
super.init()
}
func setWeightData(_ weights: [Data]) throws {
print(#function, weights)
}
func outputShapes(forInputShapes inputShapes: [[NSNumber]]) throws
-> [[NSNumber]] {
print(#function, inputShapes)
return inputShapes
}
func evaluate(inputs: [MLMultiArray], outputs: [MLMultiArray]) throws {
for i in 0..<inputs.count {
let input = inputs[i]
let output = outputs[i]
var count = Int32(input.count)
let iptr = UnsafeMutablePointer<Float>(OpaquePointer(input.dataPointer))
let optr = UnsafeMutablePointer<Float>(OpaquePointer(output.dataPointer))
vvsinf(optr, iptr, &count)
}
}
func encode(commandBuffer: MTLCommandBuffer,
inputs: [MTLTexture], outputs: [MTLTexture]) throws {
if let encoder = commandBuffer.makeComputeCommandEncoder() {
for i in 0..<inputs.count {
encoder.setTexture(inputs[i], index: 0)
encoder.setTexture(outputs[i], index: 1)
encoder.dispatch(pipeline: sinPipeline, texture: inputs[i])
encoder.endEncoding()
}
}
}
}
and in Sin.metal:
kernel void sin(
texture2d_array<half, access::read> inTexture [[texture(0)]],
texture2d_array<half, access::write> outTexture [[texture(1)]],
ushort3 gid [[thread_position_in_grid]])
{
if (gid.x >= outTexture.get_width() ||
gid.y >= outTexture.get_height()) {
return;
}
const float4 x = float4(inTexture.read(gid.xy, gid.z));
const float4 y = sin(x);
outTexture.write(half4(y), gid.xy, gid.z);
}
What I don't understand is how this would work if the custom layer has two inputs, such as I would require for FloorDiv, which returns floor(x / y).
How would I adapt the Sin class I provided to produce something like sin(x*y), even if it's just on the CPU? Are there any other good tutorials for this sort of thing?

Answer 1

The pattern is different than I expected, but is quite obvious now I've played with the code some more.
This is a class that implements FloorDiv:
import Foundation
import CoreML
import Accelerate
#objc(FloorDiv) class FloorDiv: NSObject, MLCustomLayer {
let floorDivPipeline: MTLComputePipelineState
required init(parameters: [String : Any]) throws {
print(#function, parameters)
let floorDivFunction = GPUDispatch.sharedInstance.library.makeFunction(name: "floordiv")!
floorDivPipeline = try! GPUDispatch.sharedInstance.device.makeComputePipelineState(
function: floorDivFunction)
super.init()
}
func setWeightData(_ weights: [Data]) throws {
print(#function, weights)
}
func outputShapes(forInputShapes inputShapes: [[NSNumber]]) throws
-> [[NSNumber]] {
print(#function, inputShapes)
return inputShapes
}
func evaluate(inputs: [MLMultiArray], outputs: [MLMultiArray]) throws {
let numerator = inputs[0]
let denominator = inputs[1]
var output = outputs[0]
assert(numerator.count == denominator.count)
var count = Int32(numerator.count)
let numerator_ptr = UnsafeMutablePointer<Float>(OpaquePointer(numerator.dataPointer))
let denominator_ptr = UnsafeMutablePointer<Float>(OpaquePointer(denominator.dataPointer))
let output_ptr = UnsafeMutablePointer<Float>(OpaquePointer(output.dataPointer))
vvdivf(output_ptr, numerator_ptr, denominator_ptr, &count)
vvfloorf(output_ptr, output_ptr, &count)
}
func encode(commandBuffer: MTLCommandBuffer,
inputs: [MTLTexture], outputs: [MTLTexture]) throws {
if let encoder = commandBuffer.makeComputeCommandEncoder() {
encoder.setTexture(inputs[0], index: 0)
encoder.setTexture(inputs[1], index: 1)
encoder.setTexture(outputs[0], index: 2)
encoder.dispatch(pipeline: floorDivPipeline, texture: inputs[0])
encoder.endEncoding()
}
}
}
And here is the Metal kernel:
#include <metal_stdlib>
using namespace metal;
kernel void floordiv(
texture2d_array<half, access::read> inTexture [[texture(0)]],
texture2d_array<half, access::read> inTexture2 [[texture(1)]],
texture2d_array<half, access::write> outTexture [[texture(2)]],
ushort3 gid [[thread_position_in_grid]])
{
if (gid.x >= outTexture.get_width() ||
gid.y >= outTexture.get_height()) {
return;
}
const float4 x = float4(inTexture.read(gid.xy, gid.z));
const float4 x2 = float4(inTexture2.read(gid.xy, gid.z));
const float4 y = floor(x / x2);
outTexture.write(half4(y), gid.xy, gid.z);
}