Guide d'écriture de noyaux Metal

Cette skill vous guide dans la mise en œuvre de noyaux Metal pour les opérateurs PyTorch sur Apple Silicon.

Important : L'objectif de cette skill est d'utiliser les capacités natives de Metal via l'infrastructure c10/metal/, NON MPSGraph. Les noyaux Metal natifs offrent un meilleur contrôle, une meilleure performance et une meilleure maintenabilité.

Aperçu

Deux workflows sont couverts par cette skill :

1. Ajouter un nouveau support MPS - Implémenter un nouvel opérateur à partir de zéro
2. Migrer depuis MPSGraph - Convertir les opérateurs existants basés sur MPSGraph en Metal natif

Les deux workflows impliquent :

1. Mettre à jour la dispatch dans aten/src/ATen/native/native_functions.yaml
2. Écrire le noyau Metal dans aten/src/ATen/native/mps/kernels/
3. Implémenter le stub côté hôte dans aten/src/ATen/native/mps/operations/

Étape 1 : Mettre à jour native_functions.yaml

Emplacement : aten/src/ATen/native/native_functions.yaml

Pour les nouveaux opérateurs

Trouvez l'entrée de l'opérateur et ajoutez la dispatch MPS :

# Implémentation simple spécifique à MPS
- func: my_op(Tensor self) -> Tensor
  dispatch:
    CPU: my_op_cpu
    CUDA: my_op_cuda
    MPS: my_op_mps

# Implémentation partagée entre appareils (préférée pour les noyaux structurés)
- func: my_op.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
  dispatch:
    CPU, CUDA, MPS: my_op_out

# Noyau structuré (préféré pour les nouveaux opérateurs)
- func: my_op.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
  structured: True
  structured_inherits: TensorIteratorBase
  dispatch:
    CPU, CUDA, MPS: my_op_out

Pour la migration depuis MPSGraph

Lors de la migration d'un opérateur existant de MPSGraph vers Metal natif, consolidez l'entrée de dispatch :

# AVANT (basé sur MPSGraph, dispatch séparé)
- 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  # Implémentation MPS sépaée

# APRÈS (Metal natif, dispatch partagé 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 utilise désormais le même mécanisme de stub

Changement clé : Remplacez MPS: my_op_out_mps en ajoutant MPS à la ligne de dispatch partagée (par exemple, CPU, CUDA, MPS: my_op_out).

Mettez à jour chaque surcharge. Un opérateur unique a généralement plusieurs entrées native_functions.yaml — fonctionnelle / inplace / .out, plus des variantes Tensor et Scalar. Chaque entrée a son propre bloc dispatch:, et chacune doit être migrée. Toute entrée restant pointée vers MPS: my_op_mps route encore cette surcharge vers le code MPSGraph, donc les appelants peuvent silencieusement aboutir sur l'ancien chemin selon la surcharge qu'ils utilisent. Avant de déclarer la migration terminée, cherchez le nom de la fonction héritée et confirmez qu'aucune entrée ne la référence encore.

Conventions de nommage de dispatch :

• MPS: function_name_mps - Implémentation spécifique à MPS (ancien motif MPSGraph)
• CPU, CUDA, MPS: function_name - Implémentation de stub partagée (motif Metal natif)

Étape 2 : Implémenter le noyau Metal

Emplacement : aten/src/ATen/native/mps/kernels/

Motif de noyau unaire

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

using namespace metal;
using namespace c10::metal;

// Définir le foncteur d'opération
struct my_op_functor {
  template <typename T>
  inline T operator()(const T x) {
    return /* votre opération */;
  }
};

// Enregistrer pour les types supportés
REGISTER_UNARY_OP(my_op, float, float);
REGISTER_UNARY_OP(my_op, half, half);
REGISTER_UNARY_OP(my_op, bfloat, bfloat);

Motif de noyau binaire

struct my_binary_functor {
  template <typename T>
  inline T operator()(const T a, const T b) {
    return /* votre opération */;
  }
};

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

Macros d'enregistrement de type de noyau binaire

Pour les opérations binaires, utilisez les macros de commodité définies dans BinaryKernel.metal :

// Types en virgule flottante uniquement (float, half, bfloat)
REGISTER_FLOAT_BINARY_OP(my_op);

// Types intégraux avec sortie en virgule flottante (pour les opérations mathématiques comme atan2, copysign)
// Enregistre : long->float, int->float, short->float, uchar->float, char->float, bool->float
REGISTER_INT2FLOAT_BINARY_OP(my_op);

// Types intégraux avec sortie du même type (pour les opérations bitwise/logiques)
// Enregistre : long, int, short, uchar, char, bool
REGISTER_INTEGER_BINARY_OP(my_op);

// Virgule flottante avec précision opmath (pour les opérations nécessitant une précision supérieure)
REGISTER_OPMATH_FLOAT_BINARY_OP(my_op);

