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

Offen Anfängerfreundlich
  • Stimmt etwas nicht?
#24,767 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen

Maintainer antworten meist innerhalb von 1 Tag

Dieses Issue hat noch niemand übernommen.

Bewertung

Schwierigkeit
2/5
Geschätzter Aufwand
1-3 Stunden
Anfängerfreundlichkeit
64/100
Issue-Typ
Bug
Klarheit
Größtenteils klar
Aktivitätsstatus
Aktiv
Tech-Stack
python, pytorch
Bereich
backend, machine-learning

Rechercherichtung

Lies zuerst die bestehenden CUDA-Fallback-Shims und finde heraus, wo aoti_torch_cuda_nonzero und aoti_torch_cuda_masked_select registriert oder als fehlend gemeldet werden; das Issue nennt diese Dateien nicht. Reproduziere den Fehler mit dem bereitgestellten Exportbeispiel und prüfe dann die bestehenden Fallback-Tests des CUDA-Backends oder ergänze Tests für beide direkten Aufrufe. Erledigt ist die Aufgabe, wenn beide Aufrufe über das CUDA-Backend abgesenkt werden, ohne dass der Fehler wegen des fehlenden Kernels auftritt.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Beschreibung

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.

Vorherrschende Sprache
Python
Sterne
5.2k
Forks
1.2k
Ø Merge
1 T. 23 Std.
Gemergte PRs (30 T.)
700

Entwicklungsumgebung

Erste Schritte

  1. Lesen Sie das ganze Issue und danach den Beitragsleitfaden des Projekts.
  2. Schreiben Sie ins Issue, dass Sie es übernehmen — das erspart doppelte Arbeit.
  3. Forken Sie das Repository und arbeiten Sie in einem Branch.
  4. Öffnen Sie einen Pull Request, der die Issue-Nummer nennt.

Mehr aus pytorch/executorch

Alle Issues in pytorch/executorch

Ähnliche Issues

Weitere Issues zu Python

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.