CUDA backend: CUDA graph fixes for mask indexing, .item() and changed input shapes, masked copies as where, tensor checks dropped before compiling, float64, _reshape_copy, scatter_reduce and varlen_attn, follow-ups

.ci/scripts/wheel/test_shared_libraries.py

@@ -103,6 +103,10 @@
     "aoti_torch_cuda_randint_low_out",
     "aoti_torch_cuda_sort_stable",
     "aoti_torch_cuda_guard_set_index",
+    "aoti_torch_cuda_index_Tensor",
+    "aoti_torch_index_put_out",
+    "aoti_torch_cuda_scatter_reduce_two_out",
+    "aoti_torch_cuda_scatter_src_out",
 )
 
 _QUANTIZED_KERNEL_SYMBOLS = (

backends/aoti/aoti_backend.py

@@ -53,6 +53,22 @@ def _checked_data_name(data_name: str) -> str:
     return data_name
 
 
+def _written_buffer_options(program: ExportedProgram) -> Dict[str, Any]:
+    """Inductor options for a program that writes a buffer.
+
+    They turn off Inductor's joint-graph constant folding for the whole
+    program, not only for the buffers it writes. Before torch 2.15, that
+    folding can turn the update of a one-element buffer, such as a counter,
+    into its value at export, so every call writes that same value
+    (https://github.com/pytorch/pytorch/pull/185077).
+
+    TODO: remove once the pinned torch is 2.15 or newer.
+    """
+    if program.graph_signature.buffers_to_mutate:
+        return {"joint_graph_constant_folding": False}
+    return {}
+
+
 @experimental(
     "This API and all of aoti-driven backend related functionality are experimental."
 )
@@ -286,9 +302,13 @@ def preprocess(
         Preprocess the edge program and compile it using AOTInductor.
         Weights are always separated from the SO file.
         """
+        # First, so that a backend keeping its data in the .pte refuses a data
+        # name before anything is compiled.
+        data_name = cls.data_name_from_compile_specs(compile_specs)
         device_name = cls.get_device_name()
         decomposition_table = cls.get_decomposition_table()
         options = cls.get_aoti_compile_options(compile_specs)
+        options.update(_written_buffer_options(edge_program))
 
         # Move the edge_program to the target device. Routed through a hook so
         # backends can keep large non-parameter tensors (e.g. KV-cache buffers)
@@ -390,11 +410,7 @@ def preprocess(
         named_data_store.add_named_data(so_blob_key, so_data, 1, None)
         # Determine whether to save named data externally based on backend setting
         # External: save to separate .ptd file, otherwise merge with .pte file
-        external_tag = (
-            cls.data_name_from_compile_specs(compile_specs)
-            if cls.save_data_externally()
-            else None
-        )
+        external_tag = data_name if cls.save_data_externally() else None
 
         named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag)
 
@@ -435,11 +451,7 @@ def generate_data_name_compile_spec(cls, data_name: str) -> CompileSpec:
         Generate a CompileSpec that names the external data file, written as
         ``<data_name>.ptd``, so several programs can share a folder.
         """
-        if not cls.save_data_externally():
-            raise ValueError(
-                f"{cls.__name__} keeps its data in the .pte file, so it takes no "
-                "data name"
-            )
+        cls._check_saves_data_externally()
         return CompileSpec(
             COMPILE_SPEC_KEYS.DATA_NAME.value,
             _checked_data_name(data_name).encode("utf-8"),
@@ -453,9 +465,18 @@ def data_name_from_compile_specs(cls, compile_specs: List[CompileSpec]) -> str:
         """
         for spec in compile_specs:
             if spec.key == COMPILE_SPEC_KEYS.DATA_NAME.value:
+                cls._check_saves_data_externally()
                 return _checked_data_name(spec.value.decode("utf-8"))
         return f"aoti_{cls.get_device_name()}_blob"
 
+    @classmethod
+    def _check_saves_data_externally(cls) -> None:
+        if not cls.save_data_externally():
+            raise ValueError(
+                f"{cls.__name__} keeps its data in the .pte file, so it takes no "
+                "data name"
+            )
+
     @classmethod
     def method_name_from_compile_specs(
         cls,

backends/aoti/common_shims_slim.cpp

@@ -145,6 +145,10 @@ int32_t aoti_torch_dtype_float32() {
   return 6; // ScalarType::Float
 }
 
+int32_t aoti_torch_dtype_float64() {
+  return 7; // ScalarType::Double
+}
+
 int32_t aoti_torch_dtype_bfloat16() {
   return 15; // ScalarType::BFloat16
 }

backends/aoti/common_shims_slim.h

@@ -83,6 +83,7 @@ aoti_torch_get_device_index(Tensor* tensor, int32_t* ret_device_index);
 // ============================================================
 
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float32();
+AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float64();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_bfloat16();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float16();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int64();

backends/aoti/slim/c10/core/ScalarType.h

@@ -13,12 +13,14 @@
 #include <ostream>
 
 #include <executorch/runtime/core/portable_type/bfloat16.h>
+#include <executorch/runtime/core/portable_type/half.h>
 #include <executorch/runtime/platform/assert.h>
 
 namespace executorch::backends::aoti::slim::c10 {
 
-// Import BFloat16 from ExecuTorch's portable_type
+// Import BFloat16 and Half from ExecuTorch's portable_type
 using BFloat16 = ::executorch::runtime::etensor::BFloat16;
+using Half = ::executorch::runtime::etensor::Half;
 
 /// Enum representing the scalar type (dtype) of tensor elements.
 /// Note: Enum values must match PyTorch's c10::ScalarType for compatibility.
@@ -30,7 +32,7 @@ enum class ScalarType : int8_t {
   Long = 4, // int64_t
   Half = 5, // float16
   Float = 6, // float
-  // Double = 7,   // double - not currently needed
+  Double = 7, // double
   // ComplexHalf = 8,
   // ComplexFloat = 9,
   // ComplexDouble = 10,
@@ -50,6 +52,7 @@ constexpr ScalarType kInt = ScalarType::Int;
 constexpr ScalarType kLong = ScalarType::Long;
 constexpr ScalarType kHalf = ScalarType::Half;
 constexpr ScalarType kFloat = ScalarType::Float;
+constexpr ScalarType kDouble = ScalarType::Double;
 constexpr ScalarType kBool = ScalarType::Bool;
 constexpr ScalarType kBFloat16 = ScalarType::BFloat16;
 
@@ -72,6 +75,8 @@ inline size_t elementSize(ScalarType t) {
       return 2; // sizeof(__half) = 2 bytes
     case ScalarType::Float:
       return sizeof(float);
+    case ScalarType::Double:
+      return sizeof(double);
     case ScalarType::Bool:
       return sizeof(bool);
     case ScalarType::BFloat16:
@@ -100,6 +105,8 @@ inline const char* toString(ScalarType t) {
       return "Half";
     case ScalarType::Float:
       return "Float";
+    case ScalarType::Double:
+      return "Double";
     case ScalarType::Bool:
       return "Bool";
     case ScalarType::BFloat16:
@@ -116,7 +123,7 @@ inline const char* toString(ScalarType t) {
 /// @return true if the scalar type is floating point, false otherwise.
 inline bool isFloatingType(ScalarType t) {
   return t == ScalarType::Half || t == ScalarType::Float ||
-      t == ScalarType::BFloat16;
+      t == ScalarType::Double || t == ScalarType::BFloat16;
 }
 
 /// Checks if the scalar type is an integral type (including bool optionally).
@@ -157,6 +164,7 @@ inline bool isValidScalarType(ScalarType t) {
     case ScalarType::Long:
     case ScalarType::Half:
     case ScalarType::Float:
+    case ScalarType::Double:
     case ScalarType::Bool:
     case ScalarType::BFloat16:
       return true;

backends/aoti/slim/c10/core/test/test_scalar_type.cpp

@@ -38,6 +38,7 @@ const std::vector<ScalarTypeTestData> kAllScalarTypes = {
     {ScalarType::Long, 4, 8, "Long", false, true, true, false},
     {ScalarType::Half, 5, 2, "Half", true, false, false, false},
     {ScalarType::Float, 6, 4, "Float", true, false, false, false},
+    {ScalarType::Double, 7, 8, "Double", true, false, false, false},
     {ScalarType::Bool, 11, 1, "Bool", false, false, true, true},
     {ScalarType::BFloat16, 15, 2, "BFloat16", true, false, false, false},
 };
@@ -224,6 +225,7 @@ TEST_F(IsValidScalarTypeTest, AllSupportedTypesAreValid) {
   EXPECT_TRUE(isValidScalarType(ScalarType::Long));
   EXPECT_TRUE(isValidScalarType(ScalarType::Half));
   EXPECT_TRUE(isValidScalarType(ScalarType::Float));
+  EXPECT_TRUE(isValidScalarType(ScalarType::Double));
   EXPECT_TRUE(isValidScalarType(ScalarType::Bool));
   EXPECT_TRUE(isValidScalarType(ScalarType::BFloat16));
 }

backends/aoti/slim/core/test/test_slimtensor_basic.cpp

@@ -508,18 +508,18 @@ TEST(SlimTensorValidationTest, InvalidDtypeUndefined) {
       "");
 }
 
-TEST(SlimTensorValidationTest, InvalidDtypeDouble) {
+TEST(SlimTensorValidationTest, InvalidDtypeComplexFloat) {
   std::vector<int64_t> sizes = {2, 3};
   std::vector<int64_t> strides = {3, 1};
-  size_t nbytes = 6 * sizeof(double);
+  size_t nbytes = 6 * 2 * sizeof(float);
   Storage storage = make_cpu_storage(nbytes);
 
   EXPECT_DEATH(
       SlimTensor(
           std::move(storage),
           makeArrayRef(sizes),
           makeArrayRef(strides),
-          static_cast<c10::ScalarType>(7)), // Double = 7
+          static_cast<c10::ScalarType>(9)), // ComplexFloat = 9
       "");
 }
 

backends/aoti/slim/factory/test/test_empty.cpp

@@ -249,15 +249,15 @@ TEST(EmptyStridedTest, InvalidDtypeUndefined) {
       "");
 }
 
-TEST(EmptyStridedTest, InvalidDtypeDouble) {
+TEST(EmptyStridedTest, InvalidDtypeComplexFloat) {
   std::vector<int64_t> sizes = {2, 3};
   std::vector<int64_t> strides = {3, 1};
 
   EXPECT_DEATH(
       empty_strided(
           makeArrayRef(sizes),
           makeArrayRef(strides),
-          static_cast<c10::ScalarType>(7)), // Double = 7
+          static_cast<c10::ScalarType>(9)), // ComplexFloat = 9
       "");
 }
 
@@ -280,9 +280,9 @@ TEST(EmptyTest, InvalidDtypeUndefined) {
   EXPECT_DEATH(empty({2, 3}, c10::ScalarType::Undefined), "");
 }
 
-TEST(EmptyTest, InvalidDtypeDouble) {
+TEST(EmptyTest, InvalidDtypeComplexFloat) {
   EXPECT_DEATH(
-      empty({2, 3}, static_cast<c10::ScalarType>(7)), // Double = 7
+      empty({2, 3}, static_cast<c10::ScalarType>(9)), // ComplexFloat = 9
       "");
 }
 

backends/aoti/slim/factory/test/test_from_blob.cpp

@@ -327,12 +327,12 @@ TEST(FromBlobTest, InvalidDtypeUndefined) {
       from_blob(external_data, {2, 3}, c10::ScalarType::Undefined), "");
 }
 
-TEST(FromBlobTest, InvalidDtypeDouble) {
+TEST(FromBlobTest, InvalidDtypeComplexFloat) {
   constexpr size_t kNumFloats = 6;
   float external_data[kNumFloats];
 
   EXPECT_DEATH(
-      from_blob(external_data, {2, 3}, static_cast<c10::ScalarType>(7)), "");
+      from_blob(external_data, {2, 3}, static_cast<c10::ScalarType>(9)), "");
 }
 
 TEST(FromBlobTest, InvalidDeviceType) {

backends/aoti/tests/test_common_shims_slim.cpp

@@ -143,6 +143,28 @@ void runGetSizesTest(slim_c10::DeviceType device_type) {
   delete tensor;
 }
 
+void runGetSizeTest(slim_c10::DeviceType device_type) {
+  std::vector<int64_t> sizes = {2, 3};
+  std::vector<int64_t> strides = calculateContiguousStrides(sizes);
+  slim_c10::Device device(device_type, 0);
+
+  Tensor* tensor = new Tensor(slim::empty_strided(
+      slim::makeArrayRef(sizes),
+      slim::makeArrayRef(strides),
+      slim_c10::ScalarType::Float,
+      device));
+
+  int64_t size = -1;
+  EXPECT_EQ(aoti_torch_get_size(tensor, 1, &size), Error::Ok);
+  EXPECT_EQ(size, 3);
+  EXPECT_EQ(aoti_torch_get_size(tensor, -2, &size), Error::Ok);
+  EXPECT_EQ(size, 2);
+  EXPECT_EQ(aoti_torch_get_size(tensor, 2, &size), Error::InvalidArgument);
+  EXPECT_EQ(aoti_torch_get_size(tensor, -3, &size), Error::InvalidArgument);
+
+  delete tensor;
+}
+
 void runGetStridesTest(slim_c10::DeviceType device_type) {
   std::vector<int64_t> sizes = {2, 3, 4};
   std::vector<int64_t> strides = calculateContiguousStrides(sizes);
@@ -432,6 +454,10 @@ TEST_F(CommonShimsSlimTest, GetSizes_CPU) {
   runGetSizesTest(slim_c10::DeviceType::CPU);
 }
 
+TEST_F(CommonShimsSlimTest, GetSize_CPU) {
+  runGetSizeTest(slim_c10::DeviceType::CPU);
+}
+
 TEST_F(CommonShimsSlimTest, GetStrides_CPU) {
   runGetStridesTest(slim_c10::DeviceType::CPU);
 }
@@ -483,6 +509,13 @@ TEST_F(CommonShimsSlimTest, GetSizes_CUDA) {
   runGetSizesTest(slim_c10::DeviceType::CUDA);
 }
 
+TEST_F(CommonShimsSlimTest, GetSize_CUDA) {
+  if (!isCudaAvailable()) {
+    GTEST_SKIP() << "CUDA not available";
+  }
+  runGetSizeTest(slim_c10::DeviceType::CUDA);
+}
+
 TEST_F(CommonShimsSlimTest, GetStrides_CUDA) {
   if (!isCudaAvailable()) {
     GTEST_SKIP() << "CUDA not available";
@@ -555,6 +588,7 @@ TEST_F(CommonShimsSlimTest, NullTensorArgument) {
   EXPECT_EQ(
       aoti_torch_get_data_ptr(nullptr, &data_ptr), Error::InvalidArgument);
   EXPECT_EQ(aoti_torch_get_sizes(nullptr, &sizes), Error::InvalidArgument);
+  EXPECT_EQ(aoti_torch_get_size(nullptr, 0, &numel), Error::InvalidArgument);
   EXPECT_EQ(aoti_torch_get_strides(nullptr, &strides), Error::InvalidArgument);
   EXPECT_EQ(aoti_torch_get_dtype(nullptr, &dtype), Error::InvalidArgument);
   EXPECT_EQ(aoti_torch_get_dim(nullptr, &dim), Error::InvalidArgument);
@@ -566,6 +600,7 @@ TEST_F(CommonShimsSlimTest, NullReturnPointer) {
 
   EXPECT_EQ(aoti_torch_get_data_ptr(tensor, nullptr), Error::InvalidArgument);
   EXPECT_EQ(aoti_torch_get_sizes(tensor, nullptr), Error::InvalidArgument);
+  EXPECT_EQ(aoti_torch_get_size(tensor, 0, nullptr), Error::InvalidArgument);
   EXPECT_EQ(aoti_torch_get_strides(tensor, nullptr), Error::InvalidArgument);
   EXPECT_EQ(aoti_torch_get_dtype(tensor, nullptr), Error::InvalidArgument);
   EXPECT_EQ(aoti_torch_get_dim(tensor, nullptr), Error::InvalidArgument);
@@ -663,6 +698,7 @@ TEST_F(CommonShimsSlimTest, ConsistentPointerReturn) {
 TEST_F(CommonShimsSlimTest, DTypeConstants) {
   // Verify dtype constants match expected PyTorch ScalarType values
   EXPECT_EQ(aoti_torch_dtype_float32(), 6); // ScalarType::Float
+  EXPECT_EQ(aoti_torch_dtype_float64(), 7); // ScalarType::Double
   EXPECT_EQ(aoti_torch_dtype_bfloat16(), 15); // ScalarType::BFloat16
   EXPECT_EQ(aoti_torch_dtype_int64(), 4); // ScalarType::Long
   EXPECT_EQ(aoti_torch_dtype_int32(), 3); // ScalarType::Int

backends/cuda/BUCK

@@ -79,7 +79,9 @@ fbcode_target(
         "passes/__init__.py",
         "passes/move_cond_predicate_to_cpu.py",
         "passes/lower_offgraph_kv.py",
+        "passes/remove_assert_async.py",
         "passes/replace_int64_floordiv.py",
+        "passes/replace_masked_copy_with_where.py",
     ],
     visibility = [
         "//executorch/backends/cuda/...",