metal-kernel

Metal 内核编写指南

本技能将指导你为 Apple Silicon 上的 PyTorch 运算符实现 Metal 内核。

重要： 本技能旨在通过 c10/metal/ 基础设施使用原生 Metal 功能，而非 MPSGraph。原生 Metal 内核在控制能力、性能和可维护性方面更出色。

概览

本技能涵盖两种工作流：

1. 添加新的 MPS 支持 - 从头实现新的运算符
2. 从 MPSGraph 迁移 - 将现有的基于 MPSGraph 的运算符转换为原生 Metal

这两种工作流都包括：

1. 更新分发配置：在 aten/src/ATen/native/native_functions.yaml 中操作
2. 编写 Metal 内核：在 aten/src/ATen/native/mps/kernels/ 中操作
3. 实现主机端桩函数：在 aten/src/ATen/native/mps/operations/ 中操作

第 1 步：更新 native_functions.yaml

位置： aten/src/ATen/native/native_functions.yaml

新运算符

找到运算符条目并添加 MPS 分发：

# Simple MPS-specific implementation
- func: my_op(Tensor self) -> Tensor
  dispatch:
    CPU: my_op_cpu
    CUDA: my_op_cuda
    MPS: my_op_mps

# Shared implementation across devices (preferred for structured kernels)
- func: my_op.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
  dispatch:
    CPU, CUDA, MPS: my_op_out

# Structured kernel (preferred for new ops)
- func: my_op.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
  structured: True
  structured_inherits: TensorIteratorBase
  dispatch:
    CPU, CUDA, MPS: my_op_out

从 MPSGraph 迁移

将现有运算符从 MPSGraph 迁移到原生 Metal 时，合并分发条目：

# BEFORE (MPSGraph-based, separate dispatch)
- func: atan2.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
  structured: True
  structured_inherits: TensorIteratorBase
  dispatch:
    CPU, CUDA: atan2_out
    MPS: atan2_out_mps  # Separate MPS implementation

# AFTER (native Metal, shared dispatch via stub)
- func: atan2.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
  structured: True
  structured_inherits: TensorIteratorBase
  dispatch:
    CPU, CUDA, MPS: atan2_out  # MPS now uses the same stub mechanism

主要变更： 将 MPS: my_op_out_mps 替换为在共享分发配置行中添加 MPS（例如 CPU, CUDA, MPS: my_op_out）。

更新每一个重载。 单个运算符通常在 native_functions.yaml 中有多个条目，包括函数式 / 原地 / .out 形式，以及 Tensor 和 Scalar 变体。每个条目都有自己的 dispatch: 块， 并且都必须迁移。任何仍指向 MPS: my_op_mps 的条目 都会让对应重载继续使用 MPSGraph 代码，因此调用方可能会因命中的重载不同而悄然 走上旧路径。宣布迁移完成前，请搜索旧函数名，确认没有条目仍在引用它。

分发命名约定：

• MPS: function_name_mps - MPS 专用实现（旧的 MPSGraph 模式）
• CPU, CUDA, MPS: function_name - 共享桩函数实现（原生 Metal 模式）

第 2 步：实现 Metal 内核

位置： aten/src/ATen/native/mps/kernels/

一元内核模式

// MyKernel.metal
#include <c10/metal/indexing.h>
#include <c10/metal/utils.h>
#include <metal_stdlib>

using namespace metal;
using namespace c10::metal;

// Define operation functor
struct my_op_functor {
  template <typename T>
  inline T operator()(const T x) {
    return /* your operation */;
  }
};

// Register for supported types
REGISTER_UNARY_OP(my_op, float, float);
REGISTER_UNARY_OP(my_op, half, half);
REGISTER_UNARY_OP(my_op, bfloat, bfloat);

二元内核模式

struct my_binary_functor {
  template <typename T>
  inline T operator()(const T a, const T b) {
    return /* your operation */;
  }
};

REGISTER_BINARY_OP(my_binary, float, float);
REGISTER_BINARY_OP(my_binary, half, half);

二元内核类型注册宏

对于二元运算，请使用 BinaryKernel.metal 中定义的便捷宏：

// Floating-point types only (float, half, bfloat)
REGISTER_FLOAT_BINARY_OP(my_op);

