Apple Silicon hardware acceleration backends for GoMLX / github.com/gomlx/compute.
| Backend Name | Package | Framework | Hardware |
|---|---|---|---|
mps |
github.com/gomlx/go-darwinml/compute/mps |
Metal Performance Shaders Graph (MPSGraph) |
Apple GPU (Metal) |
mlx |
github.com/gomlx/go-darwinml/compute/mlx |
Apple MLX (mlx::core) |
Apple GPU (Metal) |
- macOS 13.3+ (macOS 14.0+ Sonoma recommended for
BFloat16, nativescaledDotProductAttention, and MLX) - Apple Silicon (M1 / M2 / M3 / M4 or later)
- Xcode Command Line Tools (
xcode-select --install) for CGo + Metal/MPSGraph/MLX - Go 1.27+
- For the
"mlx"backend only: Apple's MLX library installed on the system:(See the MLX installation guide for alternative installation methods. Thebrew install mlx
"mps"backend uses macOS system frameworks only and does not requirebrew install mlx.)
go get github.com/gomlx/go-darwinmlOn Darwin (macOS) machines, the "mps" backend can be imported with _ "github.com/gomlx/go-darwinml/compute/mps" (or via GoMLX's default backends package _ "github.com/gomlx/gomlx/backends/default"):
import (
"github.com/gomlx/compute"
_ "github.com/gomlx/go-darwinml/compute/mps"
)
func main() {
backend, err := compute.NewWithConfig("mps")
if err != nil {
panic(err)
}
defer backend.Finalize()
}Import _ "github.com/gomlx/go-darwinml/compute/mlx" and select "mlx" (or "mlx:dynamic" for native dynamic shapes without recompilation):
import (
"github.com/gomlx/compute"
_ "github.com/gomlx/go-darwinml/compute/mlx"
)
func main() {
// Default: static-shape compiled execution via mlx::core::compile ("mlx").
// Or use "mlx:dynamic" for native dynamic-shape execution without recompilation.
backend, err := compute.NewWithConfig("mlx")
if err != nil {
panic(err)
}
defer backend.Finalize()
}"mlx"(default): Compiles the computation graph withmlx::core::compilefor kernel fusion (DynamicShapesNone)."mlx:dynamic": Replays the lazy MLX op tape directly on runtime inputs (DynamicShapesNative,DynamicAxes: true), supporting dynamic dimensions (shapes.DynamicDim) without recompilation."mlx:dynamic,lazy_contiguous": In dynamic mode, defers materializing row-major contiguity on view outputs (Transpose,BroadcastInDim,Slice) until CPU access (Buffer.Data()/Buffer.ToFlatData()).
You can also select either backend at runtime via the GOMLX_BACKEND environment variable (e.g., GOMLX_BACKEND=mlx, GOMLX_BACKEND="mlx:dynamic", or GOMLX_BACKEND=mps) when calling compute.New() / compute.MustNew().
See docs/MPSGraph.md and docs/MLX.md for detailed architecture and buffer management documentation.
- Unsupported DTypes: Apple Silicon Metal GPUs do not support 64-bit floats (
Float64),Complex128, or sub-byte integer types (Int2,Int4,Uint2,Uint4). (Complex64is supported.) - Unsupported Operations:
- Control flow:
If(OpTypeIf) andWhile(OpTypeWhile) are not supported (Callis supported via tape inlining, andSortis supported for standard ascending/descending comparators). - Pooling scatter sum:
SelectAndScatterSum(SelectAndScatterMaxandSelectAndScatterMinare supported). - FFT, collective & quantized fused ops:
FFT,AllReduce,FusedQuantizedDense,QuantizedEmbeddingLookup.
- Control flow:
- Unsupported DTypes: Apple Metal /
MPSGraphdoes not support 64-bit floats (Float64), complex numbers (Complex64,Complex128), or sub-byte integer types (Int2,Int4,Uint2,Uint4). - Unsupported Operations:
- Complex & FFT ops:
Complex,Real,Imag,Conj,FFT. - Bit-level & bitwise reduction ops:
Bitcast,BitCount,Clz,ReduceBitwiseAnd,ReduceBitwiseOr,ReduceBitwiseXor,ReduceLogicalXor. - Pooling scatter variants:
SelectAndScatterSum,SelectAndScatterMin(SelectAndScatterMaxis supported). - Dynamic shape ops:
DynamicShapesNone. - Collective & quantized fused ops:
AllReduce,FusedQuantizedDense,QuantizedEmbeddingLookup. - Barriers:
OptimizationBarrier,SchedulingBarrier. - (Note: Operations with GoMLX graph decompositions, such as
CumSumandCumMax, automatically fall back to their decomposed implementations).
- Complex & FFT ops:
go-darwinml/
├── compute/
│ ├── mlx/ # GoMLX compute.Backend implementation ("mlx")
│ └── mps/ # GoMLX compute.Backend implementation ("mps")
├── docs/
│ ├── MLX.md # MLX backend architecture & design notes
│ └── MPSGraph.md # MPSGraph backend architecture & buffer pooling
└── internal/
├── bridge/ # CGo / Objective-C bridge to Metal & MPSGraph
└── mlxbridge/ # CGo / C++17 bridge to Apple MLX (mlx/mlx.h)
# Run unit and backend compliance tests
go test ./...Apache 2.0 - see LICENSE file.