Motifs courants :

• Fonctions mathématiques (atan2, copysign, logaddexp) : Utilisez à la fois REGISTER_FLOAT_BINARY_OP et REGISTER_INT2FLOAT_BINARY_OP
• Opérations de comparaison/logiques (maximum, minimum) : Utilisez à la fois REGISTER_FLOAT_BINARY_OP et REGISTER_INTEGER_BINARY_OP
• Opérations arithmétiques (add, sub, mul) : Utilisez à la fois REGISTER_FLOAT_BINARY_OP et REGISTER_INTEGER_BINARY_OP

Exemple pour atan2 (supporte les entrées float et 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);

Avec paramètre scalaire

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

Foncteur avec spécialisation de type

struct special_functor {
  // Types en virgule flottante
  template <typename T, enable_if_t<is_scalar_floating_point_v<T>, bool> = true>
  inline T operator()(const T x) {
    return precise::exp(x);  // Utiliser les mathématiques précises
  }

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

  // Types complexes (float2 pour cfloat, half2 pour chalf)
  template <typename T, enable_if_t<is_complex_v<T>, bool> = true>
  inline T operator()(const T x) {
    // x.x = réel, x.y = imaginaire
    return T(/* réel */, /* imag */);
  }
};

Note sur les types complexes : Les nombres complexes en Metal sont représentés sous forme de types vecteur :

• c10::complex<float> se mappe à float2 (x = réel, y = imaginaire)
• c10::complex<half> se mappe à half2

Utilisez is_complex_v<T> pour spécialiser les foncteurs pour les types complexes.

Utilitaires c10/metal disponibles

utils.h :

• opmath_t<T> - Type mathématique d'opération (half->float)
• accum_t<T> - Type d'accumulation pour les réductions
• max(), min() avec propagation 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)

Étape 3 : Implémenter le stub côté hôte

Emplacement : aten/src/ATen/native/mps/operations/

Choisissez ou créez un fichier approprié selon le type d'opération :

• UnaryKernel.mm - Opérations à entrée unique via dispatch de stub
• BinaryKernel.mm - Opérations à deux entrées via dispatch de stub
• UnaryOps.mm / BinaryOps.mm - Implémentations MPSGraph héritées (pour référence)
• ReduceOps.mm - Réductions (sum, mean, max, etc.)
• Créer un nouveau fichier pour les catégories d'opérations distinctes

Motif d'enregistrement de stub (préféré pour Metal natif)

Pour les noyaux structurés qui utilisent le motif TensorIterator :

// Dans BinaryKernel.mm (ou fichier approprié)

static void my_op_mps_kernel(TensorIteratorBase& iter) {
  lib.exec_binary_kernel(iter, "my_op");  // "my_op" correspond au nom du foncteur dans .metal
}

// Enregistrer le stub MPS - cela se connecte au système de dispatch
REGISTER_DISPATCH(my_op_stub, &my_op_mps_kernel)

Pour les opérations unaires :

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

REGISTER_DISPATCH(my_unary_stub, &my_unary_mps_kernel)

Migration : Supprimer l'ancienne implémentation MPSGraph

Lors de la migration de MPSGraph, supprimez également l'ancienne implémentation :

1. Supprimer de BinaryOps.mm (ou UnaryOps.mm) :

   • Supprimer l'implémentation TORCH_IMPL_FUNC(my_op_out_mps)
   • Supprimer l'en-tête correspondant #include <ATen/ops/my_op_native.h>

2. Ajouter à BinaryKernel.mm (ou UnaryKernel.mm) :

   • Ajouter la fonction de noyau statique
   • Ajouter l'appel REGISTER_DISPATCH

Paramètres de noyau booléens

• Tapez les booléens en tant que bool. Le type bool de Metal est un octet, aligné sur un octet, correspondant à l'hôte (bool2 est 2/2 ; bool3 et bool4 sont tous deux 4/4). Ne dilatez jamais un drapeau à uint32_t pour que le noyau puisse tester != 0u. Une paire de drapeaux s'empile en constant bool2&, alimentée par std::array<bool, 2>.

• Placez les membres bool en dernier dans une structure de paramètres partagée. kernels/<Op>.h est compilé par le compilateur Metal et l'hôte ; un bool coincé entre des membres plus larges coûte de l'espace et suscite le doute sur le début du champ suivant. Réorganisez l'initialisateur d'agrégat hôte avec la structure -- l'initialisation entre accolades est positionnelle.

Étape 4 : Compiler

Après avoir apporté des modifications, compilez pour vérifier que tout se compile correctement :

cd build && ninja torch_cpu

Tests

Le support d'opérateur de base est déjà testé par test_output_match dans test/test_mps.py. Après avoir implémenté un opérateur, activez les tests en supprimant les défaillances attendues :

