CUDA backend: no fallback kernels for `nonzero` and `masked_select`

Ouverte Adaptée aux débutants
  • Quelque chose ne va pas ?
#24,767 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub

Les mainteneurs répondent en général sous 1 jour

Personne n'a encore pris cette issue.

Évaluation

Difficulté
2/5
Temps estimé
1-3 heures
Accessibilité débutants
64/100
Type d'issue
Bug
Clarté
Plutôt claire
Activité
Active
Stack technique
python, pytorch
Domaine
backend, machine-learning

Piste de recherche

Commencez par lire les shims de fallback CUDA existants et repérez où aoti_torch_cuda_nonzero et aoti_torch_cuda_masked_select sont enregistrés ou signalés comme manquants ; l’issue ne nomme pas ces fichiers. Reproduisez le problème avec l’exemple d’exportation fourni, puis vérifiez les tests de fallback existants du backend CUDA ou ajoutez une couverture pour les deux appels directs. Le travail est terminé lorsque les deux appels passent par le backend CUDA sans erreur de kernel manquant.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Description

What happens

A program that calls torch.nonzero or torch.masked_select directly cannot be lowered with the CUDA backend. Export stops with:

RuntimeError: missing fallback kernels (1 total): - aoti_torch_cuda_nonzero Please add them to the AOTI backend.

and the same with aoti_torch_cuda_masked_select. It fails loudly, never with wrong results. Plain boolean-mask indexing (x[x > 0]) exports fine, so only the direct calls hit it. Tools that trace models into explicit ops emit these direct calls for data-dependent shapes.

Repro
import torch
from executorch.backends.cuda.cuda_backend import CudaBackend
from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower

class M(torch.nn.Module):
    def forward(self, x):
        return torch.nonzero(x > 0).sum()   # or torch.masked_select(x, x > 0).sum()

x = torch.tensor([[1.0, -2.0, 3.0], [-4.0, 5.0, -6.0]], device="cuda")
ep = torch.export.export(M().eval().cuda(), (x,))
to_edge_transform_and_lower(
    ep,
    partitioner=[CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")])],
    compile_config=EdgeCompileConfig(_check_ir_validity=False),
).to_executorch()

Tested with executorch 1.6.0+cu132 and torch 2.15.0+cu132 on an H100.

Ask

Add aoti_torch_cuda_nonzero and aoti_torch_cuda_masked_select shims to the CUDA backend, like the existing fallback shims.

Related: #24684.

Langage dominant
Python
Étoiles
5.2k
Forks
1.2k
Merge moyen
1 j 23 h
PR mergées (30 j)
700

Préparer son environnement

Par où commencer

  1. Lisez l'issue en entier, puis le guide de contribution du projet.
  2. Signalez en commentaire que vous la prenez — cela évite que deux personnes fassent le même travail.
  3. Forkez le dépôt et travaillez sur une branche.
  4. Ouvrez une pull request qui référence le numéro de l'issue.

Autres issues de pytorch/executorch

Toutes les issues de pytorch/executorch

Issues similaires

Plus d'issues Python

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.