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

Đang mở Phù hợp với người mới
  • Có gì không ổn?
#24,767 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 1 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
2/5
Thời gian dự kiến
1-3 giờ
Mức phù hợp với người mới
64/100
Loại issue
Lỗi
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, pytorch
Lĩnh vực
backend, machine-learning

Hướng nghiên cứu

Bắt đầu bằng cách đọc các shim fallback CUDA hiện có và xác định nơi aoti_torch_cuda_nonzero và aoti_torch_cuda_masked_select được đăng ký hoặc được báo là không tồn tại; issue không nêu tên các tệp đó. Tái hiện lỗi bằng ví dụ export được cung cấp, sau đó kiểm tra các bài kiểm thử fallback hiện có của backend CUDA hoặc bổ sung kiểm thử cho cả hai lời gọi trực tiếp. Công việc hoàn tất khi cả hai lời gọi được hạ xuống qua backend CUDA mà không gặp lỗi thiếu kernel.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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.

Ngôn ngữ chính
Python
Star
5.2k
Fork
1.2k
Merge trung bình
1 ngày 23 giờ
Pull request đã merge (30 ngày)
700

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của pytorch/executorch

Tất cả issue của pytorch/executorch

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.