Follow up on pytorch#161891 move additions to stable shim and use version guards (pytorch#168025)

Address pytorch#161891 (comment)

Pull Request resolved: pytorch#168025
Approved by: https://github.com/janeyx99

Author: mikaylagawarecki

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

@@ -204,16 +204,6 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_data_ptr(
     void** ret_data_ptr // returns borrowed reference
 );
 
-AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_mutable_data_ptr(
-    AtenTensorHandle tensor,
-    void** ret_data_ptr // returns borrowed reference
-);
-
-AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_const_data_ptr(
-    AtenTensorHandle tensor,
-    const void** ret_data_ptr // returns borrowed reference
-);
-
 // Get the nbytes of the underlying storage
 AOTI_TORCH_EXPORT AOTITorchError
 aoti_torch_get_storage_size(AtenTensorHandle tensor, int64_t* ret_size);

torch/csrc/inductor/aoti_torch/shim_common.cpp

@@ -281,24 +281,6 @@ AOTITorchError aoti_torch_get_data_ptr(
   });
 }
 
-AOTITorchError aoti_torch_get_const_data_ptr(
-    AtenTensorHandle tensor,
-    const void** ret_data_ptr) {
-  AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
-    at::Tensor* t = tensor_handle_to_tensor_pointer(tensor);
-    *ret_data_ptr = t->const_data_ptr();
-  });
-}
-
-AOTITorchError aoti_torch_get_mutable_data_ptr(
-    AtenTensorHandle tensor,
-    void** ret_data_ptr) {
-  AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
-    at::Tensor* t = tensor_handle_to_tensor_pointer(tensor);
-    *ret_data_ptr = t->mutable_data_ptr();
-  });
-}
-
 AOTITorchError aoti_torch_get_storage_size(
     AtenTensorHandle tensor,
     int64_t* ret_size) {

torch/csrc/shim_common.cpp

@@ -560,3 +560,21 @@ torch_get_num_threads(uint32_t* out_num_threads) {
   AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE(
       { *out_num_threads = static_cast<uint32_t>(at::get_num_threads()); });
 }
+
+AOTI_TORCH_EXPORT AOTITorchError
+torch_get_const_data_ptr(AtenTensorHandle tensor, const void** ret_data_ptr) {
+  AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
+    at::Tensor* t =
+        torch::aot_inductor::tensor_handle_to_tensor_pointer(tensor);
+    *ret_data_ptr = t->const_data_ptr();
+  });
+}
+
+AOTI_TORCH_EXPORT AOTITorchError
+torch_get_mutable_data_ptr(AtenTensorHandle tensor, void** ret_data_ptr) {
+  AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
+    at::Tensor* t =
+        torch::aot_inductor::tensor_handle_to_tensor_pointer(tensor);
+    *ret_data_ptr = t->mutable_data_ptr();
+  });
+}

torch/csrc/stable/c/shim.h

@@ -92,6 +92,17 @@ AOTI_TORCH_EXPORT AOTITorchError torch_get_thread_idx(uint32_t* out_thread_idx);
 AOTI_TORCH_EXPORT AOTITorchError
 torch_get_num_threads(uint32_t* out_num_threads);
 
+// Get a pointer to the underlying storage data
+AOTI_TORCH_EXPORT AOTITorchError torch_get_mutable_data_ptr(
+    AtenTensorHandle tensor,
+    void** ret_data_ptr // returns borrowed reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError torch_get_const_data_ptr(
+    AtenTensorHandle tensor,
+    const void** ret_data_ptr // returns borrowed reference
+);
+
 #endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
 
 #ifdef __cplusplus

torch/csrc/stable/tensor_inl.h

@@ -33,6 +33,7 @@ inline Device Tensor::device() const {
   return Device(extension_device_type, static_cast<DeviceIndex>(device_index));
 }
 
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
 // The following data ptr cast methods mirror the methods defined in
 // aten/src/ATen/templates/TensorMethods.cpp
 #define DEFINE_DATA_PTR_CAST(T, name, PRED)               \
@@ -63,5 +64,6 @@ DEFINE_CAST(uint32_t, UInt32)
 DEFINE_CAST(uint64_t, UInt64)
 #undef DEFINE_CAST
 #undef _PRED
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
 
 HIDDEN_NAMESPACE_END(torch, stable)

torch/csrc/stable/tensor_struct.h

@@ -86,17 +86,16 @@ class Tensor {
     return data_ptr;
   }
 
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
   void* mutable_data_ptr() const {
     void* data_ptr{};
-    TORCH_ERROR_CODE_CHECK(
-        aoti_torch_get_mutable_data_ptr(ath_.get(), &data_ptr));
+    TORCH_ERROR_CODE_CHECK(torch_get_mutable_data_ptr(ath_.get(), &data_ptr));
     return data_ptr;
   }
 
   const void* const_data_ptr() const {
     const void* data_ptr{};
-    TORCH_ERROR_CODE_CHECK(
-        aoti_torch_get_const_data_ptr(ath_.get(), &data_ptr));
+    TORCH_ERROR_CODE_CHECK(torch_get_const_data_ptr(ath_.get(), &data_ptr));
     return data_ptr;
   }
 
@@ -105,6 +104,7 @@ class Tensor {
 
   template <typename T, std::enable_if_t<!std::is_const_v<T>, int> = 0>
   const T* const_data_ptr() const;
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
 
   int64_t dim() const {
     int64_t dim;