Discover compilation workflows in Metal

// Declare external functions

extern float4 foo(FragmentInput input);
extern float4 bar(FragmentInput input);

// Use functions in shader

fragment float4 main(FragmentInput input [[stage_in]])
{
    switch(condition(input)) 
    {
        case 0: 
            return foo(input);
        case 1:
            return bar(input);
    }
}

// Declare a descriptor and set CompileToBinary options

MTLFunctionDescriptor* functionDescriptor = [MTLFunctionDescriptor new];
functionDescriptor.options = MTLFunctionOptionCompileToBinary;

// Backend compile the function

functionDescriptor.name = @"foo";
id<MTLFunction> foo = [library newFunctionWithDescriptor:functionDescriptor

// Provide a list of functions that the pipeline stage may call

// AIR functions

renderPipeDesc.fragmentLinkedFunctions.functions = @[foo, bar, baz];

// Binary functions

renderPipeDesc.fragmentLinkedFunctions.binaryFunctions = @[foo, bar, baz];

// Create visible function table

    [renderPipeline newVisibleFunctionTableWithDescriptor:stage:];

// Create function handles

    [renderPipeline functionHandleWithFunction:stage:];

// Insert handles into table

    [visibleFunctionTable setFunction:atIndex:];

// Bind visible function table objects to each stage

    [renderCommandEncoder setFragmentVisibleFunctionTable:atBufferIndex:];

// Usage in shader

   fragment float4 shaderFunc(FragmentData vo[[stage_in]],
                              visible_function_table<float4(float3)>materials[[buffer(0)]])
   {
   		 //...
     
       return materials[materialSelector](coord);
   }

// Enable incrementally adding binary functions per stage

renderPipeDesc.supportAddingFragmentBinaryFunctions = YES;

// Create render pipeline functions descriptor

MTLRenderPipelineFunctionsDescriptor extraDesc;
extraDesc.fragmentAdditionalBinaryFunctions = @[bat];

// Instantiate render pipeline state

id<MTLRenderPipelineState> renderPipeline2 =
  [renderPipeline1 newRenderPipelineStateWithAdditionalBinaryFunctions:extraDesc

[[stitchable]] int FunctionA(device int*, int) {…}
[[stitchable]] int FunctionC(int, int) {…}

[[stitchable]]
int ResultFunction(device int* Input0,
                   int Input1, 
                   int Input2)
{
  int N0 = FunctionA(Input0, Input1);
  int N1 = FunctionA(Input0, Input2);
  int N2 = FunctionC(N0, N1);    
  return N2;
}

// Create input nodes

  inputs[0] = [[MTLFunctionStitchingInputNode alloc] initWithArgumentIndex:0];

// Create function nodes

  n0 = [[MTLFunctionStitchingFunctionNode alloc] initWithName:@"FunctionA"
                                                    arguments:@[inputs[0], inputs[1]]
                                          controlDependencies:@[]];
  n1 = [[MTLFunctionStitchingFunctionNode alloc] initWithName:@"FunctionA"
                                                    arguments:@[inputs[0], inputs[2]]
                                          controlDependencies:@[]];
  n2 = [[MTLFunctionStitchingFunctionNode alloc] initWithName:@"FunctionC"
                                                    arguments:@[n0, n1]
                                          controlDependencies:@[]];

// Create graph

  graph = [[MTLFunctionStitchingGraph alloc] initWithFunctionName:@"ResultFunction"
                                                            nodes:@[n0, n1]
                                                       outputNode:n2
                                                       attributes:@[]];

// Configure stitched library descriptor
  
  MTLStitchedLibraryDescriptor* descriptor = [MTLStitchedLibraryDescriptor new];

  descriptor.functions      = @[stitchableFunctions];
  descriptor.functionGraphs = @[graph];

// Create stitched function

  id<MTLLibrary> lib = [device newLibraryWithDescriptor:descriptor 
                                                  error:&error];

  id<MTLFunction> stitchedFunction = [lib newFunctionWithName:@"ResultFunction"];