In this chapter, I will introduce you to another cool tool I use every day called MTLSwift. What is MTLSwift? You might think it is something “Swifty” on the one hand and Metal-related on the other. And you will be right because this tool generates kernel encoders in Swift using Metal shaders. Same as Alloy, MTLSwift was also created by my colleague Andrey Volodin and co-maintained by me.
Before we dive into how this tool works, let’s understand how the idea of encoders code generation was born.
If you look a the shaders and the encoding code, you will see a correlation between the arguments and names used in the shaders code and the values passed and encoded on the CPU side.
This pattern repeats every time you create a kernel function and the encoder for it. Also, if you look at the kernel encoder, you will see that each one of them has the same structure: a pipeline state property, a constructor, taking a library as an argument, and the encoding logic, which takes textures, buffers, or small values as arguments. Such a simple structure of kernel encoder makes it easy to create a new encoder on the one hand. On the other hand, it contributes to the tendency when a developer starts copy-pasting the encoders and reusing them with modifications. And such behavior might become a source of bugs.
Given that, if you somehow extract the information from the kernel sources about the name of the kernel function, its arguments, and function constants, you may be able to create a generator of bugs-free encoders. The first approach that may come to mind is to parse the sources, create a logic of understanding keywords and operators, etc. But creating a source code parser from scratch is not an easy task. What if we could use a Metal compiler to help us with that?
The Metal compiler itself is a modified version of Apple’s Clang. Clang is a C language family front end for LLVM. LLVM’s front-end is responsible for parsing the source code, breaking it up into pieces according to a grammatical structure, checking it for errors. As a result, the front-end outputs an Abstract Syntax Tree (AST). The latter is a structured representation, which can be used for different purposes such as creating a symbol table, performing type checking, and finally generating code.
For example, if we dump AST from a simplified version of our adjustments kernel:
with the help of command:
we will get such AST:
As we can see, AST has a node-based structure, which can be easily parsed. Internally MTLSwift calls Metal compiler to output such AST, parses it, creates intermediate node-based representation, and extracts all needed information. In order to help MTLSwift to get the info about how we want to dispatch the kernel and also add extra info about the kernel function arguments, some custom annotations were introduced. Let’s take a look at them.
Customising of code generation
Every custom annotation starts with
mtlswift:. The program uses this declaration prefix to identify the start of a declaration. It must be written in a docstring way right before the kernel.
A dispatch type to use. All dispatch types have to be followed by either a constant amount of threads via literals (like
X, Y, Z), specifying a target texture to cover via
over:argument, or stating that amount of threads will be provided by the user by using
provided. You can see all of the examples in each section, but you can choose the combination yourself.
Dispatch threadgroups of a uniform threadgroup size.
depthdescribe the grid size.
Dispatch threads with threadgroups of non-uniform size.
exacttype if GPU supports non-uniform threadgroup size and
overif it doesn’t. This declaration requires a boolean function constant index to be passed to decide what dispatch type to use.
The dispatch type used by default. In this case, the user has to dispatch the kernel manually after calling
Specify the threadgroup size.
X, Y, Z
Allows to specify constant X, Y and Z dimensions for threadgroup size.
This parameter sets the pipeline state’s
This parameter sets the pipeline state’s
In this case, the user has to pass the threadgroup size and an argument to
The type of the buffers passed to the kernel.
The name of the buffers passed to the kernel.
Encoder’s name in generated Swift code. Must be followed by a valid Swift identifier.
Specifies the access visibility of the encoder. Must be followed by either
internalis the default.
Ok, let’s update our shaders code to support
MTLSwift. First, in
add the following code:
This new line is an entry point of the MTLSwift AST parser. Now let’s add custom annotations before the kernel:
This annotation tells that the encoder will dispatch a grid of the same dimension as destination texture with non-uniform threadgroups branching function onstant set at 0.
Next, let’s declare the types of tint and temperature values passed to the encoder:
The final version of the shader file should look like this:
Now let’s install MTLSwift.
To generate encoder for the kernel in the .metal file, you need to call MTLSwift’s
As a result you will get
Shaders.metal.swift file next to the shaders:
Import this file to the Xcode project and remove
Adjustments class doesn’t encapsulate temperature and tint properties, we need to move them to
Modify the settings in the
Inside the redraw function replace the dispatching code:
That’s it! Now you can compile and run the project. From this point, each time you modify the shaders, you don’t need to worry about the encoders at all. By calling MTLSwift, you automatically get 50% of the shaders-related job done, which means less code to maintain with fewer bugs to show up.
The final code can be found here.
Thank you for reading, see you next time 👋.