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

Aperta Adatta ai principianti
  • Qualcosa non va?
#24,767 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
2/5
Tempo stimato
1-3 ore
Idoneità per principianti
64/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
backend, machine-learning

Direzione di ricerca

Inizia leggendo gli shim di fallback CUDA esistenti e individua dove aoti_torch_cuda_nonzero e aoti_torch_cuda_masked_select vengono registrati o segnalati come mancanti; l’issue non indica quei file. Riproduci il problema con l’esempio di esportazione fornito, quindi controlla i test di fallback esistenti del backend CUDA o aggiungi test per entrambe le chiamate dirette. Il lavoro è completato quando entrambe le chiamate vengono abbassate attraverso il backend CUDA senza l’errore di kernel mancante.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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.

Lingua principale
Python
Stelle
5.2k
Fork
1.2k
Merge medio
1g 23h
PR unite (30g)
700

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di pytorch/executorch

Tutte le issue di pytorch/executorch

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.