Add beginnings of torch::stable::accelerator #159679

mikaylagawarecki · 2025-08-01T21:37:22Z

Adds

torch::stable::accelerator::DeviceGuard: std::unique_ptr to DeviceGuardOpauqe mostly copied from the below (but made generic)

pytorch/torch/csrc/inductor/aoti_runtime/utils_cuda.h

Lines 30 to 46 in 50eac81

    
           class AOTICudaGuard { 
        
            public: 
        
             AOTICudaGuard(int32_t device_index) : guard_(nullptr, delete_cuda_guard) { 
        
               CUDAGuardHandle ptr = nullptr; 
        
               AOTI_TORCH_ERROR_CODE_CHECK( 
        
                   aoti_torch_create_cuda_guard(device_index, &ptr)); 
        
               guard_.reset(ptr); 
        
             } 
        
             void set_index(int32_t device_index) { 
        
               AOTI_TORCH_ERROR_CODE_CHECK( 
        
                   aoti_torch_cuda_guard_set_index(guard_.get(), device_index)); 
        
             } 
        
            private: 
        
             std::unique_ptr<CUDAGuardOpaque, DeleterFnPtr> guard_; 
        
           };

constructor DeviceGuard(DeviceIndex) (this matches aoti but defers from the actual c10 DeviceGuard constructor that takes in device)
set_index(DeviceIndex)

torch::stable::accelerator::Stream: std::shared_ptr to StreamOpaque
- constructor Stream(StreamHandle stream) (similar to torch::stable::Tensor)
- id() -> StreamId
getCurrentStream(DeviceIndex device_index) -> stable::accelerator::Stream

Stack from ghstack (oldest at bottom):

[ghstack-poisoned]

pytorch-bot · 2025-08-01T21:37:26Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159679

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 5e8e158 with merge base 34ec5ed ():

NEW FAILURE - The following job has failed:

pull / linux-jammy-py3.13-clang12 / test (default, 2, 5, linux.4xlarge) (gh)
cpp/test_jit 1/1 failed!

This comment was automatically generated by Dr. CI and updates every 15 minutes.

ghstack-source-id: 1530b24 Pull Request resolved: #159679

github-actions · 2025-08-01T21:41:54Z

Attention! PyTorch one of the C-stable API file was changed

You MUST NOT change existing function declarations in this, as this header defines a stable C ABI. If you need to change the signature for a function, introduce a new v2 version of the function and modify code generation to target the new version of the function.

Caused by:

torch/csrc/inductor/aoti_torch/c/shim.h

[ghstack-poisoned]

ghstack-source-id: 2f2fa7b Pull Request resolved: #159679

[ghstack-poisoned]

ghstack-source-id: fae00fc Pull Request resolved: #159679

[ghstack-poisoned]

ghstack-source-id: 083b8bc Pull Request resolved: #159679

mikaylagawarecki · 2025-08-04T19:17:42Z

torch/csrc/stable/accelerator.h

