From 69783e3d024988bbee11f907a4bd3ba3b130bd73 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 4 Sep 2024 13:34:50 -0700 Subject: [PATCH] Adopt the new tensor API for aten_util. (#5062) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/5062 . Reviewed By: dbort Differential Revision: D62168422 --- CMakeLists.txt | 1 + .../make_aten_functor_from_et_functor.h | 58 +++++-------------- extension/aten_util/targets.bzl | 1 + extension/llm/custom_ops/CMakeLists.txt | 3 +- extension/tensor/CMakeLists.txt | 2 +- extension/tensor/tensor_impl_ptr.cpp | 2 +- 6 files changed, 20 insertions(+), 47 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 3618bff7677..a19f405e80c 100644 --- a/CMakeLists.txt +++ b/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() diff --git a/extension/aten_util/make_aten_functor_from_et_functor.h b/extension/aten_util/make_aten_functor_from_et_functor.h index 3b54254e8ed..d7f2906944c 100644 --- a/extension/aten_util/make_aten_functor_from_et_functor.h +++ b/extension/aten_util/make_aten_functor_from_et_functor.h @@ -20,8 +20,8 @@ #endif #include #include +#include #include -#include #include namespace executorch { @@ -105,37 +105,12 @@ struct type_convert< typename remove_const_ref::type, torch::executor::Tensor>>> final { - explicit type_convert(ATensor value) : value_(value) { - auto sizes = - std::make_shared>( - value_.sizes().begin(), value_.sizes().end()); - const ssize_t dim = sizes->size(); - auto dim_order = - std::make_shared>( - dim); - auto strides = - std::make_shared>( - 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( - static_cast(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>( - 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_; @@ -143,10 +118,7 @@ struct type_convert< private: ATensor value_; - std::unique_ptr< - torch::executor::Tensor, - std::function> - converted_; + TensorPtr converted_; }; // Tensors: ETen to ATen. @@ -158,15 +130,14 @@ struct type_convert< std::is_same_v::type, at::Tensor> && std::is_same_v< typename remove_const_ref::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(value_.scalar_type())); - } + : value_(value), + converted_(at::from_blob( + value_.mutable_data_ptr(), + std::vector{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 sizes_; }; // Optionals: ATen to ETen. diff --git a/extension/aten_util/targets.bzl b/extension/aten_util/targets.bzl index b396cb78325..f219d6253f2 100644 --- a/extension/aten_util/targets.bzl +++ b/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", diff --git a/extension/llm/custom_ops/CMakeLists.txt b/extension/llm/custom_ops/CMakeLists.txt index 41c8c0ee160..723444498a4 100644 --- a/extension/llm/custom_ops/CMakeLists.txt +++ b/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 diff --git a/extension/tensor/CMakeLists.txt b/extension/tensor/CMakeLists.txt index 4a02965c647..2cf1bf2956f 100644 --- a/extension/tensor/CMakeLists.txt +++ b/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}) diff --git a/extension/tensor/tensor_impl_ptr.cpp b/extension/tensor/tensor_impl_ptr.cpp index aa5f78e7f8d..ea4d83f5afd 100644 --- a/extension/tensor/tensor_impl_ptr.cpp +++ b/extension/tensor/tensor_impl_ptr.cpp @@ -91,7 +91,7 @@ TensorImplPtr make_tensor_impl_ptr( tensor_impl.release(), TensorImplPtrDeleter{ std::unique_ptr>( - data, std::move(deleter) ?: noop_deleter), + data, deleter ? std::move(deleter) : noop_deleter), std::move(sizes), std::move(dim_order), std::move(strides)});