From 3fe9ea1d7f742593c1b922070d3019f11c445855 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 14:17:27 -0700 Subject: [PATCH 1/4] Keep CUDA graphs on for outputs laid out like their ETensor A CUDA graph replay copies each output's raw bytes into its ETensor. The check that decides whether that copy is correct refused every output that was not row-major, before it looked at the ETensor's dim order. So a model whose output is channels-last, with a channels-last ETensor, ran without a CUDA graph, although a raw copy lands every element correctly: return y.contiguous(memory_format=torch.channels_last) # output strides {48, 1, 12, 3}, ETensor dim order {0, 2, 3, 1} # "CUDA graph: disabled ... output 0 is not laid out like its ETensor" The check now works out the strides the ETensor will have from its dim order and the output's sizes, the way resizing the ETensor computes them, and compares those with the output's strides. Size-1 dims are ignored as before, and an empty output has nothing to copy. The ETensor's own strides are not used, because they can still describe the previous shape when shapes are dynamic. The check moves into utils.h, next to the stride check of the eager copy, so it can be unit tested. Tested with test_cuda_output_copy on an A100. The new cases cover a channels-last output into a channels-last ETensor, with current and with stale strides, a row-major output into a row-major ETensor with stale strides, a size-1 dim with an unusual stride, an empty output, a row-major output into a channels-last ETensor and the reverse, a view with gaps between rows, and a rank mismatch. With the old check moved into utils.h unchanged, the two channels-last cases that should pass fail. Removing the size-1 rule, the empty-output rule or the rank check each makes one case fail. With the new check, all tests pass. --- backends/cuda/runtime/cuda_backend.cpp | 22 ------ .../runtime/test/test_cuda_output_copy.cpp | 72 ++++++++++++++++++- backends/cuda/runtime/utils.h | 37 ++++++++++ 3 files changed, 108 insertions(+), 23 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 712ed5ff545..6008d077c32 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -87,28 +87,6 @@ constexpr char kUseSharedCudaStream[] = "use_shared_cuda_stream"; constexpr char kEnableCudaGraphForMethod[] = "enable_cuda_graph_for_method"; constexpr int kCudaGraphWarmupSteps = 3; constexpr char kWeightSharingAcrossMethods[] = "weight_sharing_across_methods"; - -// A CUDA graph replay copies an output's raw bytes into its ETensor, which is -// right only when every element sits at the same offset in both. Size-1 dims -// never move an element, so their strides may differ. -bool output_is_raw_copyable(const SlimTensor* output, const Tensor* et_output) { - if (!output->is_contiguous() || - output->dim() != static_cast(et_output->dim())) { - return false; - } - // The ETensor still has its previous shape here, so its strides can be stale - // for dynamic shapes; the default dim order is enough on its own. - if (executorch::runtime::is_contiguous_dim_order( - et_output->dim_order().data(), et_output->dim())) { - return true; - } - for (size_t d = 0; d < output->dim(); d++) { - if (output->size(d) != 1 && output->stride(d) != et_output->strides()[d]) { - return false; - } - } - return true; -} } // anonymous namespace bool CudaGraphState::start_warmup() { diff --git a/backends/cuda/runtime/test/test_cuda_output_copy.cpp b/backends/cuda/runtime/test/test_cuda_output_copy.cpp index 1952b23a18f..e1b72bd2f17 100644 --- a/backends/cuda/runtime/test/test_cuda_output_copy.cpp +++ b/backends/cuda/runtime/test/test_cuda_output_copy.cpp @@ -12,7 +12,7 @@ namespace slim = ::executorch::backends::aoti::slim; using namespace ::executorch::runtime; -using ::executorch::backends::cuda::copy_slimtensor_to_device_etensor_async; +using namespace ::executorch::backends::cuda; namespace { @@ -115,4 +115,74 @@ TEST_F(CudaOutputCopyTest, RejectsAViewOutsideItsStorage) { Error::InvalidArgument); } +// Whether an output with these sizes and strides can be raw-copied into an +// ETensor that has this shape, dim order and strides from an earlier call. +bool raw_copyable( + std::initializer_list output_sizes, + std::initializer_list output_strides, + std::vector sizes, + std::vector dim_order, + std::vector strides) { + std::vector data(192); + slim::SlimTensor output = slim::from_blob( + data.data(), output_sizes, output_strides, slim::c10::ScalarType::Float); + etensor::TensorImpl impl( + etensor::ScalarType::Float, + static_cast(sizes.size()), + sizes.data(), + data.data(), + dim_order.data(), + strides.data()); + etensor::Tensor etensor(&impl); + return output_is_raw_copyable(&output, &etensor); +} + +TEST(CudaOutputRawCopyTest, AcceptsAnOutputLaidOutLikeItsETensor) { + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, + {48, 1, 12, 3}, + {2, 3, 4, 4}, + {0, 2, 3, 1}, + {48, 1, 12, 3})); + // The ETensor's strides are stale from a smaller shape. + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, {48, 1, 12, 3}, {1, 3, 2, 2}, {0, 2, 3, 1}, {12, 1, 6, 3})); + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, {48, 16, 4, 1}, {1, 3, 2, 2}, {0, 1, 2, 3}, {12, 4, 2, 1})); + // A size-1 dim may have any stride, as with a batch of one. + EXPECT_TRUE(raw_copyable( + {1, 3, 4, 4}, {1, 1, 12, 3}, {1, 3, 4, 4}, {0, 2, 3, 1}, {48, 1, 12, 3})); + // An empty output has nothing to copy, whatever its strides. + EXPECT_TRUE(raw_copyable( + {2, 0, 4, 4}, + {16, 16, 4, 1}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); +} + +TEST(CudaOutputRawCopyTest, RejectsAnOutputLaidOutUnlikeItsETensor) { + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {48, 16, 4, 1}, + {2, 3, 4, 4}, + {0, 2, 3, 1}, + {48, 1, 12, 3})); + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {48, 1, 12, 3}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); + // Every row is followed by a gap, as in a slice of a wider tensor. + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {96, 32, 8, 1}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); + EXPECT_FALSE(raw_copyable( + {3, 4, 4}, {16, 4, 1}, {1, 3, 4, 4}, {0, 1, 2, 3}, {48, 16, 4, 1})); +} + } // namespace diff --git a/backends/cuda/runtime/utils.h b/backends/cuda/runtime/utils.h index 6ca1c64d493..6dd860e11e1 100644 --- a/backends/cuda/runtime/utils.h +++ b/backends/cuda/runtime/utils.h @@ -299,6 +299,43 @@ inline executorch::runtime::Error _copy_slimtensor_to_etensor_impl( } } // namespace +/** + * Whether copying an output's raw bytes into its ETensor puts every element + * where the ETensor expects it, which is how a CUDA graph replay lands it. + * + * The ETensor may still have its previous shape, so its strides can be stale + * for dynamic shapes. Its expected strides come from its dim order and the + * output's sizes instead, the way resizing it computes them. Size-1 dims never + * move an element, so their strides may differ. + * + * @param output Pointer to the output SlimTensor (must not be null). + * @param etensor Pointer to the destination ETensor (must not be null). + * @return True if the raw copy is correct, including when there is nothing + * to copy. + */ +inline bool output_is_raw_copyable( + const executorch::backends::aoti::slim::SlimTensor* output, + const executorch::runtime::etensor::Tensor* etensor) { + const size_t ndim = output->dim(); + if (ndim != static_cast(etensor->dim())) { + return false; + } + if (output->numel() == 0) { + return true; + } + const auto dim_order = etensor->dim_order(); + int64_t expected_stride = 1; + for (size_t i = ndim; i > 0; --i) { + const int64_t d = dim_order[i - 1]; + const int64_t size = output->size(d); + if (size != 1 && output->stride(d) != expected_stride) { + return false; + } + expected_stride *= size; + } + return true; +} + /** * Copies data from a SlimTensor to an ETensor asynchronously. * From 03938c33d505084cddb5ba8ea23a508c7bd58e2b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 14:54:20 -0700 Subject: [PATCH 2/4] Give CUDA graph outputs the shape of the capture call With CUDA graphs on, the backend runs a method eagerly for a few warmup calls, captures the next call into a graph, and replays that graph on later calls. Warmup calls resize each output ETensor to the output's shape as part of copying it. The capture call and the replays copy raw bytes and resized nothing. So when the capture call's shape differed from the last warmup call's, which dynamic shapes allow, the capture call and every replay returned the old shape: inputs of batch 1, 1, 1, then 2, 2, 2 calls 4 to 6 return shape (1, 3, 4, 4) instead of (2, 3, 4, 4) The capture step now resizes each output ETensor to the shape of the captured output, in the same check that runs before the graph is instantiated. Replays never change it, so they keep that shape. A resize the ETensor does not allow, such as one past its upper bound, fails the call the same way a layout mismatch does: the outputs stay with the abandoned graph, and later calls run without a graph. Tested with a new case in test_output_stride_rearrange.py on an A100: a model with a dynamic batch whose capture call has a larger batch than its warmup calls, with row-major and channels-last outputs. Without this change the capture call and both replays return the warmup shape. With it, every call matches eager, and the rest of the file passes too. A resize failure forced at either output of a two-output model fails only that call; later calls run eagerly and match eager, and later CUDA work in the process is unaffected. --- backends/cuda/runtime/cuda_backend.cpp | 15 ++++++++++++ .../tests/test_output_stride_rearrange.py | 23 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 6008d077c32..ce429a9a65a 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -1068,6 +1068,21 @@ class ET_EXPERIMENTAL CudaBackend final handle->method_name.c_str()); return Error::NotSupported; } + // Replay copies raw bytes and never resizes, so the ETensor takes the + // captured shape here and keeps it for every replay. + const Error resize_err = + executorch::ET_RUNTIME_NAMESPACE::resize_tensor( + args[i + n_inputs]->toTensor(), slim_outputs[i]->sizes()); + if (resize_err != Error::Ok) { + std::fill(slim_outputs.begin(), slim_outputs.end(), nullptr); + ET_LOG( + Error, + "CUDA graph: cannot resize output %zu of '%s' to its captured " + "shape; this call fails and later calls run without a graph", + i, + handle->method_name.c_str()); + return resize_err; + } } gerr = cudaGraphInstantiate( diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index b1bf91638a7..431bb56a88f 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -282,6 +282,15 @@ def forward(self, x): return self.proj(x)[:, :, :14] +class _Double(nn.Module): + def __init__(self, memory_format): + super().__init__() + self.memory_format = memory_format + + def forward(self, x): + return (x * 2).contiguous(memory_format=self.memory_format) + + class TestStridedOutputWithCudaGraph(unittest.TestCase): """A CUDA graph replay copies raw bytes, so a strided output must not reach it.""" @@ -359,6 +368,20 @@ def test_strided_output_first_seen_at_capture(self): else: self._assert_matches_eager(module, x, out) + def test_output_shape_changes_at_capture(self): + torch.manual_seed(0) + # The capture call has a larger batch than the warmup calls before it. + inputs = [torch.randn(batch, 3, 4, 4) for batch in (1, 1, 1, 2, 2, 2)] + batch = torch.export.Dim("batch", min=1, max=8) + for memory_format in (torch.contiguous_format, torch.channels_last): + module = _Double(memory_format).eval() + outputs = self._run_with_cuda_graph( + module, (inputs[3].cuda(),), inputs, dynamic_shapes=({0: batch},) + ) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(memory_format=memory_format, call=i): + self._assert_matches_eager(module, x, out) + if __name__ == "__main__": unittest.main() From 270c777da27ffa7a4f70db9324fa5faebad063ed Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 14:51:59 -0700 Subject: [PATCH 3/4] Check set_output_data_ptr buffers against the output's upper bound Method::set_output_data_ptr compared the new buffer with the output's current size. A dynamic output keeps the size of the last call, so after a call that made it smaller, a buffer that fits only that smaller size was accepted. The next call that grows the output back toward its upper bound then writes past the end of that buffer. For example, with an output of at most (4, 4): execute() # output is (2, 4), 32 bytes set_output_data_ptr(buf, 32, 0) # accepted before, refused now execute() # output grows to (4, 4), 64 bytes The check now uses the output's size from MethodMeta, which is its upper bound. This matches the Python set_output binding, which already checks against MethodMeta. For static and bounded dynamic outputs the two sizes are the same before the first call, so a buffer that was valid then is still accepted. The behavior change: code that binds a buffer sized for a smaller output after a call now gets InvalidArgument, and must pass a buffer that fits the upper bound. In ATen mode, an unbounded dynamic output now also needs a buffer of its serialized size. SetOutputDataPtrChecksUpperBound shrinks the output of the dynamic cat test model and checks that a buffer for the smaller size is refused. It fails without this change. --- runtime/executor/method.cpp | 7 ++++--- runtime/executor/method.h | 4 ++-- runtime/executor/test/method_test.cpp | 27 +++++++++++++++++++++++++++ 3 files changed, 33 insertions(+), 5 deletions(-) diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index 34aca8e3856..c79b02eb52c 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -1528,13 +1528,14 @@ Method::set_output_data_ptr(void* buffer, size_t size, size_t output_idx) { return Error::InvalidArgument; } + // A later call can grow a dynamic output back to its upper bound. ET_CHECK_OR_RETURN_ERROR( - t.nbytes() <= size, + tensor_meta->nbytes() <= size, InvalidArgument, "buffer size: %" ET_PRIsize_t - " is smaller then expected tensor size: %" ET_PRIsize_t, + " is smaller than the largest output size: %" ET_PRIsize_t, size, - t.nbytes()); + tensor_meta->nbytes()); // Set data return internal::set_tensor_data(t, buffer, size); diff --git a/runtime/executor/method.h b/runtime/executor/method.h index 834325f008c..501569285f9 100644 --- a/runtime/executor/method.h +++ b/runtime/executor/method.h @@ -154,8 +154,8 @@ class Method final { * * @param[in] buffer The block of memory to point the specified tensor at. * - * @param[in] size the length of buffer in bytes, must be >= the nbytes of the - * specified tensor. + * @param[in] size the length of buffer in bytes, must be >= the largest + * nbytes the output can have, `MethodMeta::output_tensor_meta(i)->nbytes()`. * * @param[in] output_idx The index of the output to set the data_ptr for. Must * correspond to a tensor, and that tensor must not have had a buffer diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp index bcbe3e64316..49cf36ad6c7 100644 --- a/runtime/executor/test/method_test.cpp +++ b/runtime/executor/test/method_test.cpp @@ -317,6 +317,33 @@ TEST_F(MethodTest, AliasedIOTest) { } } +TEST_F(MethodTest, SetOutputDataPtrChecksUpperBound) { + // The "cat" model's output is at most (4, 4) and is not memory planned. + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["cat"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + float input[4] = {}; + int32_t sizes[2] = {1, 4}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {4, 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, 2, sizes, input, dim_order, strides); + ASSERT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), Error::Ok); + float output[16]; + ASSERT_EQ(method->set_output_data_ptr(output, sizeof(output), 0), Error::Ok); + ASSERT_EQ(method->execute(), Error::Ok); + ASSERT_EQ(method->get_output(0).toTensor().nbytes(), 8 * sizeof(float)); + + // A later call can grow the output back to (4, 4), so a buffer that only + // fits the current (2, 4) output must be refused. + float small_output[8]; + EXPECT_EQ( + method->set_output_data_ptr(small_output, sizeof(small_output), 0), + Error::InvalidArgument); +} + TEST_F(MethodTest, SetInputRejectsOverflowingSizes) { // The "cat" model (ModuleDynamicCatUnallocatedIO) has a 2D input. // set_input validates numel/nbytes overflow before resize_tensor. From 420584f8f57e9ae360c8e4552b902a63961d3d60 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 23:46:29 -0700 Subject: [PATCH 4/4] Run mask indexing in the CUDA backend Reading with a bool mask, as in x[mask], gives a result whose size depends on how many mask elements are true. Inductor cannot generate code for that, so it calls an ATen kernel through AOTInductor's C shim: aoti_torch_cuda_index_Tensor. A write through a mask, as in x[mask] = values, goes to aoti_torch_index_put_out (a single value becomes a masked fill instead). The CUDA backend had neither, so lowering failed with "missing fallback kernels". SmolVLM's vision position ids (position_ids[mask] = ...) hit this, which blocks SmolVLA. This adds both shims, with ATen's semantics: - An index is an int64 or int32 tensor, a bool or uint8 mask, or None for a full slice. A mask stands for the positions of its true elements, one index per dimension it covers. Indices broadcast together. The dimensions they produce replace the indexed ones when those are adjacent, and go first otherwise. - The read returns a new tensor with the strides index.Tensor's meta function gives, since the generated code that reads it assumes them. - Inductor also uses the write for accumulation it cannot do with atomics (int64, bool, and bfloat16 below sm_90) and for index_put in deterministic mode. So when an index tensor can send several writes to one element, the shim sorts the index points by element, then applies each element's writes in index order, adding half and bfloat16 values in float. The result is deterministic, and many small bfloat16 terms add up instead of each being rounded away. Masks cannot repeat an element, so they skip the sort. - As in ATen, an index out of range triggers a device-side assert, and an index into a dimension of size 0 is an error unless the indices select nothing. A mask makes the call read its count of true elements back to the host, which waits for the stream, so a CUDA graph cannot capture it. During a capture the shim refuses a mask before it allocates or queues anything, so the capture call fails with a clear error and later calls run eagerly; other CUDA work in the process is unaffected. The generated code reads the result's data-dependent size with aoti_torch_get_size, which the shared shims lacked, so this adds it. Lowering for a Windows target links against the checked-in import library aoti_cuda_shims.lib, so it is regenerated with the three new names, following aoti_cuda_shims.lib.md. The sed pattern in that file skipped names with capital letters, which would have dropped aoti_torch_cuda_index_Tensor on the next regeneration, so it now keeps them. test_mask_index_and_index_put lowers a model with a mask read, a mask write, a mask and an int index on dims that are not adjacent, a mask read of a transposed input (whose result is not row-major), and an int64 accumulation with repeated indices into nonzero counts. It runs three masks against eager. It fails without this change. The int64 index_put case leaves the missing-kernel test. --- backends/aoti/common_shims_slim.cpp | 13 + backends/aoti/common_shims_slim.h | 5 + backends/cuda/CMakeLists.txt | 4 +- backends/cuda/cuda_backend.py | 2 + backends/cuda/runtime/aoti_cuda_shims.lib | Bin 44602 -> 47046 bytes backends/cuda/runtime/aoti_cuda_shims.lib.md | 2 +- backends/cuda/runtime/shims/index.cu | 1062 ++++++++++++++++++ backends/cuda/runtime/shims/index.h | 82 ++ backends/cuda/runtime/targets.bzl | 2 + backends/cuda/tests/test_cuda_export.py | 65 +- backends/cuda/tests/test_sort_shim.py | 2 + 11 files changed, 1224 insertions(+), 15 deletions(-) create mode 100644 backends/cuda/runtime/shims/index.cu create mode 100644 backends/cuda/runtime/shims/index.h diff --git a/backends/aoti/common_shims_slim.cpp b/backends/aoti/common_shims_slim.cpp index d9e255bf220..e92c9d319ae 100644 --- a/backends/aoti/common_shims_slim.cpp +++ b/backends/aoti/common_shims_slim.cpp @@ -36,6 +36,19 @@ AOTITorchError aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes) { return Error::Ok; } +AOTITorchError +aoti_torch_get_size(Tensor* tensor, int64_t d, int64_t* ret_size) { + if (tensor == nullptr || ret_size == nullptr) { + return Error::InvalidArgument; + } + const auto ndim = static_cast(tensor->dim()); + if (d < -ndim || d >= ndim) { + return Error::InvalidArgument; + } + *ret_size = tensor->size(d); + return Error::Ok; +} + AOTITorchError aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides) { if (tensor == nullptr || ret_strides == nullptr) { return Error::InvalidArgument; diff --git a/backends/aoti/common_shims_slim.h b/backends/aoti/common_shims_slim.h index 0d60f578605..dea91ff8e67 100644 --- a/backends/aoti/common_shims_slim.h +++ b/backends/aoti/common_shims_slim.h @@ -39,6 +39,11 @@ aoti_torch_get_data_ptr(Tensor* tensor, void** ret_data_ptr); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes); +// Generated code reads the size of a result whose shape depends on the data +// through this, one dim at a time. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_get_size(Tensor* tensor, int64_t d, int64_t* ret_size); + AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides); diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 5f2879a7b32..fc89ab581bf 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -255,8 +255,8 @@ set(_aoti_cuda_shim_sources runtime/shims/memory.cpp # Only build CUDA-only operator shims when the CUDA toolchain is available. if(NOT EXECUTORCH_BUILD_ROCM AND CMAKE_CUDA_COMPILER) - list(APPEND _aoti_cuda_shim_sources runtime/shims/int4mm.cu - runtime/shims/sort.cu runtime/shims/rand.cu + list(APPEND _aoti_cuda_shim_sources runtime/shims/index.cu + runtime/shims/int4mm.cu runtime/shims/sort.cu runtime/shims/rand.cu ) endif() diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 0bd1599cd5a..1f55a6099a9 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -839,6 +839,8 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: "at::_ops::sort_stable::call": None, "aoti_torch_cuda_sort_stable": None, "aoti_torch_cuda_randint_low_out": None, + "aoti_torch_cuda_index_Tensor": None, + "aoti_torch_index_put_out": None, } @classmethod diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index 9f48cd086077a24be2d6cff506f6c665f6ba9b88..3294d1ec2629c60f221598224c005977bdd04da4 100644 GIT binary patch delta 1602 zcmZwGeNa?Y6aesZ7k5EmS=(I|5K3kk#zxq;EW3-2u%L(`65*nviGt!{m@Hq5qNtDo zQh}}<0yMvM5>YEXIn!8bIp&m2jp>ZzjAkEymYMlS+T@I#yVo3>dNaTG?z^Aoop4H<+GN-yHhOWv;0H_BT`v4&Jf-+$QKzg=ON%8_n$ek%1b=4Gl0I9QF z>E?0Ln;eYP%YCI_hk^wF(m;&j9l8S`y>m&~GXx;*=SihQ9JO>fOBvx{q+{G)dVh~{ zasz;L%1L)cb(nsIuccxhz8F6q_vVC0hTeSsrn*LK@H~fGbCNOAV8Ju`DD=-zWwKj;9jd< zIB>ctLFdYl58=XGQykr}=1rom$zFl=9_!PQTuwTR!4eH073#y)wEq_{oyhr6V5>fg zf*y~c_ht+V88_8W&lc31KO~SNr;?c-f&K!G60=JM-koj6i3Krae}QWk%+bg+#I@Ip z%vf$1L*~Mn0v*M6`ec(=2rF(kB+?&6(*&C4+NpCg_w4h=;hho#-cP-a=St?{sCgRY z%;koz`Bp3nH{(t1S8j;L8>KOLt~4KSFL=C$f1wG_EU?h<{5*l5FG{C`GXC&_#qrd= zkYhPeuF=PvN`;8S@52+RdU2M(H5D4w9puR$R9ev<5ka|2T>`V0+v!Lx*M3@Mp{QjU z0>}6`h$gP-tG3|RDV&}8Xmi@Hl1>PyzWE)IGVZIEpSG? zoknW8=0byoOm%X$CZ5g5ktE@^(XsShgG1nt&33iIinZ}{ERZSis}`HeUD0|Y?r0fH zU2F3>{jE0ovxOt^YNeB0PMe@hZ=j&TZl|E_J$?LNdkj?dmQ7&U zHk-~lg+A*^6d3Gu;`N{zui$w{ET~J*4_&Cp=lsc~gCZ_3s|(ar1nL^AUa45=3*a~X z33#tR2{Zd_bgqXN{|L>Lyp^+He*5t*>L0M67Hy;l!5q%bpe^+0bT+`H(Ft@7Ir1-d zsw86Jt{B`k7)8eiG6bI7tu}+Oe#_4Yu#s~`#X7=hx?7VX!kg(-8D<#^n)7Z zHhD$7eVEP(XgWBSECbh++t@NWpn delta 1260 zcmXZZZ%kWN6aesZD-=Te3MnvIWsxk~WLfrJ!43vO!6CS%(6NCmAPWmJnjv%GM#O-0 zAb-k)jr433Xb>38HNcqTs2_+XGnz#cm?nNOCX2y8@xwwkqJA+`EW~KfAor(j1CV-@MfTtG~M-V{TQKO{g0!TamRCX@`Na+*GzDfW| z&r~w@07$MBB|8B?%8gR;Hv9m28vq5X${{W+6%Q$gn*k&bCzZTT<(LgXsx&E84*{g= zFO}M90IBYRQlAGPo#r!=?}5_fpXW=nv+YVtA%JwAH%V>lO8Zv;QpbwYbqPS~Zd7{t zj5NRvNrSvs8oI1p;pU~FOS#5Vka9ym0!S10l*tbPq-kzQnsqB5*a4)-rm_&hk!KvV z(^UEtEkzHQ@%w8#=%Upt@O~tY5^Oobl$AlAc(0lMNG?jJ!qMXb0|iOA6m6ycLvBI8 z7KQ}AQ(~dU;(S406paZydBj2`hk0$TgxB&*Ewsnu7T9~lpt2V#1e!f|OgoxDvEG*i zHk2Cl>^Mh%E3?yj`O5-5M>FW#+g#^xymorSt;xs{)5oRl`sB^qiSSRNZi zwRL#{KRs#ClSSU^uCq~keUZS>DJM;}a>P-eM5fc30-IkmC_KlJM-7SewJ%?wtC3HI zIr51wk-|+mg3M>e1T~*MAn3c-#|7oKWC^_fhLfgSx#F&tZ4`UnDX^v0Nq>hq^81^K zwAOY|VD<&4%DvL2;h9d2uD53kjCo5}V+*u9usJT4Iy;I5{_wU==Q6o!O^1dz;!?=h z$d)1aUH6+-Mnqh~vLyAKE?zpe)C>bB$NfSFQz`DX9-7<9Rht4H;wV2dT1{^{ol z)b;7)3vkx;emkuW>H-f240=4wk(EI^P7OyB4(WoPTnY-RxZ)D@`EZbrSYnn2gR)pR8%-^6 zTg`Jk>P5cnqZ=B%y2RPsHx2scHfOsV(QyB=9cOochX3BVh~{~NZY}ddU_m48qrC!G L7j exports.txt # add the new names to exports.txt, then { echo 'LIBRARY aoti_cuda_shims.dll'; echo 'EXPORTS'; sed 's/^/ /' exports.txt; } \ diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu new file mode 100644 index 00000000000..97335a75ea0 --- /dev/null +++ b/backends/cuda/runtime/shims/index.cu @@ -0,0 +1,1062 @@ +/* + * 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 +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace executorch::backends::cuda { + +namespace c10_slim = executorch::backends::aoti::slim::c10; +using c10_slim::ScalarType; +using executorch::runtime::Error; + +namespace { + +constexpr int kMaxDims = 16; +constexpr int kThreads = 256; + +// One indexed dimension of the tensor being indexed. Its index is an int64 or +// int32 tensor, or the row-major positions of a mask's true elements, from +// which this dimension's coordinate is (position / mask_divisor) % mask_size. +struct DimIndex { + const void* data; + int64_t strides[kMaxDims]; // along the iteration dims, 0 where broadcast + int64_t size; // of the indexed dimension + int64_t stride; // of the indexed dimension + int64_t mask_divisor; // 0 for an index tensor + int64_t mask_size; + bool is_int32; +}; + +// The iteration space of an indexing call is the shape of the indexed result. +// Each point maps to an element of the indexed tensor and to one of the other +// tensor: the output of a gather, the values of a put. The dims the indices +// produce are [broadcast_start, broadcast_start + broadcast_ndim). +struct IndexPlan { + int ndim; + int64_t sizes[kMaxDims]; + int64_t indexed_strides[kMaxDims]; // 0 along the dims the indices produce + int64_t other_strides[kMaxDims]; + int broadcast_start; + int broadcast_ndim; + int64_t broadcast_numel; + int64_t after_numel; // of the dims after the ones the indices produce + bool from_mask[kMaxDims]; // size is a mask's count, unknown at export + int num_indexed; + DimIndex indexed[kMaxDims]; +}; + +// Adds the offsets of `linear` within dims [first, end). +__device__ __forceinline__ void add_offsets( + const IndexPlan& plan, + int first, + int end, + int64_t linear, + int64_t& indexed_offset, + int64_t& other_offset) { + for (int d = end - 1; d >= first; --d) { + const int64_t coordinate = linear % plan.sizes[d]; + linear /= plan.sizes[d]; + indexed_offset += coordinate * plan.indexed_strides[d]; + other_offset += coordinate * plan.other_strides[d]; + } +} + +// Adds the offsets of a point of the dims the indices produce, including the +// element the indices pick. +__device__ __forceinline__ void add_index_offsets( + const IndexPlan& plan, + int64_t point, + int64_t& indexed_offset, + int64_t& other_offset) { + int64_t index_offsets[kMaxDims]; + for (int i = 0; i < plan.num_indexed; ++i) { + index_offsets[i] = 0; + } + for (int d = plan.broadcast_start + plan.broadcast_ndim - 1; + d >= plan.broadcast_start; + --d) { + const int64_t coordinate = point % plan.sizes[d]; + point /= plan.sizes[d]; + other_offset += coordinate * plan.other_strides[d]; + for (int i = 0; i < plan.num_indexed; ++i) { + index_offsets[i] += coordinate * plan.indexed[i].strides[d]; + } + } + for (int i = 0; i < plan.num_indexed; ++i) { + const DimIndex& index = plan.indexed[i]; + int64_t position = index.is_int32 + ? static_cast(index.data)[index_offsets[i]] + : static_cast(index.data)[index_offsets[i]]; + if (index.mask_divisor != 0) { + position = position / index.mask_divisor % index.mask_size; + } + if (position < 0) { + position += index.size; + } + // Checked on the device, as ATen and Inductor's generated code do, since + // a check on the host would have to wait for the stream. + CUDA_KERNEL_ASSERT_MSG( + position >= 0 && position < index.size, "index out of range"); + indexed_offset += position * index.stride; + } +} + +// Adds the offsets of a point of the other dims, the slice each index picks. +__device__ __forceinline__ void add_slice_offsets( + const IndexPlan& plan, + int64_t point, + int64_t& indexed_offset, + int64_t& other_offset) { + const int end = plan.broadcast_start + plan.broadcast_ndim; + add_offsets( + plan, + end, + plan.ndim, + point % plan.after_numel, + indexed_offset, + other_offset); + add_offsets( + plan, + 0, + plan.broadcast_start, + point / plan.after_numel, + indexed_offset, + other_offset); +} + +__device__ __forceinline__ void locate( + const IndexPlan& plan, + int64_t linear, + int64_t& indexed_offset, + int64_t& other_offset) { + indexed_offset = 0; + other_offset = 0; + const int64_t rest = linear / plan.after_numel; + add_index_offsets( + plan, rest % plan.broadcast_numel, indexed_offset, other_offset); + add_slice_offsets( + plan, + rest / plan.broadcast_numel * plan.after_numel + + linear % plan.after_numel, + indexed_offset, + other_offset); +} + +template +__global__ void gather_kernel( + const IndexPlan plan, + int64_t numel, + const T* __restrict__ self, + T* __restrict__ out) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < numel; + linear += static_cast(gridDim.x) * blockDim.x) { + int64_t self_offset; + int64_t out_offset; + locate(plan, linear, self_offset, out_offset); + out[out_offset] = self[self_offset]; + } +} + +// Adds in AccT, ATen's accumulation type, and rounds once. +template +__device__ __forceinline__ void store(T* out, AccT sum) { + *out = kAccumulate ? static_cast(static_cast(*out) + sum) + : static_cast(sum); +} + +// For indices that cannot repeat an element, which is any mix of masks and +// full slices. +template +__global__ void put_kernel( + const IndexPlan plan, + int64_t numel, + T* __restrict__ out, + const T* __restrict__ values) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < numel; + linear += static_cast(gridDim.x) * blockDim.x) { + int64_t out_offset; + int64_t values_offset; + locate(plan, linear, out_offset, values_offset); + store( + out + out_offset, static_cast(values[values_offset])); + } +} + +__global__ void index_offsets_kernel( + const IndexPlan plan, + int64_t points, + int64_t* __restrict__ out_offsets, + int64_t* __restrict__ values_offsets) { + for (int64_t point = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + point < points; + point += static_cast(gridDim.x) * blockDim.x) { + int64_t out_offset = 0; + int64_t values_offset = 0; + add_index_offsets(plan, point, out_offset, values_offset); + out_offsets[point] = out_offset; + values_offsets[point] = values_offset; + } +} + +// After a stable sort of the index points by the element they pick, the +// points that pick one element are a run in index order. Each thread takes +// one element of the slice of one run, adds the run's values in that order +// or keeps the last, and writes once, so the result does not depend on +// scheduling. +template +__global__ void put_sorted_kernel( + const IndexPlan plan, + int64_t points, + int64_t slice_numel, + const int64_t* __restrict__ out_offsets, + const int64_t* __restrict__ values_offsets, + T* __restrict__ out, + const T* __restrict__ values) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < points * slice_numel; + linear += static_cast(gridDim.x) * blockDim.x) { + const int64_t i = linear / slice_numel; + const int64_t offset = out_offsets[i]; + if (i > 0 && out_offsets[i - 1] == offset) { + continue; + } + int64_t out_offset = offset; + int64_t values_slice_offset = 0; + add_slice_offsets( + plan, linear % slice_numel, out_offset, values_slice_offset); + // Starts from the element itself, so the additions associate as they do + // one after another. + AccT result = kAccumulate ? static_cast(out[out_offset]) : AccT{}; + for (int64_t j = i; j < points && out_offsets[j] == offset; ++j) { + const auto value = + static_cast(values[values_offsets[j] + values_slice_offset]); + result = kAccumulate ? static_cast(result + value) : value; + } + out[out_offset] = static_cast(result); + } +} + +// A mask's element at a row-major position is set. +struct MaskIsSet { + const uint8_t* data; + int ndim; + int64_t sizes[kMaxDims]; + int64_t strides[kMaxDims]; + + __device__ bool operator()(int64_t position) const { + int64_t offset = 0; + for (int d = ndim - 1; d >= 0; --d) { + offset += position % sizes[d] * strides[d]; + position /= sizes[d]; + } + return data[offset] != 0; + } +}; + +// Stream-ordered scratch, freed on the same stream when it goes out of scope. +class Scratch { + public: + Scratch(int32_t device, cudaStream_t stream) + : device_(device), stream_(stream) {} + + Scratch(const Scratch&) = delete; + Scratch& operator=(const Scratch&) = delete; + + ~Scratch() { + for (void* block : blocks_) { + (void)CudaAllocator::deallocate_stream_ordered(block, device_, stream_); + } + } + + void* allocate(size_t nbytes) { + auto block = CudaAllocator::allocate_stream_ordered( + std::max(nbytes, 1), device_, stream_); + if (!block.ok()) { + return nullptr; + } + blocks_.push_back(block.get()); + return block.get(); + } + + private: + int32_t device_; + cudaStream_t stream_; + std::vector blocks_; +}; + +// Writes the row-major positions of a mask's true elements to scratch and +// reads their count back, which waits for the stream. +Error mask_positions( + const Tensor& mask, + Scratch& scratch, + cudaStream_t stream, + const int64_t** positions, + int64_t* count) { + const auto numel = static_cast(mask.numel()); + *positions = nullptr; + *count = 0; + if (numel == 0) { + return Error::Ok; + } + // A graph capture cannot record the host read of the count below, so refuse + // before allocating anything. + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamIsCapturing(stream, &capture)); + ET_CHECK_OR_RETURN_ERROR( + capture == cudaStreamCaptureStatusNone, + NotSupported, + "index: a mask index cannot run inside a CUDA graph capture"); + // CUB counts items in an int on CUDA 12. + ET_CHECK_OR_RETURN_ERROR( + numel <= INT_MAX, + NotSupported, + "index: masks with more than %d elements are not supported", + INT_MAX); + MaskIsSet is_set{static_cast(mask.data_ptr())}; + is_set.ndim = static_cast(mask.dim()); + for (int d = 0; d < is_set.ndim; ++d) { + is_set.sizes[d] = mask.size(d); + is_set.strides[d] = mask.stride(d); + } + auto* selected = + static_cast(scratch.allocate(numel * sizeof(int64_t))); + auto* selected_count = + static_cast(scratch.allocate(sizeof(int64_t))); + ET_CHECK_OR_RETURN_ERROR( + selected != nullptr && selected_count != nullptr, + MemoryAllocationFailed, + "index: allocating the mask positions failed"); + thrust::counting_iterator first(0); + size_t temp_bytes = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceSelect::If( + nullptr, + temp_bytes, + first, + selected, + selected_count, + numel, + is_set, + stream)); + void* temp = scratch.allocate(temp_bytes); + ET_CHECK_OR_RETURN_ERROR( + temp != nullptr, + MemoryAllocationFailed, + "index: allocating the mask scan scratch failed"); + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceSelect::If( + temp, + temp_bytes, + first, + selected, + selected_count, + numel, + is_set, + stream)); + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync( + count, selected_count, sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamSynchronize(stream)); + *positions = selected; + return Error::Ok; +} + +struct Advanced { + int64_t dim; // of the indexed tensor + bool from_mask; + DimIndex index; + std::vector sizes; + std::vector strides; +}; + +// Lays out an indexing call the way ATen's index and index_put do: masks +// become one index per dimension they cover, all indices broadcast together, +// and the dimensions they produce replace the indexed ones if those are +// adjacent and go first otherwise. Sets *may_repeat when two points can land +// on one element, which only an index tensor allows. +Error plan_indexing( + const Tensor& indexed, + Tensor* const* indices, + int64_t num_indices, + Scratch& scratch, + cudaStream_t stream, + IndexPlan* plan, + bool* may_repeat) { + const auto ndim = static_cast(indexed.dim()); + ET_CHECK_OR_RETURN_ERROR( + ndim <= kMaxDims, + NotSupported, + "index: tensors with more than %d dims are not supported", + kMaxDims); + ET_CHECK_OR_RETURN_ERROR( + num_indices <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + *may_repeat = false; + std::vector advanced; + int64_t dim = 0; + for (int64_t i = 0; i < num_indices; ++i) { + const Tensor* index = indices[i]; + if (index == nullptr) { + ++dim; + continue; + } + ET_CHECK_OR_RETURN_ERROR( + index->is_cuda(), + NotSupported, + "index: index %" PRId64 " is not on the CUDA device", + i); + const auto dtype = index->dtype(); + if (dtype == ScalarType::Bool || dtype == ScalarType::Byte) { + const auto mask_ndim = static_cast(index->dim()); + // A 0-dim mask covers no dims, so it adds no index, as in ATen. + if (mask_ndim == 0) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + dim + mask_ndim <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + for (int64_t j = 0; j < mask_ndim; ++j) { + ET_CHECK_OR_RETURN_ERROR( + index->size(j) == indexed.size(dim + j), + InvalidArgument, + "index: mask %" PRId64 + " does not match the indexed shape at dim %" PRId64, + i, + dim + j); + } + const int64_t* positions = nullptr; + int64_t count = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + mask_positions(*index, scratch, stream, &positions, &count)); + // Row-major position p has coordinate p / divisor % size along dim j. + std::vector divisors(mask_ndim, 1); + for (int64_t j = mask_ndim - 1; j > 0; --j) { + divisors[j - 1] = divisors[j] * std::max(index->size(j), 1); + } + for (int64_t j = 0; j < mask_ndim; ++j) { + Advanced entry{dim + j, true}; + entry.index.data = positions; + entry.index.mask_divisor = divisors[j]; + entry.index.mask_size = index->size(j); + entry.sizes = {count}; + entry.strides = {1}; + advanced.push_back(entry); + } + dim += mask_ndim; + } else if (dtype == ScalarType::Long || dtype == ScalarType::Int) { + ET_CHECK_OR_RETURN_ERROR( + dim < ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + ET_CHECK_OR_RETURN_ERROR( + static_cast(index->dim()) <= kMaxDims, + NotSupported, + "index: index tensors with more than %d dims are not supported", + kMaxDims); + Advanced entry{dim, false}; + entry.index.data = index->data_ptr(); + entry.index.is_int32 = dtype == ScalarType::Int; + entry.sizes.assign(index->sizes().begin(), index->sizes().end()); + entry.strides.assign(index->strides().begin(), index->strides().end()); + advanced.push_back(entry); + *may_repeat = true; + ++dim; + } else { + ET_LOG( + Error, + "index: index %" PRId64 + " has dtype %d; indices must be int64, int32, bool or uint8", + i, + static_cast(dtype)); + return Error::InvalidArgument; + } + } + ET_CHECK_OR_RETURN_ERROR( + dim <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + // ATen's advanced indexing rejects a call with only None and 0-dim masks; + // its plan would cover the whole tensor. + ET_CHECK_OR_RETURN_ERROR( + !advanced.empty(), + InvalidArgument, + "index: at least one index must be a tensor that is not a 0-dim mask"); + + size_t broadcast_ndim = 0; + for (const auto& entry : advanced) { + broadcast_ndim = std::max(broadcast_ndim, entry.sizes.size()); + } + std::vector broadcast(broadcast_ndim, 1); + std::vector masked(broadcast_ndim, false); + std::vector fixed(broadcast_ndim, false); + for (const auto& entry : advanced) { + const size_t lead = broadcast_ndim - entry.sizes.size(); + for (size_t d = 0; d < entry.sizes.size(); ++d) { + if (entry.from_mask) { + masked[lead + d] = true; + } else if (entry.sizes[d] != 1) { + fixed[lead + d] = true; + } + int64_t& size = broadcast[lead + d]; + ET_CHECK_OR_RETURN_ERROR( + entry.sizes[d] == 1 || size == 1 || entry.sizes[d] == size, + InvalidArgument, + "index: the indices cannot be broadcast together"); + if (size == 1) { + size = entry.sizes[d]; + } + } + } + + const bool adjacent = advanced.back().dim - advanced.front().dim + 1 == + static_cast(advanced.size()); + const int64_t first = advanced.front().dim; + std::vector is_indexed(ndim, false); + for (const auto& entry : advanced) { + is_indexed[entry.dim] = true; + } + std::vector plain_dims; + for (int64_t d = 0; d < ndim; ++d) { + if (!is_indexed[d]) { + plain_dims.push_back(d); + } + } + // Dims of the indexed tensor in iteration order, -1 for each broadcast dim. + std::vector order; + for (int64_t d : plain_dims) { + if (adjacent && d < first) { + order.push_back(d); + } + } + const auto broadcast_start = static_cast(order.size()); + order.insert(order.end(), broadcast_ndim, -1); + for (int64_t d : plain_dims) { + if (!adjacent || d > first) { + order.push_back(d); + } + } + ET_CHECK_OR_RETURN_ERROR( + order.size() <= static_cast(kMaxDims), + NotSupported, + "index: results with more than %d dims are not supported", + kMaxDims); + + *plan = IndexPlan{}; + plan->ndim = static_cast(order.size()); + plan->broadcast_start = static_cast(broadcast_start); + plan->broadcast_ndim = static_cast(broadcast_ndim); + plan->broadcast_numel = 1; + for (int64_t size : broadcast) { + plan->broadcast_numel *= size; + } + // As in ATen, an index into a dim of size 0 is out of range unless there are + // no index points. The result is empty, so no kernel would assert it. + for (const auto& entry : advanced) { + ET_CHECK_OR_RETURN_ERROR( + plan->broadcast_numel == 0 || indexed.size(entry.dim) > 0, + InvalidArgument, + "index: index out of range for dim %" PRId64 " of size 0", + entry.dim); + } + plan->after_numel = 1; + for (size_t k = static_cast(broadcast_start) + broadcast_ndim; + k < order.size(); + ++k) { + plan->after_numel *= indexed.size(order[k]); + } + for (size_t k = 0; k < order.size(); ++k) { + if (order[k] >= 0) { + plan->sizes[k] = indexed.size(order[k]); + plan->indexed_strides[k] = indexed.stride(order[k]); + } else { + plan->sizes[k] = broadcast[k - broadcast_start]; + plan->from_mask[k] = + masked[k - broadcast_start] && !fixed[k - broadcast_start]; + } + } + plan->num_indexed = static_cast(advanced.size()); + for (size_t i = 0; i < advanced.size(); ++i) { + const auto& entry = advanced[i]; + DimIndex& index = plan->indexed[i]; + index = entry.index; + index.size = indexed.size(entry.dim); + index.stride = indexed.stride(entry.dim); + const size_t lead = broadcast_ndim - entry.sizes.size(); + for (size_t d = 0; d < entry.sizes.size(); ++d) { + if (entry.sizes[d] != 1) { + index.strides[broadcast_start + lead + d] = entry.strides[d]; + } + } + } + return Error::Ok; +} + +int64_t plan_numel(const IndexPlan& plan) { + int64_t numel = 1; + for (int d = 0; d < plan.ndim; ++d) { + numel *= plan.sizes[d]; + } + return numel; +} + +unsigned int grid_size(int64_t numel) { + return static_cast( + std::min((numel + kThreads - 1) / kThreads, 65535)); +} + +// The meta function decided the layout with each mask's count unknown, so a +// count is never taken to be 1, and the result is only known to have fewer +// than 2 elements when another size is 0. +bool is_one(const IndexPlan& plan, int d) { + return !plan.from_mask[d] && plan.sizes[d] == 1; +} + +bool is_contiguous(const IndexPlan& plan, const int64_t* strides) { + bool any_unknown = false; + int64_t numel = 1; + for (int d = 0; d < plan.ndim; ++d) { + if (!plan.from_mask[d] && plan.sizes[d] == 0) { + return true; + } + any_unknown |= plan.from_mask[d]; + numel *= plan.sizes[d]; + } + if (!any_unknown && numel < 2) { + return true; + } + int64_t expected = 1; + for (int d = plan.ndim - 1; d >= 0; --d) { + if (is_one(plan, d)) { + continue; + } + if (strides[d] != expected) { + return false; + } + expected *= plan.sizes[d]; + } + return true; +} + +bool is_channels_last_contiguous( + const IndexPlan& plan, + const int64_t* strides) { + if (plan.ndim != 4) { + return false; + } + int64_t expected = 1; + for (int d : {1, 3, 2, 0}) { + if (!is_one(plan, d)) { + if (strides[d] != expected) { + return false; + } + expected *= plan.sizes[d]; + } + } + return true; +} + +// The output strides index.Tensor's meta function gives, which Inductor bakes +// into the code that reads the result: those of an elementwise output for an +// input with `strides` (compute_elementwise_output_logical_to_physical_perm in +// torch/_prims_common). +std::vector output_strides(const IndexPlan& plan, bool self_is_empty) { + const int64_t* sizes = plan.sizes; + const int64_t* strides = plan.indexed_strides; + const int ndim = plan.ndim; + std::vector perm(ndim); + std::iota(perm.begin(), perm.end(), 0); + const bool contiguous = is_contiguous(plan, strides); + const bool channels_last = is_channels_last_contiguous(plan, strides); + if (self_is_empty || ndim < 2 || (contiguous && !channels_last)) { + // Row-major. + } else if (channels_last && !contiguous) { + perm = {0, 2, 3, 1}; + } else { + auto ge = [](int64_t a, int64_t b) { + return b == 0 || (a != 0 && (a >= b || a % b == 0)); + }; + auto should_swap = [&](int a, int b) { + const int64_t stride_a = strides[a]; + const int64_t stride_b = strides[b]; + if (stride_a == 0 || stride_b == 0) { + return 0; + } + if (stride_a == stride_b) { + return ge(sizes[b], sizes[a]) ? 0 : 1; + } + if (ge(stride_b, stride_a)) { + return -1; + } + return ge(stride_a, stride_b) ? 1 : 0; + }; + std::reverse(perm.begin(), perm.end()); + for (int i = 1; i < ndim; ++i) { + int dim1 = i; + for (int dim0 = i - 1; dim0 >= 0; --dim0) { + const int comparison = should_swap(perm[dim0], perm[dim1]); + if (comparison > 0) { + std::swap(perm[dim0], perm[dim1]); + dim1 = dim0; + } else if (comparison < 0) { + break; + } + } + } + std::reverse(perm.begin(), perm.end()); + } + std::vector result(ndim); + int64_t expected = 1; + for (int k = ndim - 1; k >= 0; --k) { + result[perm[k]] = expected; + expected *= std::max(sizes[perm[k]], 1); + } + return result; +} + +template +Error gather( + const IndexPlan& plan, + int64_t numel, + const void* self, + void* out, + cudaStream_t stream) { + gather_kernel<<>>( + plan, numel, static_cast(self), static_cast(out)); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + +template +Error put( + const IndexPlan& plan, + int64_t numel, + bool may_repeat, + void* out_data, + const void* values_data, + Scratch& scratch, + cudaStream_t stream) { + auto* out = static_cast(out_data); + const auto* values = static_cast(values_data); + if (!may_repeat) { + put_kernel + <<>>(plan, numel, out, values); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; + } + // One key per index point, so the scratch and the sort grow with the + // number of indices, not with the elements they write. + const int64_t points = plan.broadcast_numel; + // CUB counts items in an int on CUDA 12. + ET_CHECK_OR_RETURN_ERROR( + points <= INT_MAX, + NotSupported, + "index_put: more than %d index points are not supported", + INT_MAX); + auto* offsets = + static_cast(scratch.allocate(4 * points * sizeof(int64_t))); + ET_CHECK_OR_RETURN_ERROR( + offsets != nullptr, + MemoryAllocationFailed, + "index_put: allocating the offsets failed"); + int64_t* out_offsets = offsets; + int64_t* values_offsets = offsets + points; + int64_t* sorted_out_offsets = offsets + 2 * points; + int64_t* sorted_values_offsets = offsets + 3 * points; + index_offsets_kernel<<>>( + plan, points, out_offsets, values_offsets); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + const auto count = static_cast(points); + size_t temp_bytes = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + nullptr, + temp_bytes, + out_offsets, + sorted_out_offsets, + values_offsets, + sorted_values_offsets, + count, + 0, + static_cast(sizeof(int64_t) * 8), + stream)); + void* temp = scratch.allocate(temp_bytes); + ET_CHECK_OR_RETURN_ERROR( + temp != nullptr, + MemoryAllocationFailed, + "index_put: allocating the sort scratch failed"); + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + temp, + temp_bytes, + out_offsets, + sorted_out_offsets, + values_offsets, + sorted_values_offsets, + count, + 0, + static_cast(sizeof(int64_t) * 8), + stream)); + put_sorted_kernel + <<>>( + plan, + points, + numel / points, + sorted_out_offsets, + sorted_values_offsets, + out, + values); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + +} // namespace + +#ifdef __cplusplus +extern "C" { +#endif + +AOTITorchError aoti_torch_cuda_index_Tensor( + Tensor* self, + Tensor* const** indices, + int64_t indices_len_, + Tensor** ret0) { + ET_CHECK_OR_RETURN_ERROR( + self != nullptr && ret0 != nullptr && + (indices != nullptr || indices_len_ == 0), + InvalidArgument, + "aoti_torch_cuda_index_Tensor: null argument"); + ET_CHECK_OR_RETURN_ERROR( + self->is_cuda(), + NotSupported, + "aoti_torch_cuda_index_Tensor: self is not on the CUDA device"); + const auto device = static_cast(self->device_index()); + auto stream_result = getCurrentCUDAStream(self->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), + Internal, + "aoti_torch_cuda_index_Tensor: failed to get CUDA stream"); + cudaStream_t stream = stream_result.get(); + + std::vector index_tensors(indices_len_); + for (int64_t i = 0; i < indices_len_; ++i) { + index_tensors[i] = indices[i] == nullptr ? nullptr : *indices[i]; + } + Scratch scratch(device, stream); + IndexPlan plan; + bool may_repeat = false; + ET_CHECK_OK_OR_RETURN_ERROR(plan_indexing( + *self, + index_tensors.data(), + indices_len_, + scratch, + stream, + &plan, + &may_repeat)); + + const std::vector sizes(plan.sizes, plan.sizes + plan.ndim); + const std::vector strides = output_strides(plan, self->numel() == 0); + Tensor* result = nullptr; + ET_CHECK_OK_OR_RETURN_ERROR(aoti_torch_empty_strided( + plan.ndim, + sizes.data(), + strides.data(), + static_cast(self->dtype()), + static_cast(c10_slim::DeviceType::CUDA), + device, + &result)); + // The caller takes ownership of *ret0 only when this call succeeds. + std::unique_ptr owner(result, [](Tensor* tensor) { + (void)aoti_torch_delete_tensor_object(tensor); + }); + std::copy(strides.begin(), strides.end(), plan.other_strides); + + const int64_t numel = plan_numel(plan); + if (numel > 0) { + const void* src = self->data_ptr(); + void* dst = result->data_ptr(); + switch (self->itemsize()) { + case 1: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 2: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 4: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 8: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + default: + ET_LOG( + Error, + "aoti_torch_cuda_index_Tensor: unsupported element size %zu", + self->itemsize()); + return Error::NotSupported; + } + } + *ret0 = owner.release(); + return Error::Ok; +} + +AOTITorchError aoti_torch_index_put_out( + Tensor* out, + Tensor* self, + Tensor* const* indices, + const uint32_t num_indices, + Tensor* values, + bool accumulate) { + ET_CHECK_OR_RETURN_ERROR( + out != nullptr && values != nullptr && + (indices != nullptr || num_indices == 0), + InvalidArgument, + "aoti_torch_index_put_out: null argument"); + ET_CHECK_OR_RETURN_ERROR( + out == self, + NotSupported, + "aoti_torch_index_put_out: out must be self, as AOTInductor calls it"); + ET_CHECK_OR_RETURN_ERROR( + out->is_cuda() && values->is_cuda(), + NotSupported, + "aoti_torch_index_put_out: tensors must be on the CUDA device"); + ET_CHECK_OR_RETURN_ERROR( + values->dtype() == out->dtype(), + InvalidArgument, + "aoti_torch_index_put_out: values and out must have the same dtype"); + const auto device = static_cast(out->device_index()); + auto stream_result = getCurrentCUDAStream(out->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), + Internal, + "aoti_torch_index_put_out: failed to get CUDA stream"); + cudaStream_t stream = stream_result.get(); + + Scratch scratch(device, stream); + IndexPlan plan; + bool may_repeat = false; + ET_CHECK_OK_OR_RETURN_ERROR(plan_indexing( + *out, indices, num_indices, scratch, stream, &plan, &may_repeat)); + + const auto values_ndim = static_cast(values->dim()); + ET_CHECK_OR_RETURN_ERROR( + values_ndim <= plan.ndim, + InvalidArgument, + "aoti_torch_index_put_out: values cannot be broadcast to the indexed shape"); + const int lead = plan.ndim - values_ndim; + for (int d = 0; d < values_ndim; ++d) { + const int64_t size = values->size(d); + ET_CHECK_OR_RETURN_ERROR( + size == 1 || size == plan.sizes[lead + d], + InvalidArgument, + "aoti_torch_index_put_out: values cannot be broadcast to the indexed shape"); + plan.other_strides[lead + d] = size == 1 ? 0 : values->stride(d); + } + + const int64_t numel = plan_numel(plan); + if (numel == 0) { + return Error::Ok; + } + void* dst = out->data_ptr(); + const void* src = values->data_ptr(); + if (!accumulate) { + switch (out->itemsize()) { + case 1: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 2: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 4: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 8: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + default: + ET_LOG( + Error, + "aoti_torch_index_put_out: unsupported element size %zu", + out->itemsize()); + return Error::NotSupported; + } + } + // The accumulation types are ATen's. + switch (out->dtype()) { + case ScalarType::Float: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Half: + return put<__half, float, true>( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::BFloat16: + return put<__nv_bfloat16, float, true>( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Long: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Int: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Short: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Char: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Byte: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Bool: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + default: + ET_LOG( + Error, + "aoti_torch_index_put_out: unsupported dtype %d", + static_cast(out->dtype())); + return Error::NotSupported; + } +} + +#ifdef __cplusplus +} +#endif + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h new file mode 100644 index 00000000000..3476eaa8c88 --- /dev/null +++ b/backends/cuda/runtime/shims/index.h @@ -0,0 +1,82 @@ +/* + * 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. + */ + +#pragma once + +#include + +namespace executorch::backends::cuda { + +using executorch::backends::aoti::AOTITorchError; +using executorch::backends::aoti::Tensor; + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * Advanced indexing, self[indices], with the semantics of aten::index.Tensor. + * + * Inductor calls this fallback when an index is a mask, so the number of true + * elements decides the output size. Reading that count waits for the stream, + * so a CUDA graph cannot capture this call. As in ATen, an index out of range + * triggers a device-side assert, and an index into a dimension of size 0 is + * an error unless the indices select nothing. + * + * @param self Tensor to index (CUDA device). It, each index and the result + * can have up to 16 dims. + * @param indices indices_len_ pointers to index tensors, null for a full + * slice. An index is int64 or int32, or a bool or uint8 mask that covers as + * many dimensions of self as it has. At least one entry must be an index that + * is not a 0-dim mask. + * @param indices_len_ Number of entries in indices + * @param ret0 Output: new tensor, strided the way Inductor expects + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_index_Tensor( + Tensor* self, + Tensor* const** indices, + int64_t indices_len_, + Tensor** ret0); + +/** + * out[indices] = values, or out[indices] += values when accumulate is set, + * with the semantics of aten::index_put_. + * + * Inductor calls this fallback, with out and self the same tensor, for an + * index that is a mask, for accumulation into a dtype its generated code + * cannot add atomically, and in deterministic mode. Writes that land on the + * same element are applied in index order, and half and bfloat16 values are + * summed in float, so the result is deterministic and small terms are not + * lost. A mask makes this call wait for the stream. Indices are checked as + * for aoti_torch_cuda_index_Tensor. + * + * @param out Tensor to write (CUDA device, up to 16 dims, as for + * aoti_torch_cuda_index_Tensor) + * @param self Must be out; AOTInductor passes the tensor it updates as both + * @param indices num_indices index tensors, null for a full slice, with the + * same dtypes as for aoti_torch_cuda_index_Tensor + * @param num_indices Number of entries in indices + * @param values Values to write, broadcastable to the indexed shape, same + * dtype as out + * @param accumulate If true, add values instead of writing them + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_index_put_out( + Tensor* out, + Tensor* self, + Tensor* const* indices, + const uint32_t num_indices, + Tensor* values, + bool accumulate); + +#ifdef __cplusplus +} +#endif + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/targets.bzl b/backends/cuda/runtime/targets.bzl index 5ff2b15867c..d5635100b62 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -35,6 +35,7 @@ def define_common_targets(is_fbcode = False): name = "runtime_shims", srcs = [ "shims/cuda_guard.cpp", + "shims/index.cu", "shims/int4mm.cu", "shims/memory.cpp", "shims/rand.cu", @@ -43,6 +44,7 @@ def define_common_targets(is_fbcode = False): ], headers = [ "shims/cuda_guard.h", + "shims/index.h", "shims/int4mm.cuh", "shims/int4mm.h", "shims/memory.h", diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index d962db09117..96f637a829c 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -6,6 +6,7 @@ import contextlib import os +import tempfile import threading import unittest from typing import Tuple @@ -653,23 +654,63 @@ def test_triton_kernel_mode_off(self): "SDPA kernel export with triton_kernel_mode=OFF failed", ) - def test_missing_scatter_and_index_put_kernels_fail_at_export(self): - """Inductor leaves these int64 ops to ATen kernels the CUDA shim lacks.""" + def test_missing_scatter_kernel_fails_at_export(self): + """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks.""" - class Int64Accumulate(torch.nn.Module): + class Int64ScatterAdd(torch.nn.Module): def forward(self, counts: torch.Tensor, idx: torch.Tensor): - ones = torch.ones_like(idx) - return ( - counts.scatter_add(0, idx, ones), - counts.index_put((idx,), ones, accumulate=True), - ) + return counts.scatter_add(0, idx, torch.ones_like(idx)) inputs = (torch.zeros(8, dtype=torch.long), torch.randint(0, 8, (32,))) with self.assertRaises(RuntimeError) as context: - self._export_to_cuda_with_lower(Int64Accumulate(), inputs) - message = str(context.exception) - self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", message) - self.assertIn("aoti_torch_index_put_out", message) + self._export_to_cuda_with_lower(Int64ScatterAdd(), inputs) + self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", str(context.exception)) + + @unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", + ) + def test_mask_index_and_index_put(self): + """A mask index makes the result size depend on the data, and Inductor + has no atomic add for int64, so these run in the index shims.""" + from executorch.runtime import Runtime + + class MaskIndex(torch.nn.Module): + def forward(self, x, mask, column, idx, counts): + out = x.clone() + out[mask] = x[mask] * 2 + # The mask and column index dims 0 and 2, which are not adjacent. + picked = x[mask, :, column].sum(0) + # A transposed input gives a result that is not row-major. + rows = x.transpose(1, 2)[mask] * 3 + ones = torch.ones_like(idx) + counts = counts.index_put((idx,), ones, accumulate=True) + return out, picked, rows, counts + + model = MaskIndex() + x = torch.randn(4, 5, 6) + column = torch.tensor([2]) + idx = torch.tensor([0, 3, 3, 1, 3]) + # Nonzero counts check that repeated writes add to the old values. + counts = torch.arange(8) * 5 + masks = ( + torch.tensor([True, False, True, True]), + torch.zeros(4, dtype=torch.bool), + torch.ones(4, dtype=torch.bool), + ) + edge = self._export_to_cuda_with_lower( + model, (x, masks[0], column, idx, counts) + ) + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + edge.to_executorch().write_to_file(f) + method = Runtime.get().load_program(path).load_method("forward") + for mask in masks: + inputs = (x, mask, column, idx, counts) + outputs = method.execute(list(inputs)) + for output, expected in zip(outputs, model(*inputs), strict=True): + torch.testing.assert_close(output.cpu(), expected) def test_device_info_propagated_to_cuda_delegate_outputs(self): """ diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index 99cba99cfc5..d81d9e9e187 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -41,6 +41,8 @@ "at::_ops::sort_stable::call", "aoti_torch_cuda_sort_stable", "aoti_torch_cuda_randint_low_out", + "aoti_torch_cuda_index_Tensor", + "aoti_torch_index_put_out", } )