From 024b6c455b5d503822b91ee75466501ec1a4dc0b Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Mon, 9 Nov 2026 15:34:36 -0800 Subject: [PATCH] Add the aoti_torch_cuda_masked_select fallback shim (#24788) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/24788 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 --- .ci/scripts/wheel/test_shared_libraries.py | 1 + backends/cuda/cuda_backend.py | 1 + backends/cuda/runtime/aoti_cuda_shims.lib | Bin 59238 -> 60078 bytes backends/cuda/runtime/shims/index.cu | 61 +++++++ backends/cuda/runtime/shims/index.h | 18 +++ .../cuda/runtime/shims/tests/CMakeLists.txt | 4 +- backends/cuda/runtime/shims/tests/targets.bzl | 1 + .../test_aoti_torch_cuda_masked_select.cpp | 152 ++++++++++++++++++ backends/cuda/tests/test_cuda_export.py | 16 ++ backends/cuda/tests/test_index_shims.py | 22 +++ backends/cuda/tests/test_sort_shim.py | 1 + 11 files changed, 276 insertions(+), 1 deletion(-) create mode 100644 backends/cuda/runtime/shims/tests/test_aoti_torch_cuda_masked_select.cpp diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py index ed7134a2766..da7187d3bba 100644 --- a/.ci/scripts/wheel/test_shared_libraries.py +++ b/.ci/scripts/wheel/test_shared_libraries.py @@ -107,6 +107,7 @@ "aoti_torch_index_put_out", "aoti_torch_cuda_scatter_reduce_two_out", "aoti_torch_cuda_scatter_src_out", + "aoti_torch_cuda_masked_select", ) _QUANTIZED_KERNEL_SYMBOLS = ( diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 56cfd6b9454..d8f59bbc479 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -1219,6 +1219,7 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: "aoti_torch_index_put_out": None, "aoti_torch_cuda_scatter_reduce_two_out": None, "aoti_torch_cuda_scatter_src_out": None, + "aoti_torch_cuda_masked_select": None, } @classmethod diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index ad54a4fdb24e1e8c3395f0245d36a2d9881b5ef3..7245b856e3ee1ca12e216c16626ce6f458e50ed2 100644 GIT binary patch delta 1716 zcmZY8dr*{B6aesZSB3Slkh=(oVAN!$==Xi7$!bz2i8PZ>NR$!fA*d({fk~yiM1hJL zD@Q>Id6;i;eVozoU4Zlw#70fgurM2EY^FwQbVg-5my3T`JM)|Q=IlB5JNMlC{odTo zeYu-!jME-!1b6-tdI5k7U_dj#z%u29Tmb2%C}r>~0Md|lWmq|Yq_0s%OaPEx4OZ+v z0Mct7N_Yo=H2$6v(dq<nLAia~LOq~TF&9E!)`U6O@ZAx4tfE3@LB-8*% z$=r~XnyM^{1CW+*bJDUvCDZv3Ad_XES5|BSkUr!>X_Z~c<>^RySC#x-0MeRlW!*Rc z=`+reTsIWt#-x&C%GLq^sXR$>R&tS4&0~>vaI&qR~-a=kgrhoUmV1y0U3C_bMfEz52Ae6Jo9oTUj0UQs6K z&T_k;dmr%AiJUNj^(zgUx`Hz!9X705X{Mr&Mhik`sh~NIFhPlV{516A(E_!6gIsys z&!GZediV*~K3J{O+}3N z<#4HS78hG^@&*eI>|>^}MH2;PqDDW5@eYTQD9+&NmS|Lo++AI%6+1SYsn2F!FmS8U zIdiNi^QYEQK7|LLYZS7G^K!OX=}_4ufyXN}8dT1au~k+o+{U|iRHMKOiw8}w(gaT2 zp{c9TxU(-dbJXt(j`rNGQPvKwD%{hLuJ7c?g}q&Q)An1be9s7h?izzS_VTU`cYjRx z@}!LYZwOp>Q0qJ!>~q*g5$-5~vyZ4*U@Ljy8L!?n=raj0;ePdWxOF!C z)7wmmXL-Q!^<7=JHCQpYz6a@bJb~ejTLk^nz^D57+pU7mePb8c^h4K#r!-l~^?kUY z&CMl(lAA*L0$k*?4{6~&54zNK(z`DAqx(N{wIxiZfc$1hiM{UxHX6dxAE-`YS(GgO6bJBgF>wWU7j{(;8+_DQn*M&vLqJGGK}VA@O7lUYJ1CFjp#o_#2<70Q zpuli+gFpo`i6m4`<_<~{24cQYi4qmmA$6K^nyKmK=yWcZo%#H8f8TrWx#!+(`eDSw zt0T&0=Z6FaPa68!-2k8jU_ugrUw|^%2S9r1o-(x!K$@1L%m@OI^ht_;AAl6(R6-g7 zq_9KE90!0DS)@c|<^x1!0xXPG7TW-%7*A#C?*P(U-O6&FEUnB_;&K6`_zj9Z5BxMg-W5PSye_cGZlmV|Zws^~81Y$(2R<>l(rbxx1QxB=Nsr*XN8)&N zA8jPhlqCX_lZI>gC3y_yj5al=aw%)x(3N7c;$Z9GQyz1#a7|Q${h5{IkbOH}{4KoLj7^ZQJ%5F}>VKpO*&n@k?~o z{gP5&Eb%a4^y%Ml6RyI6GL065^2YXZBfjhKq4WbW0xzIOj~qO@=ztfRjiYIKMUX(J zLs!iYtFX|ERf`4gsnjVujkkBJ@}cJk=Lt-#=7SvOnGX+I=thkuFyb&@fqV{aH5S_c zX`mocS1$0)dRJO^l)qm%qS2Xpxz|xIn%EE^a9_QqW|$hr(xYQKx9|&{w!h4o8^Ooa8J2`}lAxKR9W{yi=}}{#77PY*q~rvrsorrs&fP z1lFF?$?4?kr=0f2VzUR$Y+1->at=4pt;L8>TTNtc<@?P%qYvG$2*uVmZ|XnCRW0~t zIN5p5MAhy5JdNKDS6;F{hA)K*djBF{?%gY) zyxL_=y`au37F_G)N=@CGpfT6@8D!lI;Lvr4pzvN^^RphEY;|1s)ZVc);uiNX^+$j8 zFrNHmq0ZY}Z(yG`^h!k(?da#|(>uKS0CzCA-$?QIV+CFt(CGF99*Mngr0~IM0vjLL ZDDD*wYlBAm<01cP6MvpbF^>~m{s&~jI?Vt8 diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 7701bb8528b..acfe8807a99 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -1236,6 +1236,43 @@ Error scatter( } } +// The shape a and b broadcast to, as in ATen. +Error broadcast_sizes( + const Tensor& a, + const Tensor& b, + std::vector* sizes) { + const size_t ndim = std::max(a.dim(), b.dim()); + sizes->assign(ndim, 1); + for (size_t i = 0; i < ndim; ++i) { + const int64_t size_a = + i < ndim - a.dim() ? 1 : a.size(i - (ndim - a.dim())); + const int64_t size_b = + i < ndim - b.dim() ? 1 : b.size(i - (ndim - b.dim())); + ET_CHECK_OR_RETURN_ERROR( + size_a == size_b || size_a == 1 || size_b == 1, + InvalidArgument, + "masked_select: self and mask cannot be broadcast together"); + (*sizes)[i] = size_a == 1 ? size_b : size_a; + } + return Error::Ok; +} + +// A view of tensor broadcast to sizes, with stride 0 along the dims it +// expands. +Tensor expanded(const Tensor& tensor, const std::vector& sizes) { + const size_t lead = sizes.size() - tensor.dim(); + std::vector strides(sizes.size(), 0); + for (size_t d = 0; d < tensor.dim(); ++d) { + if (tensor.size(d) != 1) { + strides[lead + d] = tensor.stride(d); + } + } + return tensor.as_strided( + executorch::backends::aoti::slim::makeArrayRef(sizes), + executorch::backends::aoti::slim::makeArrayRef(strides), + tensor.storage_offset()); +} + } // namespace #ifdef __cplusplus @@ -1512,6 +1549,30 @@ AOTITorchError aoti_torch_cuda_scatter_reduce_two_out( include_self != 0); } +AOTITorchError +aoti_torch_cuda_masked_select(Tensor* self, Tensor* mask, Tensor** ret0) { + ET_CHECK_OR_RETURN_ERROR( + self != nullptr && mask != nullptr && ret0 != nullptr, + InvalidArgument, + "aoti_torch_cuda_masked_select: null argument"); + ET_CHECK_OR_RETURN_ERROR( + mask->dtype() == ScalarType::Bool || mask->dtype() == ScalarType::Byte, + InvalidArgument, + "aoti_torch_cuda_masked_select: mask must be bool or uint8"); + std::vector sizes; + ET_CHECK_OK_OR_RETURN_ERROR(broadcast_sizes(*self, *mask, &sizes)); + // A 0-dim mask would cover no dims of the indexed tensor, so index a 1-D + // view of everything instead. + if (sizes.empty()) { + sizes = {1}; + } + Tensor self_view = expanded(*self, sizes); + Tensor mask_view = expanded(*mask, sizes); + Tensor* mask_index = &mask_view; + Tensor* const* indices[] = {&mask_index}; + return aoti_torch_cuda_index_Tensor(&self_view, indices, 1, ret0); +} + #ifdef __cplusplus } #endif diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h index 7a328f57478..290f7e54d97 100644 --- a/backends/cuda/runtime/shims/index.h +++ b/backends/cuda/runtime/shims/index.h @@ -139,6 +139,24 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_scatter_reduce_two_out( const char* reduce, int32_t include_self); +/** + * The elements of self where mask is set, with the semantics of + * aten::masked_select: self and mask broadcast together and the selected + * elements come out in row-major order of the broadcast shape. + * + * Inductor calls this fallback since the number of set elements decides the + * output size. Reading that count waits for the stream, so a CUDA graph cannot + * capture this call. + * + * @param self Tensor to select from (CUDA device, up to 16 dims, any dtype + * aoti_torch_cuda_index_Tensor supports) + * @param mask bool or uint8 mask on the CUDA device, broadcastable with self + * @param ret0 Output: new contiguous 1-D tensor with the dtype of self + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_cuda_masked_select(Tensor* self, Tensor* mask, Tensor** ret0); + #ifdef __cplusplus } #endif diff --git a/backends/cuda/runtime/shims/tests/CMakeLists.txt b/backends/cuda/runtime/shims/tests/CMakeLists.txt index 0e2a4bd599b..235b00586c3 100644 --- a/backends/cuda/runtime/shims/tests/CMakeLists.txt +++ b/backends/cuda/runtime/shims/tests/CMakeLists.txt @@ -50,7 +50,9 @@ set(CUDA_SHIM_TESTS ) # CUDA-specific tests requiring GPU kernels -set(CUDA_KERNEL_TESTS test_aoti_torch_cuda__weight_int4pack_mm) +set(CUDA_KERNEL_TESTS test_aoti_torch_cuda__weight_int4pack_mm + test_aoti_torch_cuda_masked_select +) enable_testing() diff --git a/backends/cuda/runtime/shims/tests/targets.bzl b/backends/cuda/runtime/shims/tests/targets.bzl index 863d0993a7c..9d412fc987e 100644 --- a/backends/cuda/runtime/shims/tests/targets.bzl +++ b/backends/cuda/runtime/shims/tests/targets.bzl @@ -47,6 +47,7 @@ def define_common_targets(is_fbcode = False): cuda_shim_cpp_unittest("aoti_torch_new_tensor_handle") cuda_shim_cpp_unittest("aoti_torch_item_bool") cuda_shim_cpp_unittest("aoti_torch_assign_tensors_out") + cuda_shim_cpp_unittest("aoti_torch_cuda_masked_select") cpp_unittest( name = "test_op__device_copy", diff --git a/backends/cuda/runtime/shims/tests/test_aoti_torch_cuda_masked_select.cpp b/backends/cuda/runtime/shims/tests/test_aoti_torch_cuda_masked_select.cpp new file mode 100644 index 00000000000..69cbcb4258c --- /dev/null +++ b/backends/cuda/runtime/shims/tests/test_aoti_torch_cuda_masked_select.cpp @@ -0,0 +1,152 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +using executorch::backends::cuda::aoti_torch_cuda_masked_select; +using executorch::backends::cuda::aoti_torch_delete_tensor_object; +using executorch::backends::cuda::aoti_torch_empty_strided; +using executorch::runtime::Error; +using Tensor = executorch::backends::aoti::slim::SlimTensor; +namespace slim_c10 = executorch::backends::aoti::slim::c10; + +class AOTITorchCudaMaskedSelectTest : public ::testing::Test { + protected: + void SetUp() override { + et_pal_init(); + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA not available, skipping CUDA tests"; + } + } + + void TearDown() override { + for (Tensor* tensor : tensors_) { + aoti_torch_delete_tensor_object(tensor); + } + } + + template + Tensor* to_device( + const std::vector& values, + const std::vector& sizes, + slim_c10::ScalarType dtype, + const std::vector& strides = {}) { + Tensor* tensor = nullptr; + EXPECT_EQ( + aoti_torch_empty_strided( + static_cast(sizes.size()), + sizes.data(), + strides.empty() ? nullptr : strides.data(), + static_cast(dtype), + static_cast(slim_c10::DeviceType::CUDA), + 0, + &tensor), + Error::Ok); + EXPECT_EQ( + cudaMemcpy( + tensor->data_ptr(), + values.data(), + values.size() * sizeof(T), + cudaMemcpyHostToDevice), + cudaSuccess); + tensors_.push_back(tensor); + return tensor; + } + + template + std::vector masked_select(Tensor* self, Tensor* mask) { + Tensor* result = nullptr; + EXPECT_EQ(aoti_torch_cuda_masked_select(self, mask, &result), Error::Ok); + if (result == nullptr) { + return {}; + } + tensors_.push_back(result); + EXPECT_EQ(result->dim(), 1); + std::vector host(result->numel()); + EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess); + EXPECT_EQ( + cudaMemcpy( + host.data(), + result->data_ptr(), + host.size() * sizeof(T), + cudaMemcpyDeviceToHost), + cudaSuccess); + return host; + } + + std::vector tensors_; +}; + +TEST_F(AOTITorchCudaMaskedSelectTest, SameShape) { + Tensor* self = to_device( + {1.0f, -2.0f, 3.0f, -4.0f, 5.0f, -6.0f}, {2, 3}, slim_c10::kFloat); + Tensor* mask = + to_device({1, 0, 1, 0, 1, 0}, {2, 3}, slim_c10::kBool); + + const std::vector expected = {1.0f, 3.0f, 5.0f}; + EXPECT_EQ(masked_select(self, mask), expected); +} + +TEST_F(AOTITorchCudaMaskedSelectTest, MaskBroadcastsOverRows) { + Tensor* self = + to_device({10, 11, 12, 20, 21, 22}, {2, 3}, slim_c10::kLong); + Tensor* mask = to_device({0, 1, 1}, {3}, slim_c10::kBool); + + const std::vector expected = {11, 12, 21, 22}; + EXPECT_EQ(masked_select(self, mask), expected); +} + +TEST_F(AOTITorchCudaMaskedSelectTest, SelfBroadcastsToMask) { + // self is a column, so each of its values repeats along the mask's rows. + Tensor* self = to_device({7, 8}, {2, 1}, slim_c10::kInt); + Tensor* mask = + to_device({1, 0, 1, 1, 1, 0}, {2, 3}, slim_c10::kByte); + + const std::vector expected = {7, 7, 8, 8}; + EXPECT_EQ(masked_select(self, mask), expected); +} + +TEST_F(AOTITorchCudaMaskedSelectTest, NonContiguousSelf) { + // Stored column-major, so the logical tensor is {{1, 2}, {3, 4}}. + Tensor* self = to_device( + {1.0f, 3.0f, 2.0f, 4.0f}, {2, 2}, slim_c10::kFloat, {1, 2}); + Tensor* mask = to_device({0, 1, 1, 0}, {2, 2}, slim_c10::kBool); + + const std::vector expected = {2.0f, 3.0f}; + EXPECT_EQ(masked_select(self, mask), expected); +} + +TEST_F(AOTITorchCudaMaskedSelectTest, ZeroDimSelfAndMask) { + Tensor* self = to_device({4.0f}, {}, slim_c10::kFloat); + Tensor* set = to_device({1}, {}, slim_c10::kBool); + Tensor* unset = to_device({0}, {}, slim_c10::kBool); + + EXPECT_EQ(masked_select(self, set), std::vector{4.0f}); + EXPECT_EQ(masked_select(self, unset), std::vector{}); +} + +TEST_F(AOTITorchCudaMaskedSelectTest, IncompatibleShapesAreRejected) { + Tensor* self = to_device({1.0f, 2.0f, 3.0f}, {3}, slim_c10::kFloat); + Tensor* mask = to_device({1, 0}, {2}, slim_c10::kBool); + Tensor* result = nullptr; + + EXPECT_EQ( + aoti_torch_cuda_masked_select(self, mask, &result), + Error::InvalidArgument); +} diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 464d62f3111..f41de8c47af 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -986,6 +986,22 @@ def forward(self, counts: torch.Tensor, idx: torch.Tensor): torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) self.assertIn("aoti_torch_cuda_scatter_value_out", str(context.exception)) + @unittest.skipIf( + torch.version.hip is not None, + "ROCm does not build the CUDA-only fallback shims", + ) + def test_masked_select(self): + """A direct masked_select call reaches its fallback kernel, unlike a mask + index, which Inductor reads as index.Tensor.""" + + class MaskedSelectModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.masked_select(x, x > 0).sum() + + self._export_to_cuda_with_lower( + MaskedSelectModule(), (torch.randn(2, 3, device="cuda"),) + ).to_executorch() + @unittest.skipIf( torch.version.hip is not None, "loads through the Python runtime, which the ROCm build does not include", diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 466bbd0fa80..0e3d14f039e 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -367,3 +367,25 @@ def forward(self, q, k, v, cu_seq): (big, big, v, torch.tensor([0, 12, 28], dtype=torch.int32)), ] self._run(Attention(), cases, atol=1e-2, rtol=1e-2) + + def test_masked_select(self): + """A direct masked_select call goes to the shim, unlike a mask index, + which Inductor reads as index.Tensor. The mask broadcasts over rows and a + transposed input checks that the shim reads strides.""" + + class MaskedSelect(torch.nn.Module): + def forward(self, x, row_mask): + return ( + torch.masked_select(x, x > 0), + torch.masked_select(x, row_mask), + torch.masked_select(x.transpose(0, 1), x.transpose(0, 1) < 0), + ) + + torch.manual_seed(0) + x = torch.randn(4, 3) + cases = [ + (x, torch.tensor([True, False, True])), + (torch.zeros(4, 3), torch.zeros(3, dtype=torch.bool)), + (torch.randn(4, 3), torch.ones(3, dtype=torch.bool)), + ] + self._run(MaskedSelect(), cases) diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index 4241405b9a0..39978a536dc 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -45,6 +45,7 @@ "aoti_torch_index_put_out", "aoti_torch_cuda_scatter_reduce_two_out", "aoti_torch_cuda_scatter_src_out", + "aoti_torch_cuda_masked_select", } )