metal-kernel

Metal Kernel 編寫指南

本 Skill 將引導您在 Apple Silicon 上為 PyTorch 算子實作 Metal kernel。

**重要事項：**本 Skill 的目標是透過 c10/metal/ 基礎架構直接使用原生 Metal 功能，而非 MPSGraph。原生 Metal kernel 能提供更好的掌控度、效能與可維護性。

概覽

本 Skill 涵蓋兩種工作流程：

1. 新增 MPS 支援 — 從頭實作全新算子
2. 從 MPSGraph 遷移 — 將現有基於 MPSGraph 的算子轉為原生 Metal

兩種流程皆包含以下步驟：

1. 更新分派設定（dispatch）：位於 aten/src/ATen/native/native_functions.yaml
2. 編寫 Metal kernel：位於 aten/src/ATen/native/mps/kernels/
3. 實作主機端 stub：位於 aten/src/ATen/native/mps/operations/

步驟 1：更新 native_functions.yaml

檔案路徑： aten/src/ATen/native/native_functions.yaml

新增算子

找到對應的算子項目並加入 MPS dispatch：

# 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 時，需合併分派（dispatch）項目：

# 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 替換為在共享的 dispatch 行中加上 MPS（例如 CPU, CUDA, MPS: my_op_out）。

請務必更新每一個多載（overload）。 單一算子在 native_functions.yaml 中通常會有數個項目——包含 functional、inplace、.out 版本，以及 Tensor 與 Scalar 變體。每個項目都有各自的 dispatch: 區塊，且每個區塊都必須轉移。若有任何項目仍指向 MPS: my_op_mps，該多載就依然會路由至舊有 MPSGraph 程式碼，導致呼叫端依據所觸發的多載種類而默默走到舊路徑。在宣告遷移完成前，請 grep 搜尋舊版函式名稱，確認沒有任何項目仍引用它。

Dispatch 命名規範：

• MPS: function_name_mps — 專用於 MPS 的實作（舊版 MPSGraph 模式）
• CPU, CUDA, MPS: function_name — 共享 stub 的實作（原生 Metal 模式）

步驟 2：實作 Metal Kernel

檔案路徑： aten/src/ATen/native/mps/kernels/

一元算子（Unary Kernel）模式

// 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);

二元算子（Binary Kernel）模式

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);

二元算子型態註冊巨集（Registration Macros）

針對二元運算，請使用 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);

帶有標量（Scalar）參數

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);

型態特化（Type-Specialized）Functor

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

可在 functor 中使用 is_complex_v<T> 針對複數型態進行特化。

可用的 c10/metal 工具函式

utils.h：

• opmath_t<T> — 運算用數學型態（half -> float）
• accum_t<T> — 用於規約（reduction）的累加型態
• 帶有 NaN 傳播處理的 max(), min()

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：實作主機端 Stub

檔案路徑： aten/src/ATen/native/mps/operations/

依據運算類型選擇或建立合適的檔案：

• UnaryKernel.mm — 透過 stub 分派的單一輸入運算
• BinaryKernel.mm — 透過 stub 分派的雙輸入運算
• UnaryOps.mm / BinaryOps.mm — 舊版 MPSGraph 實作（供參考）
• ReduceOps.mm — 規約運算（sum、mean、max 等）
• 若為獨立的運算類別，可建立新檔案

Stub 註冊模式（原生 Metal 推薦方式）

適用於使用 TensorIterator 模式的結構化 kernel：

// 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）：

   • 新增 static kernel 函式
   • 新增 REGISTER_DISPATCH 呼叫

步驟 4：編譯

完成變更後，進行編譯以確認建置成功：

cd build && ninja torch_cpu

測試

基本的算子支援已透過 test/test_mps.py 中的 test_output_match 進行測試。實作完成算子後，請移除預期失敗的設定以啟用測試：

1. 從 common_mps.py 中移除

檔案路徑： torch/testing/_internal/common_mps.py

找到並將算子從 skip / xfail 列表中刪除：

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

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

2. 從 OpInfo 修飾器（decorator）中移除

檔案路徑： 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 Kernel

使用 torch.mps.compile_shader 可獨立對單一 Metal kernel 進行 JIT 即時編譯與測試。當你需要獨立驗證多個 kernel 管道（pipeline）中的各個階段時，這項工具極具價值。

基本用法

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')

分派語意（Dispatch Semantics）

compile_shader 採用 dispatchThreads 語意（與 PyTorch 中的 mtl_dispatch1DJob 相同）：

• threads=[N, 1, 1] — 總執行緒數量（非 threadgroup 數量）
• group_size=[G, 1, 1] — 每個 threadgroup 的執行緒數量

這與部分主機端程式碼使用的 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])

常數緩衝區（Constant Buffer）參數

將純量（scalar）常數作為單一元素的 tensor 傳入：

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])

多 Kernel 管道的除錯策略

When a pipeline of kernels (e.g., histogram → prefix_sum → scatter) pro

<!-- truncated for translation batch; full body continues in source -->