// Integral types with float output (for math ops like atan2, copysign)
// Registers: long->float, int->float, short->float, uchar->float, char->float, bool->float
REGISTER_INT2FLOAT_BINARY_OP(my_op);

// Integral types with same-type output (for bitwise/logical ops)
// Registers: long, int, short, uchar, char, bool
REGISTER_INTEGER_BINARY_OP(my_op);

// Floating-point with opmath precision (for ops needing higher precision)
REGISTER_OPMATH_FLOAT_BINARY_OP(my_op);

常见模式：

• 数学函数（atan2、copysign、logaddexp）：同时使用 REGISTER_FLOAT_BINARY_OP 和 REGISTER_INT2FLOAT_BINARY_OP
• 比较 / 逻辑运算（maximum、minimum）：同时使用 REGISTER_FLOAT_BINARY_OP 和 REGISTER_INTEGER_BINARY_OP
• 算术运算（add、sub、mul）：同时使用 REGISTER_FLOAT_BINARY_OP 和 REGISTER_INTEGER_BINARY_OP

atan2 示例（同时支持浮点和整数输入）：

struct atan2_functor {
  template <typename T, enable_if_t<is_floating_point_v<T>, bool> = true>
  inline T operator()(const T a, const T b) {
    return static_cast<T>(precise::atan2(float(a), float(b)));
  }
  template <typename T, enable_if_t<is_integral_v<T>, bool> = true>
  inline float operator()(const T a, const T b) {
    return precise::atan2(float(a), float(b));
  }
};

REGISTER_FLOAT_BINARY_OP(atan2);
REGISTER_INT2FLOAT_BINARY_OP(atan2);

带标量参数

struct my_alpha_functor {
  template <typename T>
  inline T operator()(const T a, const T b, const T alpha) {
    return a + c10::metal::mul(alpha, b);
  }
};

REGISTER_UNARY_ALPHA_OP(my_alpha, float, float, float);
REGISTER_UNARY_ALPHA_OP(my_alpha, half, half, half);

类型专用函子

struct special_functor {
  // Floating point types
  template <typename T, enable_if_t<is_scalar_floating_point_v<T>, bool> = true>
  inline T operator()(const T x) {
    return precise::exp(x);  // Use precise math
  }

  // Integral types
  template <typename T, enable_if_t<is_scalar_integral_v<T>, bool> = true>
  inline float operator()(const T x) {
    return precise::exp(float(x));
  }

  // Complex types (float2 for cfloat, half2 for chalf)
  template <typename T, enable_if_t<is_complex_v<T>, bool> = true>
  inline T operator()(const T x) {
    // x.x = real, x.y = imaginary
    return T(/* real */, /* imag */);
  }
};

关于复数类型的说明： Metal 中的复数以向量类型表示：

• c10::complex<float> 映射为 float2（x = 实部，y = 虚部）
• c10::complex<half> 映射为 half2

在函子中使用 is_complex_v<T> 为复数类型提供专门处理。

可用的 c10/metal 工具

utils.h：

• opmath_t<T> - 运算数学类型（half->float）
• accum_t<T> - 归约的累加类型
• max()、min()，支持传播 NaN

special_math.h：

• precise::exp()、precise::log()、precise::sqrt()
• precise::sin()、precise::cos()、precise::tan()
• erf()、erfc()、erfinv()

indexing.h：

• REGISTER_UNARY_OP(name, in_type, out_type)
• REGISTER_BINARY_OP(name, in_type, out_type)
• REGISTER_UNARY_ALPHA_OP(name, in_type, alpha_type, out_type)

第 3 步：实现主机端存根

位置： aten/src/ATen/native/mps/operations/

根据操作类型选择或创建合适的文件：

• UnaryKernel.mm — 通过存根分发处理单输入操作
• BinaryKernel.mm — 通过存根分发处理双输入操作
• UnaryOps.mm / BinaryOps.mm — 旧版 MPSGraph 实现（供参考）
• ReduceOps.mm — 归约操作（sum、mean、max 等）
• 针对不同的操作类别创建新文件

存根注册模式（原生 Metal 优先）

对于使用 TensorIterator 模式的结构化内核：

// In BinaryKernel.mm (or appropriate file)

