metal-kernel

Metal カーネルの記述ガイド

このスキルでは、Apple Silicon 上で PyTorch 演算子向けの Metal カーネルを実装する方法を説明します。

重要： このスキルの目標は、MPSGraph ではなく、c10/metal/ のインフラストラクチャを介してネイティブ Metal の機能を使用することです。ネイティブ Metal カーネルを使用すると、制御性、パフォーマンス、保守性が向上します。

概要

このスキルでは、次の 2 つのワークフローを扱います。

1. 新しい MPS サポートの追加 - 演算子を一から実装する
2. MPSGraph からの移行 - 既存の MPSGraph ベースの演算子をネイティブ Metal に変換する

どちらのワークフローにも、次の手順が含まれます。

1. aten/src/ATen/native/native_functions.yaml のディスパッチを更新する
2. aten/src/ATen/native/mps/kernels/ に Metal カーネルを記述する
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）に置き換えます。

すべてのオーバーロードを更新します。 1 つの演算子には通常、native_functions.yaml に複数のエントリがあります。functional、inplace、.out に加え、Tensor と Scalar のバリアントもあります。各エントリには独自の dispatch: ブロックがあり、すべて移行する必要があります。MPS: my_op_mps を参照するエントリが残っていると、そのオーバーロードは引き続き MPSGraph のコードにルーティングされるため、呼び出し側は使用したオーバーロードによって、気付かないうちに古い経路に進む可能性があります。移行完了を宣言する前に、レガシー関数名を grep し、それを参照するエントリが残っていないことを確認してください。

ディスパッチの命名規則：

• 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 の例（float と int の両方の入力をサポート）：

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> - リダクションの累積型
• 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：ホスト側スタブを実装する

場所： 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）。カーネルで != 0u と判定できるようにするため、フラグを uint32_t に拡張しないでください。2 つのフラグは 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

スキップ／xfail リストから演算子を探して削除します。

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

定数バッファーパラメーター

スカラー定数は、要素数 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 はスレッドグループ数ではなく、スレッドの総数です。256 スレッドのスレッドグループを 5 つ使用する場合は、threads=[1280, 1, 1] を指定します。
• スレッドグループメモリ — compile_shader は [[threadgroup(N)]] パラメーターを直接サポートしていません。カーネルにスレッドグループメモリが必要な場合は、カーネル本体の中で宣言した threadgroup 配列を使うように構成を変更します。

TensorIterator の使用

REGISTER_UNARY_OP／REGISTER_BINARY_OP はイテレーターの処理を隠蔽します。追加パラメーターや要素ごとではないレイアウトを扱うカーネルでは、TensorIterator を直接操作する必要があります。注意が必要なルールは次のとおりです。

• Tensor& ではなく TensorIteratorBase& を渡します。 Tensor& では、with_32bit_indexing() が生成するオフセット／形状情報が失われます。スタブから const TensorBase&（例：bernoulli_scalar_stub）が渡される場合は、const_cast するのではなく、at::TensorIterator::borrowing_nullary_op(self) でイテレーターをインラインで構築します。

• サブイテレーターでも iter.tensor(0) は テンソル全体 を返します。 with_32bit_indexing() の後も、サブイテレーターは元のストレージを参照するため、iter.tensor(0) をそのままバインドすると、すべてのサブイテレーターが同じ先頭部分を上書きし、末尾が未初期化のままになります。チャンクごとのオフセットを計算してバッファー 0 をスライスにバインドする bind_iter_tensors を使用します。ディスパッチには iter.tensor(0).numel() ではなく、必ず iter.numel() を使用します。

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

• 連鎖した mtl_setBytes より mtl_setArgs<N> を優先し、同じ型のスカラーは 1 つのベクター引数にまとめます。 mtl_setArgs<1>(encoder, a, b, c) は、マクロと同じオーバーロード解決で a／b／c をスロット 1／2／3 にバインドするため、std::array<T, N> は N 個のバインディングではなく、単一の constant T2&／T3& バインディングとして渡されます（std::array<long,2> → constant long2&）。2 つの int は 2 つのバッファースロットではなく、int2 です。

• ceil_div を使用します。 ホスト側では <ATen/ceil_div.h> の at::ceil_div を使用します。Metal 側では <c10/metal/common.h> の c10::metal::ceil_div を使用します（using namespace c10::metal; の後は名前空間修飾なしで呼び出せます）。どちらも (a + b - 1) / b をラップします。

• Metal 3／4 との互換性には IF_CONSTEXPR を使用します。 Metal 4 は if constexpr に対応していますが、Metal 3 は対応していません。<c10/metal/common.h> のマクロを使用します。例：if IF_CONSTEXPR (sizeof(T) == 8) { ... }。

• REGISTER_MPS_DISPATCH を使って CPU／CUDA スタブを共有します。 上流に DECLARE_DISPATCH スタブがある演算（分布、融合演算など）では、MPS 固有の native_functions.yaml エントリーを追加する代わりに、REGISTER_MPS_DISPATCH(stub_name, &fn) で MPS を登録します。このスタブはすでに TensorIteratorBase& を受け取ります。

大きなテンソルの処理

ほとんどの Metal カーネルは numel を uint32_t として受け取り、32 ビットでインデックスを計算します。そのため、INT32_MAX を超える場合はホスト側で分割する必要があります。手動でスライスするのではなく、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;
}

生成されるすべてのサブイテレーターは can_use_32bit_indexing() を満たすため、再帰は 1 段階で終了します。しきい値は UINT32_MAX ではなく INT32_MAX です（TensorIterator は符号付き 32 ビットオフセットを使用します）。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, ...) を使ってください。可変長引数は連結され、整数は 10 進数でフォーマットされます。通知は非同期です。エラーが発生したスレッドは実行を続け、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 から skip/xfail デコレーターを削除した