Add aten::record_stream (#1047)

- [x] record_stream
intel · Nov 6, 2024 · 597d8e4 · 597d8e4
1 parent fda86cf
commit 597d8e4
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/src/ATen/native/xpu/sycl/RecordStream.cpp b/src/ATen/native/xpu/sycl/RecordStream.cpp
@@ -0,0 +1,16 @@
+#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+#include <ATen/core/Tensor.h>
+#include <c10/xpu/XPUCachingAllocator.h>
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/NativeFunctions.h>
+#else
+#include <ATen/ops/record_stream_native.h>
+#endif
+
+namespace at::native {
+void record_stream_xpu(Tensor& self, c10::Stream stream) {
+  struct c10::StreamData3 data = stream.pack3();
+  c10::xpu::XPUCachingAllocator::recordStream(self.storage().data_ptr(), at::xpu::XPUStream::unpack3(data.stream_id, data.device_index, data.device_type));
+}
+}  // namespace at::native
diff --git a/yaml/native/native_functions.yaml b/yaml/native/native_functions.yaml
@@ -7164,6 +7164,11 @@
 - func: index_copy.dimname(Tensor self, Dimname dim, Tensor index, Tensor source) -> Tensor
   variants: function, method
 
+- func: record_stream(Tensor(a!) self, Stream s) -> ()
+  variants: method
+  dispatch:
+    XPU: record_stream_xpu
+
 - func: i0(Tensor self) -> Tensor
   structured_delegate: i0.out
   variants: function, method