Fix CUDA mask indexing, CUDA graph outputs and the runtime output-buffer check

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.md

@@ -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/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<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(