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

Abierto Apto para principiantes
  • ¿Algo no está bien?
#24,767 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 1 día

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
2/5
Tiempo estimado
1-3 horas
Aptitud para principiantes
64/100
Tipo de issue
Error
Claridad
Bastante claro
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
backend, machine-learning

Línea de trabajo

Empieza leyendo los shims de fallback de CUDA existentes y localiza dónde se registran aoti_torch_cuda_nonzero y aoti_torch_cuda_masked_select o dónde se informa de que faltan; el issue no indica esos archivos. Reproduce el problema con el ejemplo de exportación proporcionado y luego revisa las pruebas de fallback existentes del backend CUDA o añade cobertura para ambas llamadas directas. Se considera terminado cuando ambas llamadas se reducen a través del backend CUDA sin el error de kernel ausente.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
Python
Estrellas
5.2k
Forks
1.2k
Merge medio
1 d 23 h
PR fusionados (30 d)
700

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de pytorch/executorch

Todos los issues de pytorch/executorch

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.