[Gasoonjia]

Summary:

A program that calls torch.masked_select directly could not be lowered with the CUDA backend; export stopped with "missing fallback kernels: aoti_torch_cuda_masked_select" (P2594125604). Mask indexing (x[x > 0]) was unaffected because Inductor reads it as index.Tensor, but tools that trace models into explicit ops emit direct masked_select calls.

masked_select(self, mask) is self_b[mask_b] for self and mask broadcast together, so the shim builds stride-0 broadcast views of both and calls the existing aoti_torch_cuda_index_Tensor with the mask as its index. That reuses its CUB mask scan, dtype coverage and CUDA graph refusal, and gives the contiguous 1-D result ATen does. A 0-dim broadcast shape is viewed as 1-D, since a 0-dim mask indexes no dims.

Registered in get_supported_fallback_kernels, the wheel's shim symbol check, and the Windows import library (regenerated per aoti_cuda_shims.lib.md with mingw64 dlltool 2.41, which reproduces the checked-in file byte for byte before the change; superset of the old exports, zero timestamps).

Reviewed By: shoumikhin

Differential Revision: D130366694

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24788

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 4 Pending

As of commit 024b6c4 with merge base fce7099 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-codesync]

@Gasoonjia has exported this pull request. If you are a Meta employee, you can view the originating Diff in D130366694.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[shoumikhin]

@claude review effort=max

[claude]

Claude finished @shoumikhin's task in 20m 15s —— View job

---

Request changes: the new test_cuda_export.py::test_masked_select has no ROCm skip, so it will fail the ROCm unittest job on main once merged (that job was skipped on this PR). The other findings are minor.

Detailed review

PR Review: #24788

Summary

This PR adds aoti_torch_cuda_masked_select. The shim builds stride-0 broadcast views of self and mask and passes them to aoti_torch_cuda_index_Tensor. It is registered in the fallback list, the wheel symbol check and the Windows import library. One finding will break CI: the new export test fails on ROCm. The rest are minor.

What I could not verify:

• Nothing was built or run. There is no GPU here, and running Python was not permitted.
• For the binary aoti_cuda_shims.lib, I only confirmed that the new symbol name is in it. I did not check the claim that it exports everything the old file did.
• The ATen behaviour below comes from my knowledge of PyTorch's source, which is not in this checkout.

Infrastructure

• backends/cuda/tests/test_cuda_export.py:989: test_masked_select fails on ROCm (blocking).
  • Why it runs there: .ci/scripts/test-rocm-aoti.sh:154-157 runs pytest backends/cuda/tests and does not exclude this file. TestCudaExport.setUp only checks torch.cuda.is_available(), which is true on ROCm.
  • Why it fails: under HIP, CudaBackend.get_supported_fallback_kernels() returns {} (cuda_backend.py:1208-1210). collect_unsupported_fallback_kernels then records aoti_torch_cuda_masked_select, and AotiBackend.preprocess raises missing fallback kernels (backends/aoti/aoti_backend.py:366-372). That is the same error the PR description quotes.
  • Why this PR's CI stays green: the unittest-rocm-gfx950 job was skipped on this PR (run 42576457354), so the green "Test ROCm AOTI" check does not cover the new test. The job's if: lets push events through (rocm.yml:8-10, 165-168), so the failure will show up on main after merge.
  • Pattern mismatch: the other tests in this file that ROCm cannot pass have @unittest.skipIf(torch.version.hip is not None, ...). TestIndexShims skips the whole class.
  • The test does not check what its docstring says. The docstring says the call "reaches its fallback kernel", but the test only checks that lowering does not raise. test_index_shims.py::test_masked_select already lowers, runs and compares the same masked_select(x, x > 0).
  • Fix: rewrite it like test_missing_scatter_kernel_fails_at_export. Patch CudaBackend.get_supported_fallback_kernels to leave out aoti_torch_cuda_masked_select, then assertIn that name in the raised RuntimeError. This proves the op is routed to the shim, and it also passes on ROCm. The minimal fix is the hip skipIf. Fix this →

API Design

