MetalLink` - General Purpose GPU Computation on Apple Silicon

Soumya Sambeet Mohapatra
Application Developer, ADDITIVE GmbH

(RowDefault[, (Spacer1), , (Spacer1), ])  #WolframTechConf

ABSTRACT

Apple has dropped support for CUDA, and has fully transitioned to its own API called Metal. Wolfram Language lacks Metal support and this limits users from utilizing powerful features Apple Silicon offers. This has motivated the development of MetalLink, to use the ‘M’ GPUs from Mathematica.

Overview

Metal

Architecture

Features

Compute Workflow

Development

LibraryLink

Objective-C

Swift + C Bridge

Foreign Function Interface

Demonstrations

Basic functions: MetalMap, MetalAdd, MetalMultiply

MonteCarlo simulations for estimation of Pi

Image processing examples

Image Classification with Neural Networks

Metal

What is Metal?

Graphics & Compute API optimized for Apple Hardware

Ships with its own Shading Language

Low-overhead

Architecture

Apple Silicon chips come with integrated GPU

GPU core count ranges from 7 on base models to 64 on high-end hardware

Each core is split into 16 execution units, which each contain 8 Arithmetic Logic Units

Memory is unified: Both CPU and GPU have access to the same memory (Additionally the GPU has some extremely high speed cache memory)

Features

Metal ships with the Metal Shading Language, a variant of C++ with special datatypes to handle images, textures, etc.

Atomic datatypes  take care of synchronization and race conditions

On top of the Metal API, the Metal Performance Shaders is built. MPS provides highly optimized kernels for several operations

Like MPS, there is also MPSGraph which is developed to build and train Neural Networks on the GPU

Compute workflow

Write a kernel in MSL

"
kernel void add_arrays(
	device const float *inA,
	device const float *inB,
	device float *result,
	uint index[[thread_position_in_grid]])
{
	/**
	* for-loop is replaced with a collection of threads,
	* each of which calls this function.
	*/
	result[index] = inA[index] + inB[index];
}
"

Create a reference to Metal device

Create library and load it with the kernel

Create a pipeline state object with the library

Create a command queue

Create data buffers and load data into them

Create command buffer and command encoder

Organize threads and threadgroups

Encode the compute commands to execute the threads and commit the command buffer

Command Queue, Command Buffer & Command Encoder

Railway track === Command Queue

Cargo === Command Buffer

Cargo loader === Command Encoder

Threads & Threadgroups organization

Threads

Kernel functions are executed over 1D, 2D or 3D grids

Each point in the grid represents an instance of kernel

Threads are identified by their position in the grid using the [[thread_position_in_grid]] attribute qualifier

Threadgroups

Threadgroups are a group of threads that execute together & can share a common block of memory

Although kernel functions run independently, they can collaborate within a threadgroup

Threads within a threadgroup are identified using the [[thread_position_in_threadgroup]] attribute qualifier

Calculation of Threads & Threadgroups

Two important values are required

maximumTotalThreadsPerThreadgroup (sufficient for 1D)

threadExecutionWidth (required for 2D)

These values are calculated when a PSO is created and remain constant

2D Example with Uniform threadgroup size

Threadgroup width w = Thread execution width

Threadgroup height = Maximum threads per threadgroup / w

Threadgroups in the grid span over

Width = ((Image width + threadgroup width)/(Image width))

Height = ((Image height + threadgroup height)/(Image height))

Dispatched using the API dispatchThreadgroups

Metal also supports non-uniform threadgroups and this is recommended for better performance

In this case, only the total number threads and the threads per thread group is passed to the API dispatchThreads

Development Strategies

LibraryLink

LibraryLink enables loading external functions in the Wolfram Language

External functions must conform to certain signature

Wolfram Kernel’s memory space is shared, thus eliminating data transfers

But: Only possible to create external functions in C/C++

Swift + C  Bridge

Metal API is available only in Objective-C and Swift

Technically, it is possible to mix Objective-C and C, but has several drawbacks (codebase becomes messy, hard to maintain, poor documentation, lack of support for newer APIs, etc.)

Swift provides more features, enabling faster development and better maintainability. It has better documentation as well. Many newer API’s can be used only from Swift

Solution: Functions are written Swift and dynamically linked to a bridge written in C, which is then loaded in Mathematica

Drawbacks:

Multiple points of failure

Jumping around multiple frameworks

No support for bidirectional referencing

Swift + Foreign  Function  Interface

Swift functions can be exposed to C style languages using the “@_cdecl” attribute

"
@_cdecl(\"metalDeviceCreate\")
public func createMetalDevice() -> UnsafeMutableRawPointer? {
    guard let deviceLocal = MTLCreateSystemDefaultDevice() else {
        return nil
    }
    device = deviceLocal
    return Unmanaged.passUnretained(device!).toOpaque()
}
"

This makes it possible to load the functions directly in Mathematica with Foreign Function Interface

MetalDeviceCreate = ForeignFunctionLoad[
	$MetalLinkLibraryPath, 
	"metalDeviceCreate", {} -> "OpaqueRawPointer"
]

Caveats:

Slightly increased overhead due to function calls

Still requires manual memory management when transferring data

MetalCPP (To be researched)

MetalCPP is a low-overhead API for Metal

This will make it possible to develop everything in C++

Unfortunately, this is not so well documented and it is unclear if all the latest APIs are supported

Demonstrations

This section demonstrates the performance gain with GPU computation. The AbosluteTiming value is used as a comparison metric.

Device Information

MetalDeviceInformation[]//Dataset

Basic Demo Functions

MetalMap

Similar to CUDAMap

Supports 15 functions: (RefLink[Cos, paclet:ref/Cos]), (RefLink[Sin, paclet:ref/Sin]), (RefLink[Tan, paclet:ref/Tan]), (RefLink[ArcCos, paclet:ref/ArcCos]), (RefLink[ArcSin, paclet:ref/ArcSin]), (RefLink[ArcTan, paclet:ref/ArcTan]), (RefLink[Cosh, paclet:ref/Cosh]), (RefLink[Sinh, paclet:ref/Sinh]), (RefLink[Exp, paclet:ref/Exp]), (RefLink[Log, paclet:ref/Log]), (RefLink[Log10, paclet:ref/Log10]), (RefLink[Sqrt, paclet:ref/Sqrt]), (RefLink[Ceiling, paclet:ref/Ceiling]), (RefLink[Floor, paclet:ref/Floor]), or (RefLink[Abs, paclet:ref/Abs])

cpuComputationTimes=Table[{i,First[AbsoluteTiming[Sin[RandomReal[{-1,1},10^i]];]]},{i,1,8}];
gpuComputationTimes=Table[{i,First[AbsoluteTiming[MetalMap[Sin,NumericArray[RandomReal[{-1,1},10^i],"Real32"]];]]},{i,1,8}];
ListLinePlot[{cpuComputationTimes,gpuComputationTimes},PlotRange->Full,PlotMarkers->Automatic,PlotLegends->{"CPU","GPU"},AxesLabel->{Style["Exponent",16],Style["Absolute time",16]}]

MetalAdd

Adds two 1D arrays

cpuComputationTimes=Table[{
i,
First@AbsoluteTiming[
RandomReal[{-1,1},10^i]+RandomReal[{-1,1},10^i];
]},
{i,1,8}
];

gpuComputationTimes=Table[{
i,
First@AbsoluteTiming[
MetalAdd[
NumericArray[RandomReal[{-1,1},10^i],"Real32"],
NumericArray[RandomReal[{-1,1},10^i],"Real32"]
];]},
{i,1,8}
];
ListLinePlot[{cpuComputationTimes,gpuComputationTimes},PlotRange->Full,PlotMarkers->Automatic,PlotLegends->{"CPU","GPU"},AxesLabel->{Style["Exponent",16],Style["Absolute time",16]}]

The compute kernels are handcrafted, not very optimized

Buffer allocation takes up some time

CPU evaluation is in 64 bit precision while Metal supports only 32 bits

Operation mode is Synchronous. Asynchronous mode can provide significant boost.

Matrix operations

MetalMatrixAdd

matrices=Table[RandomReal[{-1,1},{1000,1000}],1000];
nMatrices=Table[NumericArray[RandomReal[{-1,1},{1000,1000}],"Real32"],1000];

MetalMatrixAdd[nMatrices];//AbsoluteTiming

Total[matrices];//AbsoluteTiming

MetalMatrixMultiply

m1=RandomReal[{-1,1},{10000,10000}];
m2=RandomReal[{-1,1},{10000,10000}];
n1=NumericArray[m1,"Real32"];
n2=NumericArray[m2,"Real32"];

m1.m2;//AbsoluteTiming

MetalMatrixMultiply[n1,n2];//AbsoluteTiming

MPSMatrix is used under the hood

MPSMatrix can store very large matrices efficiently in comparison to directly storing bytes in a Buffer

Work in progress: Matrix decomposition, Eigenvalues computation

Monte Carlo Simulation for estimation of π

Single Thread computation on CPU

approxPi[n_]:=4. Count[Map[Norm,RandomReal[{-1,1},{n,2}]],_?(#<=1&)]/n

Mean[Table[approxPi[10000],{5000}]]//AbsoluteTiming

Multi Thread computation on CPU

LaunchKernels[]

Mean[ParallelTable[approxPi[10000],{5000}]]//AbsoluteTiming

Computation on GPU

approxPiGPU[n_]:=(4./n) *MetalCount[
MetalNorm[MetalRandomMatrix[{-1.0,1.0},{n,2}]]]

approxPiGPU[50000000]//AbsoluteTiming

Image Processing

Gaussian Blur

i=ExampleData[{"TestImage","Marruecos"}];

ImageDimensions[i]

cpuComputationTimes=Table[{r,First@AbsoluteTiming[GaussianFilter[i,2^r];]},{r,1,9}];
gpuComputationTimes=Table[{r,First@AbsoluteTiming[MetalGaussianBlur[i,2^r];]},{r,9}];
ListLinePlot[{cpuComputationTimes,gpuComputationTimes},PlotRange->Full,PlotMarkers->Automatic,PlotLegends->{"CPU","GPU"},AxesLabel->{Style["Exponent",16],Style["Absolute time",16]}]

Work in progress: Histogram Equalization, Generic convolution, Laplacian Filter

Thank you
.
.
.
Questions?

References

https://developer.apple.com/metal/

https://developer.apple.com/documentation/metal/performing_calculations_on_a_gpu

https://developer.apple.com/documentation/metal/compute_passes/creating_threads_and_threadgroups

https://developer.apple.com/documentation/metal/compute_passes/calculating_threadgroup_and_grid_sizes