Eliminate sizes, strides, mutable_data_ptr, const_data_ptr ops

Author: pearu

src/libtorchaudio/rnnt/cpu/compute.cpp

@@ -114,9 +114,10 @@ std::tuple<Tensor, Tensor> compute(
   // when stable ABI Tensor supports mutable_data_ptr templates.
   Workspace<float> workspace(
       /*options=*/options,
-      /*dtype_data=*/reinterpret_cast<float*>(float_workspace.data_ptr()),
+      /*dtype_data=*/
+      reinterpret_cast<float*>(float_workspace.mutable_data_ptr()),
       /*dtype_size=*/float_workspace.numel(),
-      /*int_data=*/reinterpret_cast<int*>(int_workspace.data_ptr()),
+      /*int_data=*/reinterpret_cast<int*>(int_workspace.mutable_data_ptr()),
       /*int_size=*/int_workspace.numel());
 
   THO_DISPATCH_V2(

src/libtorchaudio/stable/ops.h

@@ -23,46 +23,11 @@ namespace torchaudio::stable {
 
 using Layout = int32_t;
 
-// TODO: When sizes and strides are implemented in torch::stable,
-// eliminate sizes and strides function below.
-inline std::vector<int64_t> sizes(const Tensor& t) {
-  int64_t* ptr;
-  TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes(t.get(), &ptr));
-  std::vector<int64_t> r(ptr, ptr + t.dim());
-  return r;
-}
-
-inline std::vector<int64_t> strides(const Tensor& t) {
-  int64_t* ptr;
-  TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(t.get(), &ptr));
-  std::vector<int64_t> r(ptr, ptr + t.dim());
-  return r;
-}
-
-// TODO: When https://github.com/pytorch/pytorch/pull/161891 lands,
-// eliminate mutable_data_ptr and const_data_ptr templates.
-#define aoti_torch_get_mutable_data_ptr aoti_torch_get_data_ptr
-#define aoti_torch_get_const_data_ptr aoti_torch_get_data_ptr
-template <typename T>
-T* mutable_data_ptr(const Tensor& t) {
-  void* data_ptr{};
-  TORCH_ERROR_CODE_CHECK(aoti_torch_get_mutable_data_ptr(t.get(), &data_ptr));
-  return reinterpret_cast<T*>(data_ptr);
-}
-
-template <typename T>
-const T* const_data_ptr(const Tensor& t) {
-  const void* data_ptr{};
-  TORCH_ERROR_CODE_CHECK(
-      aoti_torch_get_const_data_ptr(t.get(), const_cast<void**>(&data_ptr)));
-  return reinterpret_cast<const T*>(data_ptr);
-}
-
 // TODO: When cpu is implemented in torch::stable, eliminate
 // cpu function below.
 inline Tensor cpu(const Tensor& self) {
-  auto sizes_ = sizes(self);
-  auto cpu_type = aoti_torch_device_type_cpu();
+  auto sizes_ = self.sizes();
+  int32_t cpu_type = static_cast<int32_t>(torch::stable::DeviceType::CPU);
   int32_t dtype;
   TORCH_ERROR_CODE_CHECK(aoti_torch_get_dtype(self.get(), &dtype));
   int32_t layout;
@@ -85,8 +50,8 @@ inline Tensor cpu(const Tensor& self) {
 
 // TODO:
 inline Tensor cuda(const Tensor& self, int32_t cuda_index) {
-  auto sizes_ = sizes(self);
-  auto cuda_type = aoti_torch_device_type_cuda();
+  auto sizes_ = self.sizes();
+  int32_t cuda_type = static_cast<int32_t>(torch::stable::DeviceType::CUDA);
   int32_t dtype;
   TORCH_ERROR_CODE_CHECK(aoti_torch_get_dtype(self.get(), &dtype));
   int32_t layout;
@@ -169,7 +134,9 @@ T item(const Tensor& self) {
   STD_TORCH_CHECK(
       self.numel() == 1, "item requires single element tensor input");
   if (self.is_cpu()) {
-    return torchaudio::stable::const_data_ptr<T>(self)[0];
+    // TODO: use `return self.const_data_ptr<T>()[0];` after torch
+    // stable supports const_data_ptr templates.
+    return reinterpret_cast<const T*>(self.const_data_ptr())[0];
 #ifdef USE_CUDA
   } else if (self.is_cuda()) {
     T value;