1. Supprimer de common_mps.py

Emplacement : torch/testing/_internal/common_mps.py

Trouvez et supprimez l'opérateur des listes de skip/xfail :

# Supprimer les entrées comme :
MPS_XFAILLIST = {
    "my_op": ...,  # Supprimer cette ligne
}

MPS_SKIPLIST = {
    "my_op": ...,  # Supprimer cette ligne
}

2. Supprimer des décorateurs OpInfo

Emplacement : torch/testing/_internal/common_methods_invocations.py (ou fichiers connexes)

Supprimez les décorateurs spécifiques à MPS de l'OpInfo :

OpInfo(
    "my_op",
    # Supprimer les décorateurs comme :
    # decorators=[skipMPS, expectedFailureMPS("reason")],
    ...
)

3. Exécuter les tests pour vérifier

# Exécuter le test de l'opérateur spécifique
python test/test_mps.py -k test_output_match_my_op

# Ou exécuter la suite de tests MPS complète
python test/test_mps.py

Déboguer les noyaux Metal avec torch.mps.compile_shader

Utilisez torch.mps.compile_shader pour compiler JIT et tester les noyaux Metal individuels en isolation. C'est inestimable pour déboguer les pipelines multi-noyaux où vous devez vérifier chaque étape indépendamment.

Utilisation de base

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

Sémantique de dispatch

compile_shader utilise la sémantique de dispatchThreads (identique à mtl_dispatch1DJob dans PyTorch) :

• threads=[N, 1, 1] — nombre total de threads (NON groupes de threads)
• group_size=[G, 1, 1] — threads par groupe de threads

Cela diffère de l'API dispatchThreadgroups utilisée par du code côté hôte. Pour correspondre à dispatchThreadgroups:MTLSizeMake(num_tgs, num_slices, 1) threadsPerThreadgroup:MTLSizeMake(TG_SIZE, 1, 1) :

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

Paramètres de buffer constant

Passez les constantes scalaires comme tenseurs d'un seul élément :

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

Stratégie de débogage pour les pipelines multi-noyaux

Lorsqu'un pipeline de noyaux (par exemple, histogram → prefix_sum → scatter) produit des résultats incorrects, testez chaque noyau individuellement et vérifiez sa sortie par rapport à une référence Python/NumPy :

# 1. Exécuter le noyau GPU
lib.histogram(keys, hist, ..., threads=[N, 1, 1], group_size=[256, 1, 1])
torch.mps.synchronize()

# 2. Calculer la référence en Python
ref_hist = compute_histogram_cpu(keys.cpu().numpy(), ...)

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

Cela isole quel noyau du pipeline est cassé, plutôt que de déboguer l'ensemble du pipeline à la fois.

Pièges courants

• Nombre de threads incorrect — threads est le nombre total de threads, pas les groupes de threads. Pour 5 groupes de threads de 256, utilisez threads=[1280, 1, 1].
• Mémoire du groupe de threads — compile_shader ne supporte pas directement les paramètres [[threadgroup(N)]]. Si votre noyau a besoin de mémoire de groupe de threads, restructurez pour utiliser des tableaux threadgroup déclarés dans le corps du noyau à la place.

Travailler avec les TensorIterators

REGISTER_UNARY_OP / REGISTER_BINARY_OP cachent la plomberie d'itérateur. Les noyaux avec des paramètres supplémentaires ou des mises en page non élémentaires doivent piloter TensorIterator directement. Les règles non évidentes :

• Passez TensorIteratorBase&, pas Tensor&. Tensor& perd les informations d'offset/shape que with_32bit_indexing() produit. Quand le stub vous remet un const TensorBase& (par exemple bernoulli_scalar_stub), construisez l'itérateur en ligne avec at::TensorIterator::borrowing_nullary_op(self) plutôt que de supprimer const.

• iter.tensor(0) retourne le tenseur entier pour les sous-itérateurs. Après with_32bit_indexing(), les sous-itérateurs référencent toujours le stockage original, donc lier iter.tensor(0) naïvement fait à chaque sous-itérateur d'écraser le même préfixe et laisse la queue non initialisée. Utilisez bind_iter_tensors, qui calcule l'offset par chunk et lie le buffer 0 à la tranche. Dispatchez avec iter.numel(), jamais iter.tensor(0).numel() :

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

• Préférez mtl_setArgs<N> aux mtl_setBytes chaînés, et emballez les scalaires du même type dans un argument vecteur. mtl_setArgs<1>(encoder, a, b, c) lie a/b/c aux emplacements 1/2/3 avec la même résolution de surcharge que les macros, donc un std::array<T, N> arrive comme un seul binding constant T2&/T3& (std::array<long,2> → constant long2&) plutôt que N d'entre eux. Deux ints sont un int2, pas deux emplacements de buffer.

