metal-kernel

Implementing Metal kernels for PyTorch operators on Apple Silicon is the focus, adding MPS (Metal Performance Shaders) device support through PyTorch's native c10/metal infrastructure rather than MPSGraph, which the guide favors for control, performance, and maintainability. Two workflows are covered: standing up brand-new MPS support for an operator, and migrating an existing MPSGraph-based operator to native Metal.

Both follow the same three stages. Dispatch is first declared in aten/src/ATen/native/native_functions.yaml, either as an MPS-specific entry or, preferably, as a shared line routing CPU, CUDA, and MPS to a single stub. When migrating, the guidance stresses updating every overload (functional, inplace, and .out, plus Tensor and Scalar variants) and grepping the legacy function name to confirm nothing still routes to the old MPSGraph path. Next, the Metal kernel is written under aten/src/ATen/native/mps/kernels/ using functor structs, and finally a host-side stub is added under aten/src/ATen/native/mps/operations/.

Concrete patterns are given for unary, binary, and scalar-parameter (alpha) operations, plus type-specialized functors that branch across floating-point, integral, and complex types. Registration relies on macros like REGISTER_UNARY_OP and REGISTER_BINARY_OP together with convenience variants such as REGISTER_FLOAT_BINARY_OP, REGISTER_INT2FLOAT_BINARY_OP, and REGISTER_INTEGER_BINARY_OP that fan out across supported dtypes. Complex numbers map to vector types (float2 and half2), and utilities from utils.h, special_math.h, and indexing.h, including opmath_t, precise math functions, and NaN-propagating max and min, round out the toolkit for contributors extending PyTorch's Apple Silicon backend.