+
+using DeviceIndex = int8_t;
+using StreamId = int64_t;
+class DeviceGuard {


something that I'm not sure about -- do the copy / move semantics need to match the real device guard although this just a wrapper that is holding a std::unique_ptr to DeviceGuardOpaque?

pytorch/c10/core/DeviceGuard.h

Lines 37 to 46 in 444e238

~DeviceGuard() = default;

/// Copy is disallowed

DeviceGuard(const DeviceGuard&) = delete;

DeviceGuard& operator=(const DeviceGuard&) = delete;

/// Move is disallowed, as DeviceGuard does not have an uninitialized state,

/// which is required for moves on types with nontrivial destructors.

DeviceGuard(DeviceGuard&& other) = delete;

DeviceGuard& operator=(DeviceGuard&& other) = delete;

It doesn't need to copy the existing DeviceGuard. For example, I agree with deleting the default constructor for this API. We can disallow copy and move if it's easiest to maintain anyway.

[ghstack-poisoned]

ghstack-source-id: afc07a8 Pull Request resolved: #159679

[ghstack-poisoned]

ghstack-source-id: 23b8be1 Pull Request resolved: #159679

Adds - `torch::stable::accelerator::DeviceGuard`: `std::unique_ptr` to `DeviceGuardOpauqe` copied from https://github.com/pytorch/pytorch/blob/50eac811a68e63e96ad56c11c983bfe298a0bb8a/torch/csrc/inductor/aoti_runtime/utils_cuda.h#L30-L46 - constructor `DeviceGuard(DeviceIndex)` (this matches aoti but defers from the actual c10 DeviceGuard constructor that takes in device - `set_index(DeviceIndex)` - `torch::stable::accelerator::Stream`: `std::shared_ptr` to `StreamOpaque` - constructor `Stream(StreamHandle stream)` (similar to torch::stable::Tensor) - `id() -> StreamId` - `getCurrentStream(DeviceIndex device_index) -> stable::accelerator::Stream` [ghstack-poisoned]

ghstack-source-id: 655868b Pull Request resolved: #159679

Adds - `torch::stable::accelerator::DeviceGuard`: `std::unique_ptr` to `DeviceGuardOpauqe` copied from https://github.com/pytorch/pytorch/blob/50eac811a68e63e96ad56c11c983bfe298a0bb8a/torch/csrc/inductor/aoti_runtime/utils_cuda.h#L30-L46 - constructor `DeviceGuard(DeviceIndex)` (this matches aoti but defers from the actual c10 DeviceGuard constructor that takes in device - `set_index(DeviceIndex)` - `torch::stable::accelerator::Stream`: `std::shared_ptr` to `StreamOpaque` - constructor `Stream(StreamHandle stream)` (similar to torch::stable::Tensor) - `id() -> StreamId` - `getCurrentStream(DeviceIndex device_index) -> stable::accelerator::Stream` [ghstack-poisoned]

ghstack-source-id: 94eaf15 Pull Request resolved: #159679

albanD · 2025-08-04T22:07:36Z

FYI @guangyey @EikanWang in case you have some feedback here.

Adds - `torch::stable::accelerator::DeviceGuard`: `std::unique_ptr` to `DeviceGuardOpauqe` copied from https://github.com/pytorch/pytorch/blob/50eac811a68e63e96ad56c11c983bfe298a0bb8a/torch/csrc/inductor/aoti_runtime/utils_cuda.h#L30-L46 - constructor `DeviceGuard(DeviceIndex)` (this matches aoti but defers from the actual c10 DeviceGuard constructor that takes in device - `set_index(DeviceIndex)` - `torch::stable::accelerator::Stream`: `std::shared_ptr` to `StreamOpaque` - constructor `Stream(StreamHandle stream)` (similar to torch::stable::Tensor) - `id() -> StreamId` - `getCurrentStream(DeviceIndex device_index) -> stable::accelerator::Stream` [ghstack-poisoned]

ghstack-source-id: 2ede272 Pull Request resolved: #159679

guangyey

Thanks.

test/cpp_extensions/libtorch_agnostic_extension/libtorch_agnostic/csrc/kernel.cpp

test/cpp_extensions/libtorch_agnostic_extension/libtorch_agnostic/ops.py

janeyx99 · 2025-08-06T18:42:59Z

test/cpp_extensions/libtorch_agnostic_extension/setup.py

+    if torch.cuda.is_available():
+        extra_compile_args["cxx"].append("-DUSE_CUDA")
+        extension = CUDAExtension
+


Hmmm maybe this would be a good call to move the CUDA stuff into its own C++ file and build them separately..

Hmmm, I don't see the issue with the current approach, so gonna keep it as is unless there's something specific you're concerned about! :)

No issues yet, but if we have more "cuda only" code in the future it may be a good way to split the ever growing kernel.cpp

janeyx99 · 2025-08-06T18:49:14Z

torch/csrc/inductor/aoti_torch/shim_common.cpp

  });
 }