static void my_op_mps_kernel(TensorIteratorBase& iter) {
  lib.exec_binary_kernel(iter, "my_op");  // "my_op" matches the functor name in .metal
}

// Register the MPS stub - this connects to the dispatch system
REGISTER_DISPATCH(my_op_stub, &my_op_mps_kernel)

对于一元操作：

static void my_unary_mps_kernel(TensorIteratorBase& iter) {
  lib.exec_unary_kernel(iter, "my_unary");
}

REGISTER_DISPATCH(my_unary_stub, &my_unary_mps_kernel)

迁移：移除旧版 MPSGraph 实现

从 MPSGraph 迁移时，也要移除旧实现：

1. 从 BinaryOps.mm（或 UnaryOps.mm）中移除：

   • 删除 TORCH_IMPL_FUNC(my_op_out_mps) 实现
   • 移除对应的 #include <ATen/ops/my_op_native.h> 头文件

2. 添加到 BinaryKernel.mm（或 UnaryKernel.mm）：

   • 添加静态内核函数
   • 添加 REGISTER_DISPATCH 调用

布尔内核参数

• 布尔类型使用 bool。 Metal 的 bool 占 1 字节，按 1 字节对齐， 与主机端匹配（bool2 为 2/2；bool3 和 bool4 均为 4/4）。切勿将标志 扩宽为 uint32_t，以便内核检查 != 0u。一对标志可打包到 constant bool2&， 并由 std::array<bool, 2> 提供。

• 将共享参数结构体中的 bool 成员放在最后。 kernels/<Op>.h 会由 Metal 和 主机编译器共同编译；将 bool 塞在较宽成员之间会产生填充，也会让后续字段 的起始位置难以确定。结构体调整后，要同步调整主机端聚合初始化器—— 花括号初始化按位置对应。

第 4 步：编译

完成更改后，编译以确认所有内容都能正确构建：

cd build && ninja torch_cpu

测试

test/test_mps.py 中的 test_output_match 已测试基本算子支持。实现算子后，移除预期失败项以启用测试：

1. 从 common_mps.py 中移除

位置： torch/testing/_internal/common_mps.py

找到并从跳过/预期失败列表中移除该算子：

# Remove entries like:
MPS_XFAILLIST = {
    "my_op": ...,  # Remove this line
}

MPS_SKIPLIST = {
    "my_op": ...,  # Remove this line
}

2. 从 OpInfo 装饰器中移除

位置： torch/testing/_internal/common_methods_invocations.py（或相关文件）

从 OpInfo 中移除 MPS 专用装饰器：

OpInfo(
    "my_op",
    # Remove decorators like:
    # decorators=[skipMPS, expectedFailureMPS("reason")],
    ...
)

3. 运行测试进行验证

# Run the specific operator test
python test/test_mps.py -k test_output_match_my_op

# Or run full MPS test suite
python test/test_mps.py

使用 torch.mps.compile_shader 调试 Metal 内核

使用 torch.mps.compile_shader 对单个 Metal 内核进行 JIT 编译并独立测试。这对于调试多内核流水线非常有用，可逐个验证各个阶段。

基本用法

import torch

source = '''
#include <metal_stdlib>
using namespace metal;

kernel void my_kernel(
    const device float* input [[buffer(0)]],
    device float* output [[buffer(1)]],
    uint tid [[thread_position_in_grid]]) {
  output[tid] = input[tid] * 2.0;
}
'''

lib = torch.mps.compile_shader(source)

inp = torch.tensor([1.0, 2.0, 3.0], device='mps')
out = torch.zeros(3, device='mps')
lib.my_kernel(inp, out, threads=[3, 1, 1], group_size=[3, 1, 1])
torch.mps.synchronize()
print(out)  # tensor([2., 4., 6.], device='mps:0')

分发语义

compile_shader 使用 dispatchThreads 语义（与 PyTorch 中的 mtl_dispatch1DJob 相同）：

• threads=[N, 1, 1] — 线程总数（不是线程组数）
• group_size=[G, 1, 1] — 每个线程组中的线程数

这与某些主机端代码使用的 dispatchThreadgroups API 不同。要匹配 dispatchThreadgroups:MTLSizeMake(num_tgs, num_slices, 1) threadsPerThreadgroup:MTLSizeMake(TG_SIZE, 1, 1)：

