Add the aoti_torch_cuda_masked_select fallback shim (#24788)

.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 = (

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

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<int64_t>* 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<int64_t>& sizes) {
+  const size_t lead = sizes.size() - tensor.dim();
+  std::vector<int64_t> 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<int64_t> 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

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

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()
 

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",

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 <cuda_runtime.h>
+#include <executorch/backends/aoti/slim/c10/core/DeviceType.h>
+#include <executorch/backends/aoti/slim/c10/core/ScalarType.h>
+#include <executorch/backends/cuda/runtime/shims/index.h>
+#include <executorch/backends/cuda/runtime/shims/memory.h>
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/platform/platform.h>
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <vector>
+
+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 <typename T>
+  Tensor* to_device(
+      const std::vector<T>& values,
+      const std::vector<int64_t>& sizes,
+      slim_c10::ScalarType dtype,
+      const std::vector<int64_t>& strides = {}) {
+    Tensor* tensor = nullptr;
+    EXPECT_EQ(
+        aoti_torch_empty_strided(
+            static_cast<int64_t>(sizes.size()),
+            sizes.data(),
+            strides.empty() ? nullptr : strides.data(),
+            static_cast<int32_t>(dtype),
+            static_cast<int32_t>(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 <typename T>
+  std::vector<T> 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<T> 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<Tensor*> tensors_;
+};
+
+TEST_F(AOTITorchCudaMaskedSelectTest, SameShape) {
+  Tensor* self = to_device<float>(
+      {1.0f, -2.0f, 3.0f, -4.0f, 5.0f, -6.0f}, {2, 3}, slim_c10::kFloat);
+  Tensor* mask =
+      to_device<uint8_t>({1, 0, 1, 0, 1, 0}, {2, 3}, slim_c10::kBool);
+
+  const std::vector<float> expected = {1.0f, 3.0f, 5.0f};
+  EXPECT_EQ(masked_select<float>(self, mask), expected);
+}
+
+TEST_F(AOTITorchCudaMaskedSelectTest, MaskBroadcastsOverRows) {
+  Tensor* self =
+      to_device<int64_t>({10, 11, 12, 20, 21, 22}, {2, 3}, slim_c10::kLong);
+  Tensor* mask = to_device<uint8_t>({0, 1, 1}, {3}, slim_c10::kBool);
+
+  const std::vector<int64_t> expected = {11, 12, 21, 22};
+  EXPECT_EQ(masked_select<int64_t>(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<int32_t>({7, 8}, {2, 1}, slim_c10::kInt);
+  Tensor* mask =
+      to_device<uint8_t>({1, 0, 1, 1, 1, 0}, {2, 3}, slim_c10::kByte);
+
+  const std::vector<int32_t> expected = {7, 7, 8, 8};
+  EXPECT_EQ(masked_select<int32_t>(self, mask), expected);
+}
+
+TEST_F(AOTITorchCudaMaskedSelectTest, NonContiguousSelf) {
+  // Stored column-major, so the logical tensor is {{1, 2}, {3, 4}}.
+  Tensor* self = to_device<float>(
+      {1.0f, 3.0f, 2.0f, 4.0f}, {2, 2}, slim_c10::kFloat, {1, 2});
+  Tensor* mask = to_device<uint8_t>({0, 1, 1, 0}, {2, 2}, slim_c10::kBool);
+
+  const std::vector<float> expected = {2.0f, 3.0f};
+  EXPECT_EQ(masked_select<float>(self, mask), expected);
+}
+
+TEST_F(AOTITorchCudaMaskedSelectTest, ZeroDimSelfAndMask) {
+  Tensor* self = to_device<float>({4.0f}, {}, slim_c10::kFloat);
+  Tensor* set = to_device<uint8_t>({1}, {}, slim_c10::kBool);
+  Tensor* unset = to_device<uint8_t>({0}, {}, slim_c10::kBool);
+
+  EXPECT_EQ(masked_select<float>(self, set), std::vector<float>{4.0f});
+  EXPECT_EQ(masked_select<float>(self, unset), std::vector<float>{});
+}
+
+TEST_F(AOTITorchCudaMaskedSelectTest, IncompatibleShapesAreRejected) {
+  Tensor* self = to_device<float>({1.0f, 2.0f, 3.0f}, {3}, slim_c10::kFloat);
+  Tensor* mask = to_device<uint8_t>({1, 0}, {2}, slim_c10::kBool);
+  Tensor* result = nullptr;
+
+  EXPECT_EQ(
+      aoti_torch_cuda_masked_select(self, mask, &result),
+      Error::InvalidArgument);
+}

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",

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)

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",
     }
 )