[executorch][cuda] Add AOTI tensor zero shim (pytorch#23130)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#23083 by
@Gasoonjia
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/27/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/27/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/27/orig
Differential Revision:
[D120071924](https://our.internmc.facebook.com/intern/diff/D120071924/)
@diff-train-skip-merge

Co-authored-by: gasoonjia <gasoonjia@icloud.com>

Author: pytorchbot

backends/cuda/runtime/shims/memory.cpp

@@ -10,6 +10,11 @@
 #include <executorch/backends/cuda/runtime/shims/memory.h>
 #include <type_traits>
 
+#if defined(CUDA_AVAILABLE)
+#include <executorch/backends/aoti/slim/c10/cuda/Exception.h>
+#include <executorch/backends/aoti/slim/cuda/guard.h>
+#endif
+
 #include <executorch/backends/aoti/slim/factory/empty.h>
 #include <executorch/backends/aoti/slim/factory/from_blob.h>
 #include <executorch/backends/aoti/slim/util/array_ref_util.h>
@@ -402,6 +407,27 @@ aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking) {
   return Error::Ok;
 }
 
+AOTITorchError aoti_torch_zero_(SlimTensor* self) {
+  ET_CHECK_OR_RETURN_ERROR(
+      self != nullptr && self->is_contiguous(),
+      InvalidArgument,
+      "aoti_torch_zero_: expected a contiguous tensor");
+  if (self->nbytes() == 0) {
+    return Error::Ok;
+  }
+#if defined(CUDA_AVAILABLE)
+  if (self->device().is_cuda()) {
+    auto stream = getCurrentCUDAStream(self->device_index());
+    ET_CHECK_OK_OR_RETURN_ERROR(stream.error());
+    ET_CUDA_CHECK_OR_RETURN_ERROR(
+        cudaMemsetAsync(self->data_ptr(), 0, self->nbytes(), stream.get()));
+    return Error::Ok;
+  }
+#endif
+  ET_CHECK_OR_RETURN_ERROR(
+      false, NotSupported, "aoti_torch_zero_: only CUDA tensors are supported");
+}
+
 #define ET_CUDA_DEFINE_ITEM_SHIM(SUFFIX, CTYPE)            \
   AOTITorchError aoti_torch_item_##SUFFIX(                 \
       SlimTensor* tensor, CTYPE* ret_value) {              \

backends/cuda/runtime/shims/memory.h

@@ -195,6 +195,10 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch__reinterpret_tensor(
 AOTI_SHIM_EXPORT AOTITorchError
 aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking);
 
+/// Zeros a contiguous CUDA tensor on the current AOTI stream. Inductor uses it
+/// to clear split-scan workspaces.
+AOTI_SHIM_EXPORT AOTITorchError aoti_torch_zero_(SlimTensor* self);
+
 /// See aoti_torch_item_uint8.
 AOTI_SHIM_EXPORT AOTITorchError
 aoti_torch_item_bool(SlimTensor* tensor, bool* ret_value);

backends/cuda/runtime/shims/tests/CMakeLists.txt

@@ -42,6 +42,7 @@ set(CUDA_SHIM_TESTS
     test_aoti_torch_delete_tensor_object
     test_aoti_torch__reinterpret_tensor
     test_aoti_torch_copy_
+    test_aoti_torch_zero_
     test_aoti_torch_cuda_rand
     test_aoti_torch_new_tensor_handle
     test_aoti_torch_item_bool

backends/cuda/runtime/shims/tests/targets.bzl

@@ -40,6 +40,7 @@ def define_common_targets(is_fbcode = False):
     cuda_shim_cpp_unittest("aoti_torch_create_tensor_from_blob_v2")
     cuda_shim_cpp_unittest("aoti_torch__reinterpret_tensor")
     cuda_shim_cpp_unittest("aoti_torch_copy_")
+    cuda_shim_cpp_unittest("aoti_torch_zero_")
     cuda_shim_cpp_unittest("aoti_torch_cuda_guard")
     cuda_shim_cpp_unittest("aoti_torch_cuda__weight_int4pack_mm")
     cuda_shim_cpp_unittest("aoti_torch_cuda_rand")

backends/cuda/runtime/shims/tests/test_aoti_torch_zero_.cpp

@@ -0,0 +1,85 @@
+/*
+ * 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 <gtest/gtest.h>
+#include <vector>
+
+#include <executorch/backends/aoti/common_shims_slim.h>
+#include <executorch/backends/aoti/slim/c10/core/Device.h>
+#include <executorch/backends/aoti/slim/c10/core/ScalarType.h>
+#include <executorch/backends/cuda/runtime/shims/memory.h>
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/platform/platform.h>
+
+using namespace executorch::backends::cuda;
+using namespace executorch::backends::aoti;
+using executorch::runtime::Error;
+
+namespace slim_c10 = executorch::backends::aoti::slim::c10;
+
+class AOTITorchZeroTest : 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";
+    }
+  }
+
+  Tensor* create(
+      const std::vector<int64_t>& sizes,
+      const std::vector<int64_t>& strides,
+      slim_c10::DeviceType device_type) {
+    Tensor* tensor = nullptr;
+    EXPECT_EQ(
+        aoti_torch_empty_strided(
+            sizes.size(),
+            sizes.data(),
+            strides.data(),
+            static_cast<int32_t>(slim_c10::ScalarType::Float),
+            static_cast<int32_t>(device_type),
+            0,
+            &tensor),
+        Error::Ok);
+    return tensor;
+  }
+};
+
+TEST_F(AOTITorchZeroTest, ZerosContiguousCudaTensor) {
+  Tensor* tensor = create({2, 3}, {3, 1}, slim_c10::DeviceType::CUDA);
+  const std::vector<float> ones(6, 1.0f);
+  ASSERT_EQ(
+      cudaMemcpy(
+          tensor->data_ptr(),
+          ones.data(),
+          sizeof(float) * ones.size(),
+          cudaMemcpyHostToDevice),
+      cudaSuccess);
+
+  EXPECT_EQ(aoti_torch_zero_(tensor), Error::Ok);
+
+  std::vector<float> actual(6, 1.0f);
+  ASSERT_EQ(
+      cudaMemcpy(
+          actual.data(),
+          tensor->data_ptr(),
+          sizeof(float) * actual.size(),
+          cudaMemcpyDeviceToHost),
+      cudaSuccess);
+  EXPECT_EQ(actual, std::vector<float>(6, 0.0f));
+  aoti_torch_delete_tensor_object(tensor);
+}
+
+TEST_F(AOTITorchZeroTest, RejectsNullAndNonContiguousTensors) {
+  EXPECT_EQ(aoti_torch_zero_(nullptr), Error::InvalidArgument);
+  Tensor* tensor = create({2, 3}, {1, 2}, slim_c10::DeviceType::CUDA);
+  EXPECT_EQ(aoti_torch_zero_(tensor), Error::InvalidArgument);
+  aoti_torch_delete_tensor_object(tensor);
+}

extension/cuda/runtime_api.h

@@ -122,6 +122,11 @@ inline cudaError_t cudaMemGetInfo(size_t* free, size_t* total) {
   return hipMemGetInfo(free, total);
 }
 
+inline cudaError_t
+cudaMemsetAsync(void* ptr, int value, size_t size, cudaStream_t stream) {
+  return hipMemsetAsync(ptr, value, size, stream);
+}
+
 inline cudaError_t cudaPointerGetAttributes(
     cudaPointerAttributes* attributes,
     const void* ptr) {