+
+AOTITorchError aoti_torch_create_device_guard(


I'm guessing this code is inspired by https://github.com/pytorch/pytorch/blob/main/torch/csrc/inductor/aoti_torch/shim_cuda.cpp#L8?

janeyx99 · 2025-08-06T18:53:17Z

torch/csrc/inductor/aoti_torch/shim_common.cpp

+    c10::Stream* stream_ptr = new c10::Stream(stream);
+    *ret_stream = reinterpret_cast<StreamHandle>(stream_ptr);


Creating a new stream on the heap is gonna leak memory, we should be able to assign the stream in line 1642 into the pointer (tho idk the cast semantics).

Ahhh after reading the rest of the code, I see what the user flow is intended to be. That said, I'm not sure if we want the semantics of get_current_stream to create a new stream on the heap (this is likely not expected from the user POV), and if we do end up sticking with this, we need to loudly document this as.

I think we do not want to mess with the memory of the stream at all...cc @albanD for thoughts

Hmm looks like @albanD had no thoughts :b

I'm not sure what to do here, this one looks different from

pytorch/torch/csrc/inductor/aoti_torch/shim_cuda.cpp

Lines 49 to 55 in 5f5f508

AOTITorchError aoti_torch_get_current_cuda_stream(

int32_t device_index,

void** ret_stream) {

AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({

*(cudaStream_t*)(ret_stream) = at::cuda::getCurrentCUDAStream(device_index);

});

}

As that is just directly returning the id. However, the actual getCurrentStream returns a c10::Stream, so I'm trying to match the semantic in the stable ABI (so in future users can add other stream methods on a stream object)

pytorch/aten/src/ATen/DeviceAccelerator.cpp

Lines 106 to 110 in 2ee22e4

c10::Stream getCurrentStream(c10::DeviceIndex device_index) {

const auto device_type = getAccelerator(true).value();

c10::impl::VirtualGuardImpl impl(device_type);

return impl.getStream({device_type, device_index});

}

I think at::acceerator::getCurrentStream is returning c10::Stream by value which is why we will need to create a new object on the heap if we want to return it to the caller and transfer ownership to the stable::Stream

I can add a comment for sure, but would be curious how else we can tweak this in a way that's more user/memory friendly, would it make more sense to just return the .id() directly?

test/cpp_extensions/libtorch_agnostic_extension/libtorch_agnostic/csrc/kernel.cpp

janeyx99 · 2025-08-06T19:00:24Z

torch/csrc/stable/accelerator.h

+
+using DeviceIndex = int8_t;
+using StreamId = int64_t;
+class DeviceGuard {


It doesn't need to copy the existing DeviceGuard. For example, I agree with deleting the default constructor for this API. We can disallow copy and move if it's easiest to maintain anyway.

janeyx99 · 2025-08-06T19:05:04Z

torch/csrc/stable/accelerator.h

+
+  // Construct a stable::Stream from a StreamHandle
+  // Steals ownership from the StreamHandle
+  explicit Stream(StreamHandle stream)


Ah, I see, so the expected use case is:

user calls getCurrentStream which creates a Stream

then they can call id on it

janeyx99 · 2025-08-12T18:39:17Z

@pytorchbot merge -i

pytorchmergebot · 2025-08-12T18:41:24Z

Merge started

Your change will be merged while ignoring the following 4 checks: pull / linux-jammy-py3.9-clang12 / test (default, 4, 5, linux.4xlarge), pull / linux-jammy-py3.9-gcc11 / test (default, 4, 5, linux.2xlarge), inductor / unit-test / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_amx, 1, 2, linux.8xlarge.amx), inductor / unit-test / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_avx2, 1, 2, linux.10xlarge.avx2)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-08-12T18:42:19Z

The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command
For more information see pytorch-bot wiki.

mikaylagawarecki · 2025-08-12T20:02:22Z

torch/csrc/stable/accelerator.h

@@ -0,0 +1,71 @@
+#pragma once
+
+#include <torch/csrc/inductor/aoti_runtime/utils.h>


note to self to remove when rebasing and use TORCH_ERROR_CODE_CHECK instead

torch/csrc/stable/accelerator.h

Adds - `torch::stable::accelerator::DeviceGuard`: `std::unique_ptr` to `DeviceGuardOpauqe` mostly copied from the below (but made generic) https://github.com/pytorch/pytorch/blob/50eac811a68e63e96ad56c11c983bfe298a0bb8a/torch/csrc/inductor/aoti_runtime/utils_cuda.h#L30-L46 - constructor `DeviceGuard(DeviceIndex)` (**this matches aoti but defers from the actual c10 DeviceGuard constructor that takes in device**) - `set_index(DeviceIndex)` - `torch::stable::accelerator::Stream`: `std::shared_ptr` to `StreamOpaque` - constructor `Stream(StreamHandle stream)` (similar to torch::stable::Tensor) - `id() -> StreamId` - `getCurrentStream(DeviceIndex device_index) -> stable::accelerator::Stream` [ghstack-poisoned]

pytorchmergebot · 2025-08-13T16:16:38Z