Metal cleanup (#81)

* Create shaders.metal

* Cleanup: use FragmentBuilder

* Remove unused code

* Remove unused functions

* Cleanup

* Add extension

* Reorder

* Use genericFragment

* Delete FragmentBuilder.swift

* Remove unused code

Sources/AudioKitUI/AudioKitUI.swift

@@ -14,6 +14,16 @@ public extension Color {
     var cg: CGColor {
         return CrossPlatformColor(self).cgColor
     }
+
+    var simd: SIMD4<Float> {
+        if let comps = cg.components {
+            return .init(Float(comps[0]),
+                         Float(comps[1]),
+                         Float(comps[2]),
+                         Float(comps[3]))
+        }
+        return .zero
+    }
 }
 
 public extension EnvironmentValues {

Sources/AudioKitUI/Visualizations/FloatPlot.swift

@@ -4,47 +4,32 @@ import AudioKit
 import Metal
 import MetalKit
 
+// This must be in sync with the definition in shaders.metal
+public struct FragmentConstants {
+    public var foregroundColor: SIMD4<Float>
+    public var backgroundColor: SIMD4<Float>
+    public var isFFT: Bool
+    public var isCentered: Bool
+    public var isFilled: Bool
+
+    // Padding is required because swift doesn't pad to alignment
+    // like MSL does.
+    public var padding: Int = 0
+}
+
 public class FloatPlot: MTKView, MTKViewDelegate {
     let waveformTexture: MTLTexture!
     let commandQueue: MTLCommandQueue!
     let pipelineState: MTLRenderPipelineState!
     let bufferSampleCount: Int
-    let parameterBuffer: MTLBuffer!
-    let colorParameterBuffer: MTLBuffer!
     var dataCallback: () -> [Float]
-
-    let metalHeader = """
-    #include <metal_stdlib>
-    using namespace metal;
-
-    struct VertexOut {
-    float4 position  [[ position ]];
-    float2 t;
-    };
-
-    constant float2 verts[4] = { float2(-1, -1), float2(1, -1), float2(-1, 1), float2(1, 1) };
-
-    vertex VertexOut textureVertex(uint vid [[ vertex_id ]]) {
-
-    VertexOut out;
-    out.position = float4(verts[vid], 0.0, 1.0);
-    out.t = (verts[vid] + float2(1)) * .5;
-    out.t.y = 1.0 - out.t.y;
-    return out;
-
-    }
-
-    constexpr sampler s(coord::normalized,
-    filter::linear);
-
-    fragment half4 textureFragment(VertexOut in [[ stage_in ]],
-    texture1d<float, access::sample> waveform, device float* parameters, device float4* colorParameters) {
-    """
+    var constants: FragmentConstants
 
     public init(frame frameRect: CGRect,
-                fragment: String? = nil,
+                constants: FragmentConstants,
                 dataCallback: @escaping () -> [Float]) {
         self.dataCallback = dataCallback
+        self.constants = constants
         bufferSampleCount = Int(frameRect.width)
 
         let desc = MTLTextureDescriptor()
@@ -58,20 +43,10 @@ public class FloatPlot: MTKView, MTKViewDelegate {
         waveformTexture = device?.makeTexture(descriptor: desc)
         commandQueue = device!.makeCommandQueue()
 
-        let defaultFragment = """
-        float sample = waveform.sample(s, in.t.x).x;
-        float y = (in.t.y - .5);
-        float d = fabs(y - sample);
-        float alpha = fabs(1/(50 * d));
-        return alpha;
-        """
-
-        let metal = metalHeader + (fragment ?? defaultFragment) + "}"
-//        let library = device!.makeDefaultLibrary()!
-        let library = try! device?.makeLibrary(source: metal, options: nil)
+        let library = try! device?.makeDefaultLibrary(bundle: Bundle.module)
 
-        let fragmentProgram = library!.makeFunction(name: "textureFragment")!
-        let vertexProgram = library!.makeFunction(name: "textureVertex")!
+        let fragmentProgram = library!.makeFunction(name: "genericFragment")!
+        let vertexProgram = library!.makeFunction(name: "waveformVertex")!
 
         let pipelineStateDescriptor = MTLRenderPipelineDescriptor()
         pipelineStateDescriptor.vertexFunction = vertexProgram
@@ -88,11 +63,6 @@ public class FloatPlot: MTKView, MTKViewDelegate {
 
         pipelineState = try! device!.makeRenderPipelineState(descriptor: pipelineStateDescriptor)
 
-        parameterBuffer = device!.makeBuffer(length: 128 * MemoryLayout<Float>.size,
-                                             options: .storageModeShared)
-        colorParameterBuffer = device!.makeBuffer(length: 128 * MemoryLayout<SIMD4<Float>>.size,
-                                                  options: .storageModeShared)
-
         super.init(frame: frameRect, device: device)
 
         clearColor = .init(red: 0.0, green: 0.0, blue: 0.0, alpha: 0)
@@ -140,8 +110,8 @@ public class FloatPlot: MTKView, MTKViewDelegate {
 
                 encoder.setRenderPipelineState(pipelineState)
                 encoder.setFragmentTexture(waveformTexture, index: 0)
-                encoder.setFragmentBuffer(parameterBuffer, offset: 0, index: 0)
-                encoder.setFragmentBuffer(colorParameterBuffer, offset: 0, index: 1)
+                assert(MemoryLayout<FragmentConstants>.size == 48)
+                encoder.setFragmentBytes(&constants, length: MemoryLayout<FragmentConstants>.size, index: 0)
                 encoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4)
                 encoder.endEncoding()
 
@@ -154,16 +124,4 @@ public class FloatPlot: MTKView, MTKViewDelegate {
             commandBuffer.waitUntilCompleted()
         }
     }
-
-    func setParameter(address: Int, value: Float) {
-        if address >= 0, address < 128 {
-            parameterBuffer.contents().assumingMemoryBound(to: Float.self)[address] = value
-        }
-    }
-
-    func setColorParameter(address: Int, value: SIMD4<Float>) {
-        if address >= 0, address < 128 {
-            colorParameterBuffer.contents().assumingMemoryBound(to: SIMD4<Float>.self)[address] = value
-        }
-    }
 }

Sources/AudioKitUI/Visualizations/NodeFFTView.swift

@@ -9,37 +9,23 @@ public struct NodeFFTView: ViewRepresentable {
     var nodeTap: FFTTap
     let bufferSampleCount = 128
 
-    let foregroundColorAddress = 0
-    let backgroundColorAddress = 1
-
     public init(_ node: Node) {
         nodeTap = FFTTap(node, bufferSize: UInt32(bufferSampleCount), callbackQueue: .main) { _ in }
     }
 
     internal var plot: FloatPlot {
         nodeTap.start()
 
-        let metalFragmentOrig = """
-        float sample = waveform.sample(s, (pow(10, in.t.x) - 1.0) / 9.0).x;
-
-        half4 backgroundColor = half4(colorParameters[1]);
-        half4 foregroundColor = half4(colorParameters[0]);
+        let constants = FragmentConstants(foregroundColor: Color.yellow.simd,
+                                          backgroundColor: Color.black.simd,
+                                          isFFT: true,
+                                          isCentered: false,
+                                          isFilled: true)
 
-        float y = (in.t.y - 1);
-        bool isFilled = parameters[0] != 0;
-        float d = isFilled ? fmax(fabs(y) - fabs(sample), 0) : fabs(y - sample);
-        float alpha = fabs(1/(50 * d));
-        return { mix(foregroundColor, backgroundColor, alpha) };
-        """
-
-        let plot = FloatPlot(frame: CGRect(x: 0, y: 0, width: 1024, height: 1024), fragment: metalFragmentOrig) {
+        let plot = FloatPlot(frame: CGRect(x: 0, y: 0, width: 1024, height: 1024), constants: constants) {
             nodeTap.fftData
         }
 
-        plot.setParameter(address: 0, value: 1)
-        plot.setColorParameter(address: foregroundColorAddress, value: SIMD4<Float>(1, 1, 0, 1))
-        plot.setColorParameter(address: backgroundColorAddress, value: SIMD4<Float>(0, 0, 0, 1))
-
         return plot
     }
 

Sources/AudioKitUI/Visualizations/NodeOutputView.swift

@@ -7,20 +7,21 @@ import SwiftUI
 
 public struct NodeOutputView: ViewRepresentable {
     private var nodeTap: RawDataTap
-    private var metalFragment: FragmentBuilder
+    private let constants: FragmentConstants
 
     public init(_ node: Node, color: Color = .gray, backgroundColor: Color = .clear, bufferSize: Int = 1024) {
-        metalFragment = FragmentBuilder(foregroundColor: color.cg,
-                                        backgroundColor: backgroundColor.cg,
-                                        isCentered: true,
-                                        isFilled: false)
+        constants = FragmentConstants(foregroundColor: color.simd,
+                                      backgroundColor: backgroundColor.simd,
+                                      isFFT: false,
+                                      isCentered: true,
+                                      isFilled: false)
         nodeTap = RawDataTap(node, bufferSize: UInt32(bufferSize), callbackQueue: .main)
     }
 
     var plot: FloatPlot {
         nodeTap.start()
 
-        return FloatPlot(frame: CGRect(x: 0, y: 0, width: 1024, height: 1024), fragment: metalFragment.stringValue) {
+        return FloatPlot(frame: CGRect(x: 0, y: 0, width: 1024, height: 1024), constants: constants) {
             return nodeTap.data
         }
     }

Sources/AudioKitUI/Visualizations/NodeRollingView.swift

@@ -42,27 +42,28 @@ public class RollingViewData {
 
 public struct NodeRollingView: ViewRepresentable {
     private let nodeTap: RawDataTap
-    private let metalFragment: FragmentBuilder
     private let rollingData: RollingViewData
+    private let constants: FragmentConstants
 
     public init(_ node: Node,
                 color: Color = .gray,
                 backgroundColor: Color = .clear,
                 isCentered: Bool = false,
                 isFilled: Bool = false,
                 bufferSize: UInt32 = 1024) {
-        metalFragment = FragmentBuilder(foregroundColor: color.cg,
-                                        backgroundColor: backgroundColor.cg,
-                                        isCentered: isCentered,
-                                        isFilled: isFilled)
+        constants = FragmentConstants(foregroundColor: color.simd,
+                                      backgroundColor: backgroundColor.simd,
+                                      isFFT: false,
+                                      isCentered: isCentered,
+                                      isFilled: isFilled)
         nodeTap = RawDataTap(node, bufferSize: bufferSize, callbackQueue: .main)
         rollingData = RollingViewData(bufferSize: bufferSize)
     }
 
     var plot: FloatPlot {
         nodeTap.start()
 
-        return FloatPlot(frame: CGRect(x: 0, y: 0, width: 1024, height: 1024), fragment: metalFragment.stringValue) {
+        return FloatPlot(frame: CGRect(x: 0, y: 0, width: 1024, height: 1024), constants: constants) {
             rollingData.calculate(nodeTap)
         }
     }

Sources/AudioKitUI/Visualizations/shaders.metal

@@ -0,0 +1,44 @@
+
+#include <metal_stdlib>
+using namespace metal;
+
+struct VertexOut {
+    float4 position  [[ position ]];
+    float2 t;
+};
+
+constant float2 verts[4] = { float2(-1, -1), float2(1, -1), float2(-1, 1), float2(1, 1) };
+
+vertex VertexOut waveformVertex(uint vid [[ vertex_id ]]) {
+
+    VertexOut out;
+    out.position = float4(verts[vid], 0.0, 1.0);
+    out.t = (verts[vid] + float2(1)) * .5;
+    out.t.y = 1.0 - out.t.y;
+    return out;
+
+}
+
+constexpr sampler s(coord::normalized,
+                    filter::linear);
+
+// This must be in sync with the definition in FloatPlot.swift
+struct FragmentConstants {
+    float4 foregroundColor;
+    float4 backgroundColor;
+    bool isFFT;
+    bool isCentered;
+    bool isFilled;
+};
+
+fragment half4 genericFragment(VertexOut in [[ stage_in ]],
+                               texture1d<float, access::sample> waveform,
+                               constant FragmentConstants& c) {
+
+    float sample = waveform.sample(s, c.isFFT ? (pow(10, in.t.x) - 1.0) / 9.0 : in.t.x).x;
+
+    float y = (-in.t.y + (c.isCentered ? 0.5 : 1));
+    float d = c.isFilled ? fmax(fabs(y) - fabs(sample), 0) : fabs(y - sample);
+    float alpha = c.isFFT ? fabs(1/(50 * d)) : smoothstep(0.01, 0.02, d);
+    return half4( mix(c.foregroundColor, c.backgroundColor, alpha) );
+}