Search code examples
neural-networkshadermetalraytracing

How can you invoke a MetalPerformanceShadersGraph directly from Metal Shaders?


There are many resources on Apple's developer documentation showing how to use Metal from Swift for ML to infer, or more specifically how to call a Metal Performance Shader graph of a neural network to predict some result. However, all of these examples, like the one linked below start from Swift.

https://developer.apple.com/documentation/metalperformanceshadersgraph/training_a_neural_network_using_mps_graph

    func encodeInferenceCase(sourceTensorData: MPSGraphTensorData) -> MPSGraphTensorData {
        let fetch = graph.run(with: gCommandQueue,
                              feeds: [sourcePlaceholderTensor: sourceTensorData],
                              targetTensors: targetInferenceTensors,
                              targetOperations: targetInferenceOps)

        return fetch[targetInferenceTensors[0]]!
    }

How can you invoke a graph, like the one above, when already in Metal Shader? This way you don't have to pay the cost of switching back and forth between the CPU and GPU. The context for this is for ray tracing, and using models to help predict global illumination. NVidia's more complicated version is linked below.

https://research.nvidia.com/publication/2021-06_real-time-neural-radiance-caching-path-tracing


Solution

  • The answer is you cannot invoke a Metal Performance Shader from a Metal shader.

    But assuming you already have a Metal device object, a command queue and tensors the CPU overhead should be minimal.

    If that isn't acceptable for your application, you would have to implement your own graph operation(s) as Metal function(s).