# Equivalent compile_shader call:
lib.kernel(args...,
    threads=[num_tgs * TG_SIZE, num_slices, 1],
    group_size=[TG_SIZE, 1, 1])

常量缓冲区参数

将标量常量作为单元素张量传入：

slice_size = torch.tensor([1024], dtype=torch.int32, device='mps')
lib.my_kernel(data, output, slice_size, threads=[1024, 1, 1], group_size=[256, 1, 1])

多内核流水线的调试策略

如果内核流水线（例如 histogram → prefix_sum → scatter）产生错误结果，应逐个测试各个内核，并将输出与 Python/NumPy 参考结果进行验证：

# 1. Run GPU kernel
lib.histogram(keys, hist, ..., threads=[N, 1, 1], group_size=[256, 1, 1])
torch.mps.synchronize()

# 2. Compute reference in Python
ref_hist = compute_histogram_cpu(keys.cpu().numpy(), ...)

# 3. Compare
assert np.array_equal(hist.cpu().numpy(), ref_hist), "Histogram mismatch!"

这样可以定位流水线中出错的内核，而不必一次调试整个流水线。

常见问题

• threads 数量错误 — threads 表示线程总数，而不是线程组数。对于 5 个线程组、每组 256 个线程，应使用 threads=[1280, 1, 1]。
• 线程组内存 — compile_shader 不直接支持 [[threadgroup(N)]] 参数。如果内核需要线程组内存，请调整结构，改用在内核主体内声明的 threadgroup 数组。

使用 TensorIterators

REGISTER_UNARY_OP / REGISTER_BINARY_OP 封装了迭代器处理逻辑。带有额外参数或非逐元素布局的内核必须直接驱动 TensorIterator。以下规则不太直观：

• 传入 TensorIteratorBase&，而不是 Tensor&。 Tensor& 会丢失 with_32bit_indexing() 生成的偏移量和形状信息。当 stub 传给你 const TensorBase&（例如 bernoulli_scalar_stub）时，应通过 at::TensorIterator::borrowing_nullary_op(self) 内联构建 iter，而不是去掉 const 限定。

• 对于子 iter，iter.tensor(0) 返回的是整个张量。 调用 with_32bit_indexing() 后，子 iter 仍引用原始存储，因此直接绑定 iter.tensor(0) 会导致每个子 iter 都覆盖同一段前缀，并使尾部未初始化。请使用 bind_iter_tensors，它会计算每个分块的偏移量，并将缓冲区 0 绑定到对应切片。分发时使用 iter.numel()，绝不要使用 iter.tensor(0).numel()：

  bind_iter_tensors(computeEncoder, iter, /*ntensors=*/1);
  mtl_setArgs<1>(computeEncoder, params, ..., numel);
  mtl_dispatch1DJob(computeEncoder, pso, threads);
  

• 优先使用 mtl_setArgs<N>，而不是链式调用 mtl_setBytes，并将类型相同的标量打包到一个向量参数中。 mtl_setArgs<1>(encoder, a, b, c) 会按照与宏相同的重载解析方式，将 a / b / c 绑定到槽位 1 / 2 / 3，因此 std::array<T, N> 会作为单个 constant T2& / T3& 绑定传入（std::array<long,2> → constant long2&），而不是传入 N 个参数。两个整数会作为一个 int2，而不是占用两个缓冲区槽位。

• 使用 ceil_div。 Host 端：使用 <ATen/ceil_div.h> 中的 at::ceil_div。Metal 端：使用 <c10/metal/common.h> 中的 c10::metal::ceil_div（在 using namespace c10::metal; 后可不加限定名）。两者都封装了 (a + b - 1) / b。

• 使用 IF_CONSTEXPR 确保 Metal 3/4 兼容性。 Metal 4 支持 if constexpr；Metal 3 不支持。请使用 <c10/metal/common.h> 中的宏，例如 if IF_CONSTEXPR (sizeof(T) == 8) { ... }。

• 通过 REGISTER_MPS_DISPATCH 共享 CPU/CUDA stub。 当该算子在上游有 DECLARE_DISPATCH stub（如分布、融合算子）时，应使用 REGISTER_MPS_DISPATCH(stub_name, &fn) 接入 MPS，而不是在 native_functions.yaml 中添加仅供 MPS 使用的条目。该 stub 已接受 TensorIteratorBase&。

