Fix CUDA mask indexing, CUDA graph outputs and the runtime output-buffer check (#24299)

## What is wrong today

Robot policies and follow-up checks on #24222 found these bugs in the
CUDA backend and the runtime:

1. **CUDA mask indexing:** `x[mask]` and `x[mask] = values` have no C
shims, so lowering fails with "missing fallback kernels". This blocks
SmolVLA.
2. **CUDA graphs:** a channels-last output turns CUDA graphs off,
although a raw copy is correct.
3. **CUDA graph shapes:** when the captured call has a new batch size,
it and every replay return the previous output shape.
4. **Output buffers:** `Method::set_output_data_ptr` checks a buffer
against the output's current size. After a dynamic output shrinks, a
too-small buffer is accepted and a later, larger call writes past it.

## What this change does

One commit per fix:

1. New shims for index and index_put follow ATen's indexing rules. The
Windows import library gains the new names. Inside a CUDA graph capture,
a mask index fails with a clear error.
2. The check compares output strides with the strides the ETensor's dim
order gives.
3. The capture step resizes each output to the captured shape.
4. The buffer is checked against the output's upper bound. Code that
binds a smaller buffer after a call now gets `InvalidArgument`.

## Tests

Each new test fails without its fix.

- **CUDA (A100):** mask reads and writes, also inside a capture; the
raw-copy check; a capture call with a new batch size.
- **Runtime:** a too-small output buffer is refused after a dynamic
output shrinks; the runtime tests pass.

On LeRobot policies, SmolVLA now lowers to CUDA and runs in 43 ms per
call on an A100, against 504 ms in eager PyTorch (cosine similarity
0.99996).

Author: shoumikhin

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<int64_t>(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;

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

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

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

backends/cuda/runtime/aoti_cuda_shims.lib

@@ -10,7 +10,7 @@ pin moves and the generated wrapper starts calling something new:
 
 ```
 nm --defined-only aoti_cuda_shims.lib \
-  | sed -n 's/.* T _\?\(aoti_torch_[a-z_0-9]*\)$/\1/p' \
+  | sed -n 's/.* T _\?\(aoti_torch_[A-Za-z_0-9]*\)$/\1/p' \
   | sort -u > exports.txt
 # add the new names to exports.txt, then
 { echo 'LIBRARY aoti_cuda_shims.dll'; echo 'EXPORTS'; sed 's/^/    /' exports.txt; } \

backends/cuda/runtime/aoti_cuda_shims.lib.md

@@ -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<size_t>(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() {
@@ -1090,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(