• index.cu:1560-1563, index.h:153: uint8 masks are accepted, but aten::masked_select rejects them.
  • ATen behaviour: both the CPU and CUDA kernels run TORCH_CHECK(mask.scalar_type() == ScalarType::Bool, "masked_select: expected BoolTensor for mask"). Only index.Tensor still accepts the deprecated uint8 masks.
  • Contradiction: the header promises "the semantics of aten::masked_select" but lists uint8 as valid.
  • Test impact: SelfBroadcastsToMask (test file line 117) is the only C++ case where self broadcasts, and it uses kByte. It locks in behaviour ATen does not have.
  • Fix: accept only Bool and switch that test to kBool. Fix this →
• index.h:151-153: the @param self limits describe self, but they apply to the broadcast shape.
  • "up to 16 dims": plan_indexing checks the broadcast self_view (index.cu:419-423). A 2-D self with a 17-D mask returns NotSupported.
  • "any dtype aoti_torch_cuda_index_Tensor supports": this points to a doc (index.h:22-40) that lists no dtypes. The real rule is an element size of 1, 2, 4 or 8 bytes (index.cu:1341-1363), so complex128 returns NotSupported.
  • Undocumented INT_MAX limit: the CUB scan's INT_MAX limit (index.cu:346-350) applies to the broadcast element count. Two modest inputs such as (65536, 1) and (1, 65536) are rejected, and the doc does not say so.
  • Fix this →

Testing

• test_aoti_torch_cuda_masked_select.cpp: tests for the broadcast handling this PR adds before calling index_Tensor are incomplete.
  • A 0-dim operand against an N-d one is untested: for example a scalar mask on a (2, 3) self (an all-stride-0 mask view), or a 0-dim self with a (2, 3) mask. ZeroDimSelfAndMask covers only the case where both are 0-dim, which is the only case that reaches sizes = {1}.
  • A size-0 dim is untested: for example self (0, 3) with mask (3,).
  • Two-sided broadcast is untested: for example self (2, 1) with mask (1, 3).
  • A non-zero storage offset is untested: expanded forwards tensor.storage_offset() (index.cu:1273), but every test tensor has offset 0, so a regression there would not be caught.
  • The output check is weak: the masked_select<T> helper (line 80) only asserts dim() == 1. It does not check stride(0) == 1 or that the dtype matches self, although index.h:154 promises both.
  • Fix this →
• test_index_shims.py:371-391: no end-to-end case where self broadcasts. Only the mask is broadcast, so no case selects more than self.numel() elements.
  • Suggested case: add an output such as torch.masked_select(x[:, :1], x >= 0). On the torch.zeros(4, 3) input it selects 12 elements from a 4-element self.
  • Possible upstream bound (low confidence, not checked): PyTorch's fake impl for masked_select may cap the unbacked output size at self.numel(). If it does, export adds a runtime assert that this case would trip even though the shim handles it, and the header should document that limit.
  • Fix this →

Code Quality

• index.cu:1254, index.cu:1575: masked_select failures are reported under other op names.
  • broadcast_sizes: the helper has a generic name but hard-codes "masked_select:". The shim's own checks say "aoti_torch_cuda_masked_select:".
  • Errors after broadcasting: every later failure comes from aoti_torch_cuda_index_Tensor, prefixed "aoti_torch_cuda_index_Tensor:" or "index:". A CPU mask reports "index: index 0 is not on the CUDA device" (index.cu:441). A large broadcast reports "index: masks with more than 2147483647 elements…" (index.cu:349) about a mask the user never passed at that size.
  • Pattern mismatch: shared helpers in this file take the public op's name: scatter(const char* name, …) (line 1094) and sort_pairs(…, const char* op) (line 778).
  • Fix: pass the name into broadcast_sizes, and move the body of aoti_torch_cuda_index_Tensor into an internal helper that takes the name, called by both shims. Fix this →

Recommendation

Request Changes

The ROCm failure in test_cuda_export.py::test_masked_select is confirmed by reading the code: the ROCm test script, setUp, the HIP branch of get_supported_fallback_kernels, and the missing-kernel check. This PR's own ROCm unittest job was skipped, so CI will not catch it before merge. Either the one-line skip or the rewrite above is enough. The other findings are minor and could be follow-ups.