Programming Apple GPUs through Go and Metal Shading Language

Investigating Go, Cgo, Metal Shading Language, Metal Performance Shaders, and benchmarking different approaches to matrix multiplication

Below I’ll describe the process of using framework, how to interface with custom GPU code (shaders) written in the Go-based matrix multiplication operations. This was written to run on my M2 MacBook.

The layout of the source, .

GPUs and Floating-Point Parallelism

I assume most people at this point are intuitively familiar with the concept that GPUs are incredibly powerful at certain kinds of computational tasks; especially some which support machine learning. It wasn’t until I started playing around with Metal that I understood first-hand how much more powerful they can be than CPUs.

GPUs by design are extremely efficient at massively parallel floating-point operations which demand high memory bandwidth. My MacBook M2 has 8 CPU cores and 8 GPU cores, but for comparison, the contains 16896 CUDA cores with hundreds of additional specialized tensor cores. GPUs usually support ] [ is probably the most well-known GPU programming platform, which specific to Nvidia hardware. There are also mathematical frameworks available for and reasonably transparently with GPU hardware. .

Metal GPU Basics

Programming direct GPU computation is not as simple as writing code for on-device CPUs. When working with Apple’s Metal framework, a rough series of operations for executing code on the GPU looks like:

• Find an appropriate GPU device
• Create a queue for executing commands (ie the , otherwise a , meaning we don’t need to actually copy any data into GPU-specific physical memory
• is a derivative of ().

  Debugging of MSL code is pretty difficult. You can use the algorithm, combined with some loop unrolling which surprisingly improved its performance significantly. This is just for comparison purposes; normally the MPSMatrixMultiplication functionality from MPS would be more suitable.

  kernel void matrix_multiply_naive(
    device const MatrixParams *params,
    constant float *A,
    constant float *B,
    device float *C,
    // Indicates the thread's unique position within the entire grid of 
    // threads being executed. The uint2 type is a 2D coordinate, with 
    // fields x and y representing its indices on each axis.
    // This parameter is not directly provided from the calling code, 
    // but provided by the Metal framework
    uint2 gid [[thread_position_in_grid]]
  ) {
    if (gid.x >= params->a_rows || gid.y >= params->b_cols) {
      return; // This thread is out of matrix dimensionality range, do nothing
    }
  
    float sum = 0.0;
    int k;
  
    // Loop unrolling; improves performance by a notable margin
    for (k = 0; k <= params->a_cols - 4; k += 4) {
      sum += A[gid.x * params->a_cols + k] 
         * B[k * params->b_cols + gid.y];
      sum += A[gid.x * params->a_cols + k + 1] 
         * B[(k + 1) * params->b_cols + gid.y];
      sum += A[gid.x * params->a_cols + k + 2] 
         * B[(k + 2) * params->b_cols + gid.y];
      sum += A[gid.x * params->a_cols + k + 3] 
         * B[(k + 3) * params->b_cols + gid.y];
    }
  
    // Handle any remaining elements
    for (; k < params->a_cols; ++k) {
      sum += A[gid.x * params->a_cols + k] * B[k * params->b_cols + gid.y];
    }
  
    C[gid.x * params->b_cols + gid.y] = sum;
  }

  I implemented a as a pretty easy way to improve the scalar performance of the naive algorithm, at least, on CPUs. More on that later.

  Objective-C Bindings

  The Metal framework provides the ability to , and initializes a new MTLComputePipelineState representing the compiled function code.

  id<MTLDevice> device = MTLCreateSystemDefaultDevice();
  
  // Compile and initialize a new library located at the provided source path.
  MTLCompileOptions *compileOptions = [MTLCompileOptions new];
  compileOptions.languageVersion = MTLLanguageVersion3_0;
  
  // Wrap input source path string
  NSString *ss = [NSString stringWithUTF8String:source_path];
  
  // Initialize new library containing compiled shader functions
  id<MTLLibrary> lib = [device newLibraryWithSource:ss
    options:compileOptions
    error:&error];
  
  // Create a representation of the naive multiplication public shader function in 
  // the Metal library created above
  id<MTLFunction> naiveFunction =
      [lib newFunctionWithName:@"matrix_multiply_naive"];
  
  // Create the new compute pipeline state
  id<MTLComputePipelineState> pipelineStateNaive = [device newComputePipelineStateWithFunction:naiveFunction
    error:&error];

  To actually call the native Metal code, the thread configuration is a high-performance library provided by Apple for use with its .

  APIs are primarily available through Swift or Objective-C, though there is also a is is a Go language feature which allows the Go compiler to understand compiler directives contained within comments related to native C code. It supports a version of when we want to pass the contents of the MSL source file (mm.metal) to the Metal framework, as discussed above, for compilation.

  //go:embed mm.metal
  var source string
  
  // Compile the shader source code and initialize pipelines. The metalSource
  // param contains the contents of an embedded Metal Shading Language file.
  func Compile (metalSource string) {
   // Wrap string in a C string
   src := C.CString(metalSource)
  
   // Free the above string after command queue is initialized
   defer C.free(unsafe.Pointer(src))
  
   // Compile the source, initialize pipelines and command queue
   C.initializePipelineAndCommandQueue(src)
  }

  The references to C above are interfacing with C APIs through cgo, for example:

  // Calls initializeMTLBuffers from Obj-C bindings
  C.initializeMTLBuffers(
   a_data,                  // Input opaque pointer for A
   b_data,                  // Input opaque pointer for B
   C.int(4),                // Converts 4 into C integer type
   C.int(a.Size()),         
   C.int(b.Size()),         
   C.int(a.Rows * b.Cols))
  
  params := MatrixParams{
   a_rows: int32(a.Rows),
   a_cols: int32(a.Cols),
   b_rows: int32(b.Rows),
   b_cols: int32(b.Cols),
  }
  
  // Return an unsafe pointer to this MatrixParams struct, cast to 
  // the native C representation defined in the shared header file
  return (*C.MatrixParams)(unsafe.Pointer(&params));

  Note that this means that C is a reserved keyword and cannot be used as a variable name.

  Go Implementation Baselines and OpenBLAS

  I wanted to compare the performance of GPU-based matrix multiplication to both higher-level implementations, such as the . However, it can also be configured to divert algebraic operations to a native-code BLAS implementation like .

  I showed above how cgo can be configured to properly link to an OpenBLAS installation on MacOS. Within the application code, the preferred BLAS implementation can be set directly. From the benchmark code:

  // Convert primitive arrays into gonum dense matrix types
  gonum_a := mat.NewDense(a_rows, a_cols, a64_data)
  gonum_b := mat.NewDense(b_rows, b_cols, b64_data)
  gonum_c := mat.NewDense(a_rows, b_cols, nil)
  gonum_d := mat.NewDense(a_rows, b_cols, nil)
  
  // Configure Gonum to use Gonum-default Go implementation
  blas64.Use(gonum.Implementation{})
  
  // Run a multiplication using Gonum BLAS impl
  start = time.Now()
  gonum_c.Mul(gonum_a, gonum_b)
  bdata.TimeGonumNative(start)
  
  // Configure Gonum to use Netlib which forwards operations to a 
  // native C-code BLAS implementation (OpenBLAS in our case)
  blas64.Use(netlib.Implementation{})
  
  // Run a multiplication using OpenBLAS impl through Gonum API
  start = time.Now()
  gonum_d.Mul(gonum_a, gonum_b)
  bdata.TimeGonumOpenBLAS(start)

  Results

  My through matplotlib

  

  As one might expect, my hand-written Go implementations are comparatively out of control. In fact, the other approaches are so fast, you can’t even tell them apart in the graph. Here’s the sliding histogram of GPU usage during this run

  

  You can see the GPU isn’t particularly busy, because time is mostly being spent on CPU operations. Here’s another run, excluding the slowest three multiplication techniques:

  

  Around 16M elements (4k x 4k), Gonum begins to degrade. You can see clearly here that the GPU-based and OpenBLAS operations outperform the pure Go implementations. Looking just at the GPU-based approaches:

  

  A couple interesting notes here:

  • The Metal Performance Shaders library is amazingly fast
  • There’s no real performance difference between the naive and transposed naive approaches

  For the second point: this is unlike the performance characteristics of the Go-based pair of implementations above. It turns out that favorable cache access patterns for CPUs do not work the same way for GPUs and the way their SIMD groups ( was originally published in Towards Data Science on Medium, where people are continuing the conversation by highlighting and responding to this story.

  Vollständiger Original-Bericht

  Ausführliche Details, Code-Beispiele & Hersteller-Stellungnahme auf towardsdatascience.com.

  ↗ Original-Artikel auf towardsdatascience.com lesen