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

オープン 初心者向け
  • 何か問題がありますか？
#24,767 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
2/5
見積もり時間
1〜3時間
初心者へのやさしさ
64/100
issue の種類
バグ
明瞭さ
おおむね明確
活発さ
活発
技術スタック
python, pytorch
領域
backend, machine-learning

調査の方向性

まず既存のCUDAフォールバックshimを読み、aoti_torch_cuda_nonzeroとaoti_torch_cuda_masked_selectが登録されている箇所、または見つからないと報告される箇所を特定してください。このissueでは対象ファイルは示されていません。提供されたエクスポート例で再現し、その後、CUDAバックエンドの既存のフォールバックテストを確認するか、両方の直接呼び出しをテストするカバレッジを追加してください。両方の呼び出しが、カーネルが見つからないというエラーなしにCUDAバックエンドを通じてlowerされれば完了です。

索引モデルが issue の本文から書いたものです。

説明

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.

主要言語
Python
スター
5.2k
フォーク
1.2k
平均マージ
1日 23時間
マージ済み PR（30日）
700

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

pytorch/executorch のほかの issue

pytorch/executorch の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。