处理大型张量

大多数 Metal 内核将 numel 设为 uint32_t 并使用 32 位索引，因此对于超过 INT32_MAX 的张量，需要在 host 端拆分。应通过 TensorIterator 处理，而不是手动切片。

使用 iter.with_32bit_indexing() 拆分：

if (!iter.can_use_32bit_indexing()) {
  for (auto&& sub_iter : iter.with_32bit_indexing()) {
    my_kernel_impl(sub_iter, ...);
  }
  return;
}

每个生成的子 iter 都满足 can_use_32bit_indexing()，因此递归会在一层后终止。阈值是 INT32_MAX（TensorIterator 使用有符号 32 位偏移量），而不是 UINT32_MAX——刚超过 INT32_MAX 的测试会触发拆分，但不会触发 uint32_t 转换；只有 numel > 2^32 时转换本身才会产生影响。

缩窄类型时使用经过检查的转换：

const uint32_t numel = c10::checked_convert<uint32_t>(iter.numel(), "uint32_t");

请使用 c10::checked_convert（来自 <c10/util/TypeCast.h>），这样溢出回绕会变成 TORCH_CHECK 错误，而不会造成输出损坏。

从内核报告错误

绝不要为了检查错误而将结果复制到 CPU。 对 GPU 张量执行任何 .item()、.cpu() 或其他主机端读取，都会强制执行完整的 GPU->CPU 同步，排空流中所有正在执行的操作，而不仅仅是你正在检查的归约操作。在实际流水线中，这会让整个队列停滞，代价远高于检查本身。将同步放在 is_mps() 判断后面也无济于事；每次运行该操作都会发生停顿。应改为在设备上验证，并通过以下机制报告错误。

GPU 代码无法抛出异常，但内核可以将错误写入共享错误缓冲区，主机随后会将其作为 c10::AcceleratorError 抛出。请使用 <c10/metal/error.h> 中的 TORCH_REPORT_ERROR(error_buf, ...)。可变参数会拼接在一起；整数会以十进制格式化。错误会异步传递——发生错误的线程会继续运行；当 MPSStream::checkLastError() 下次运行时（在 synchronize() 之后，或下一个排空流的操作中），错误才会显现，因此不要依赖它来控制同一次派发中的流程。关键是，这不会强制同步：错误会搭载在用户代码本来就会执行的同步上。

**内核侧：**接收一个 device ErrorMessages* error_buf 参数，并在出错路径上调用 TORCH_REPORT_ERROR。之后跳过有问题的元素，以免同时破坏内存：

#include <c10/metal/error.h>

kernel void index_set_1d(
    device float* self,
    constant float* values,
    constant long* indices,
    constant uint& self_numel,
    device ::c10::metal::ErrorMessages* error_buf,
    uint tid [[thread_position_in_grid]]) {
  long idx = indices[tid];
  if (idx < 0 || idx >= long(self_numel)) {
    TORCH_REPORT_ERROR(
        error_buf, "index ", idx, " out of bounds for size ", long(self_numel));
    return;
  }
  self[idx] = values[tid];
}

**主机侧：**将流的错误缓冲区绑定为对应参数。mtl_setArgs 可直接接受该参数：

auto* stream = getCurrentMPSStream();
mtl_setArgs(encoder, self, values, indices, uint32_t(self.numel()),
            stream->getErrorBuffer());

该缓冲区由 MPSStream 所有，可容纳 30 条消息，并会在每次 checkLastError() 排空后重置。只会报告第一条消息；后续消息会保留以供调试，但 AcceleratorError 中包含的是 msg[0]。

检查清单

• 已在 native_functions.yaml 中添加 MPS 派发
• 已在 kernels/ 中实现 Metal 内核
• 已在 operations/ 中实现主机端运算符
• 标量参数已打包到向量参数中；布尔值类型为 bool，并放在共享结构体的最后
• 能处理空张量
• 能处理非连续张量
• 支持所需数据类型（float32、float16、bfloat16，以及通常通过 float2/half2 支持的复数类型）
• 已从 torch/testing/_internal/common_mps.py 中移除预期失败项
• 已从 OpInfo 中移除跳过/预期失败装饰器（如适用）