Feature Request: Support for Metal 4 Tensor APIs for Apple M5 chip

[wanderingmeow]

Summary

Metal 4 introduces the tensor resource and the Metal Performance Primitives (MPP) framework, enabling efficient development of machine learning kernels that leverage GPU neural accelerators in the Apple M5 chip. Leveraging these APIs in ds4 can significantly boost inference performance and efficiency on this Apple silicon.

Key Details

• Metal Tensor APIs:
  • The mpp::tensor_ops APIs provide compute primitives for threadgroup- and simdgroup-scoped operations on tensors, making it possible to write optimized matrix and tensor kernels.
  • Tensors encapsulate a base storage pointer, rank, data type, dimensions, and strides, supporting slicing and static/dynamic extents for high performance.
• GEMM Kernel Optimizations:
  • The MPP guide discusses best practices for implementing generalized matrix multiplication (GEMM) using hierarchical tiling (simdgroup, threadgroup, and grid scopes).
  • Tuning threadgroup and simdgroup tile sizes can maximize data reuse and occupancy, and locality-preserving walk orders (e.g., Morton ordering) can increase cache locality.
  • The Apple M5 chip's architecture allows optimal GEMM performance without the need for threadgroup memory staging or explicit software pipelining, relying instead on hardware-managed caching and sufficient thread occupancy.
• Advanced Kernel Applications:
  • New constructs like cooperative tensors and postfix fusion allow for additional fused operations with minimal memory traffic.
  • Theoretical performance models and optimal tiling strategies.

References

• Apple Metal Performance Primitives (MPP) Programming Guide (2026-03-16)