• Utilisez ceil_div. Hôte : at::ceil_div de <ATen/ceil_div.h>. Metal : c10::metal::ceil_div de <c10/metal/common.h> (non qualifié après using namespace c10::metal;). Les deux enrobent (a + b - 1) / b.

• IF_CONSTEXPR pour la portabilité Metal 3/4. Metal 4 a if constexpr ; Metal 3 ne l'a pas. Utilisez la macro de <c10/metal/common.h>, par exemple if IF_CONSTEXPR (sizeof(T) == 8) { ... }.

• Partager le stub CPU/CUDA via REGISTER_MPS_DISPATCH. Quand l'op a un stub DECLARE_DISPATCH en amont (distributions, opérations fusionnées), câblez MPS avec REGISTER_MPS_DISPATCH(stub_name, &fn) au lieu d'une entrée native_functions.yaml spécifique à MPS. Le stub prend déjà TensorIteratorBase&.

Travailler avec les grands tenseurs

La plupart des noyaux Metal prennent numel comme uint32_t et indexent sur 32 bits, donc tout ce qui dépasse INT32_MAX a besoin de division côté hôte. Pilotez cela via TensorIterator plutôt que de découper manuellement.

Décomposer via iter.with_32bit_indexing() :

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

Chaque sous-itérateur généré satisfait can_use_32bit_indexing(), donc la récursion se termine après un niveau. Le seuil est INT32_MAX (TensorIterator utilise des offsets signés 32 bits), pas UINT32_MAX — un test juste après INT32_MAX exerce la division mais pas le cast uint32_t ; le cast lui-même n'importe que pour numel > 2^32.

Utilisez un cast vérifié lors du rétrécissement :

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

Utilisez c10::checked_convert (<c10/util/TypeCast.h>) pour que le dépassement devienne une défaillance TORCH_CHECK au lieu d'une sortie corrompue.

Rapport d'erreurs depuis les noyaux

NE JAMAIS copier les résultats sur CPU à des fins de vérification d'erreur. Tout .item(), .cpu(), ou autre lecture hôte sur un tenseur GPU force une synchronisation complète GPU→CPU qui draine chaque op en vol sur le flux — pas seulement la réduction que vous inspectez. Dans un pipeline réaliste, cela immobilise l'ensemble de la file d'attente, et le coût dépasse de loin la vérification réelle. Garder la synchronisation derrière is_mps() n'aide pas ; l'immobilisation se produit chaque fois que l'op s'exécute. Validez sur l'appareil à la place et signalez les erreurs via le mécanisme ci-dessous.

Le code GPU ne peut pas lever d'exception, mais les noyaux peuvent écrire dans un buffer d'erreur partagé que l'hôte lève en tant que c10::AcceleratorError à la prochaine synchronisation. Utilisez TORCH_REPORT_ERROR(error_buf, ...) de <c10/metal/error.h>. Les arguments variadiques sont concaténés ; les entiers sont formatés en base 10. La livraison est asynchrone — le thread défaillant continue de s'exécuter et l'erreur apparaît quand MPSStream::checkLastError() s'exécute ensuite (après une synchronize() ou l'op suivante qui draine le flux), donc ne vous appuyez pas sur elle pour le flux de contrôle à l'intérieur du même dispatch. Crucialement, cela n'ajoute aucune synchronisation forcée : l'erreur s'appuie sur toute synchronisation que le code de l'utilisateur fait déjà.

Côté noyau : prenez un argument device ErrorMessages* error_buf et appelez TORCH_REPORT_ERROR sur le mauvais chemin. Sautez l'élément offensant par la suite pour ne pas aussi corrompre la mémoire :

#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];
}

Côté hôte : liez le buffer d'erreur du flux en tant qu'argument correspondant. mtl_setArgs l'accepte directement :

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

Le buffer est détenu par MPSStream, dimensionné pour 30 messages, et réinitialisé après chaque vidage checkLastError(). Seul le premier message est signalé ; les suivants sont conservés à des fins de débogage mais l'AcceleratorError transporte msg[0].

Liste de vérification

• [ ] Ajout de dispatch MPS à native_functions.yaml
• [ ] Implémentation du noyau Metal dans kernels/
• [ ] Implémentation de l'opérateur côté hôte dans operations/
• [ ] Paramètres scalaires emballés dans des arguments vecteur ; booléens tapés bool et placés en dernier dans les structures partagées
• [ ] Gère les tenseurs vides
• [ ] Gère les tenseurs non contigus
• [ ] Supporte les dtypes requis (float32, float16, bfloat16, et souvent les types complexes via float2/half2)
• [ ] Suppression des défaillances attendues de torch/testing/_internal/common_mps.py
• [ ] Suppression des décorateurs skip/xfail de l'OpInfo (si applicable)