Adopt the new tensor API for aten_util.

Differential Revision: D62168422

Pull Request resolved: #5062

Author: shoumikhin

CMakeLists.txt

@@ -228,6 +228,7 @@ cmake_dependent_option(
 )
 
 if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
+  set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON)
   set(EXECUTORCH_BUILD_KERNELS_CUSTOM ON)
 endif()
 

extension/aten_util/make_aten_functor_from_et_functor.h

@@ -20,8 +20,8 @@
 #endif
 #include <ATen/native/Resize.h>
 #include <executorch/extension/kernel_util/type_list.h>
+#include <executorch/extension/tensor/tensor.h>
 #include <executorch/runtime/core/evalue.h>
-#include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
 #include <torch/torch.h>
 
 namespace executorch {
@@ -105,48 +105,20 @@ struct type_convert<
             typename remove_const_ref<ETensor>::type,
             torch::executor::Tensor>>>
     final {
-  explicit type_convert(ATensor value) : value_(value) {
-    auto sizes =
-        std::make_shared<std::vector<torch::executor::Tensor::SizesType>>(
-            value_.sizes().begin(), value_.sizes().end());
-    const ssize_t dim = sizes->size();
-    auto dim_order =
-        std::make_shared<std::vector<torch::executor::Tensor::DimOrderType>>(
-            dim);
-    auto strides =
-        std::make_shared<std::vector<torch::executor::Tensor::StridesType>>(
-            dim);
-
-    std::iota(dim_order->begin(), dim_order->end(), 0);
-    ::executorch::runtime::dim_order_to_stride_nocheck(
-        sizes->data(), dim_order->data(), dim, strides->data());
-
-    auto tensor_impl = std::make_shared<torch::executor::TensorImpl>(
-        static_cast<torch::executor::ScalarType>(value_.scalar_type()),
-        sizes->size(),
-        sizes->data(),
-        value_.mutable_data_ptr(),
-        dim_order->data(),
-        strides->data());
-
-    converted_ = std::unique_ptr<
-        torch::executor::Tensor,
-        std::function<void(torch::executor::Tensor*)>>(
-        new torch::executor::Tensor(tensor_impl.get()),
-        [sizes, dim_order, strides, tensor_impl](
-            torch::executor::Tensor* pointer) { delete pointer; });
-  }
+  explicit type_convert(ATensor value)
+      : value_(value),
+        converted_(from_blob(
+            value_.mutable_data_ptr(),
+            {value_.sizes().begin(), value_.sizes().end()},
+            ::torch::executor::ScalarType(value_.scalar_type()))) {}
 
   ETensor call() {
     return *converted_;
   }
 
  private:
   ATensor value_;
-  std::unique_ptr<
-      torch::executor::Tensor,
-      std::function<void(torch::executor::Tensor*)>>
-      converted_;
+  TensorPtr converted_;
 };
 
 // Tensors: ETen to ATen.
@@ -158,15 +130,14 @@ struct type_convert<
         std::is_same_v<typename remove_const_ref<ATensor>::type, at::Tensor> &&
         std::is_same_v<
             typename remove_const_ref<ETensor>::type,
-            torch::executor::Tensor>>>
+            ::torch::executor::Tensor>>>
     final {
   explicit type_convert(ETensor value)
-      : value_(value), sizes_(value_.sizes().begin(), value_.sizes().end()) {
-    converted_ = at::from_blob(
-        value_.mutable_data_ptr(),
-        sizes_,
-        static_cast<c10::ScalarType>(value_.scalar_type()));
-  }
+      : value_(value),
+        converted_(at::from_blob(
+            value_.mutable_data_ptr(),
+            std::vector<int64_t>{value_.sizes().begin(), value_.sizes().end()},
+            c10::ScalarType(value_.scalar_type()))) {}
 
   ATensor call() {
     return converted_;
@@ -175,7 +146,6 @@ struct type_convert<
  private:
   ETensor value_;
   at::Tensor converted_;
-  std::vector<int64_t> sizes_;
 };
 
 // Optionals: ATen to ETen.

extension/aten_util/targets.bzl

@@ -27,6 +27,7 @@ def define_common_targets():
         ],
         exported_deps = [
             "//executorch/extension/kernel_util:kernel_util",
+            "//executorch/extension/tensor:tensor",
             "//executorch/runtime/core:core",
             "//executorch/runtime/core:evalue",
             "//executorch/runtime/core/exec_aten:lib",

extension/llm/custom_ops/CMakeLists.txt

@@ -94,7 +94,8 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
   endif()
 
   target_link_libraries(
-    custom_ops_aot_lib PUBLIC cpublas torch extension_threadpool
+    custom_ops_aot_lib PUBLIC cpublas torch extension_tensor
+                              extension_threadpool
   )
   if(WIN32)
     # There is no direct replacement for libpthread.so on Windows. For the

extension/tensor/CMakeLists.txt

@@ -18,7 +18,7 @@ endif()
 
 list(TRANSFORM _extension_tensor__srcs PREPEND "${EXECUTORCH_ROOT}/")
 add_library(extension_tensor ${_extension_tensor__srcs})
-target_link_libraries(extension_tensor executorch)
+target_link_libraries(extension_tensor executorch_no_prim_ops)
 target_include_directories(extension_tensor PUBLIC ${EXECUTORCH_ROOT}/..)
 target_compile_options(extension_tensor PUBLIC ${_common_compile_options})
 

extension/tensor/tensor_impl_ptr.cpp

@@ -91,7 +91,7 @@ TensorImplPtr make_tensor_impl_ptr(
       tensor_impl.release(),
       TensorImplPtrDeleter{
           std::unique_ptr<void, std::function<void(void*)>>(
-              data, std::move(deleter) ?: noop_deleter),
+              data, deleter ? std::move(deleter) : noop_deleter),
           std::move(sizes),
           std::move(dim_order),
           std::move(strides)});