Skip to content

Commit 432e251

Browse files
swolchokswolchok
committed
Remove ExecuTorch copy of Vectorized
Pull Request resolved: #7042 All uses are outside ExecuTorch core, so we can just use ATen Vectorized. ghstack-source-id: 289279745 @exported-using-ghexport Differential Revision: [D66396016](https://our.internmc.facebook.com/intern/diff/D66396016/)
1 parent ef1d2ff commit 432e251

32 files changed

+212
-5561
lines changed

‎extension/llm/custom_ops/CMakeLists.txt‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,10 +70,11 @@ else()
7070
endif()
7171

7272
add_library(custom_ops ${_custom_ops__srcs})
73-
73+
find_package_torch_headers()
7474
target_include_directories(custom_ops PUBLIC "${_common_include_directories}")
7575
target_include_directories(
7676
custom_ops PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/../../../include"
77+
${TORCH_INCLUDE_DIRS}
7778
)
7879
target_link_libraries(custom_ops PUBLIC ${custom_ops_libs} executorch_core)
7980

@@ -99,6 +100,7 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
99100
)
100101
target_include_directories(
101102
custom_ops_aot_lib PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/../../../include"
103+
${TORCH_INCLUDE_DIRS}
102104
)
103105
# TODO: This only works if we install portable_lib.so to
104106
# <site-packages>/executorch/extension/pybindings/.

‎extension/llm/custom_ops/op_sdpa.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,9 @@
99
#include <executorch/extension/llm/custom_ops/op_sdpa.h>
1010
#include <executorch/extension/llm/custom_ops/op_sdpa_impl.h>
1111

12+
#include <ATen/cpu/vec/functional.h>
13+
#include <ATen/cpu/vec/vec.h>
1214
#include <executorch/kernels/optimized/blas/CPUBlas.h>
13-
#include <executorch/kernels/optimized/vec/functional.h>
14-
#include <executorch/kernels/optimized/vec/vec.h>
1515
#include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
1616
// @lint-ignore CLANGTIDY facebook-unused-include-check
1717
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>

‎extension/llm/custom_ops/op_sdpa_impl.h‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@
88

99
#pragma once
1010

11+
#include <ATen/cpu/vec/vec.h>
1112
#include <executorch/kernels/optimized/blas/CPUBlas.h>
1213
#include <executorch/kernels/optimized/vec/functional.h>
13-
#include <executorch/kernels/optimized/vec/vec.h>
1414
#include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
1515
// @lint-ignore CLANGTIDY facebook-unused-include-check
1616
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
@@ -319,7 +319,7 @@ void _qk_at_v_gemm(
319319
constexpr size_t kKVDim = 4;
320320

321321
template <typename T>
322-
inline void _store(T* dst, ::executorch::vec::Vectorized<T> src) {
322+
inline void _store(T* dst, ::at::vec::Vectorized<T> src) {
323323
src.store(dst);
324324
}
325325

@@ -354,7 +354,7 @@ inline double calculate_scale(const Tensor& query, optional<double> scale) {
354354
return softmax_scale;
355355
}
356356

357-
namespace vec = ::executorch::vec;
357+
namespace vec = ::at::vec;
358358
using Tensor = ::executorch::aten::Tensor;
359359

360360
// 1) out = exp(a - val)

‎kernels/optimized/cpu/moments_utils.h‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
// for use in optimized ExecuTorch ops. Template specializations of BFloat16
1313
// are excluded.
1414

15-
#include <executorch/kernels/optimized/vec/vec.h>
15+
#include <ATen/cpu/vec/vec.h>
1616

1717
#include <executorch/kernels/optimized/utils/math_utils.h>
1818
#include <executorch/runtime/platform/compiler.h>
@@ -47,12 +47,12 @@ void AddMoments(
4747
template <typename T>
4848
ET_INLINE void AddMomentsVec(
4949
int64_t m0_add,
50-
const executorch::vec::Vectorized<T>& m1_add,
51-
const executorch::vec::Vectorized<T>& m2_add,
50+
const at::vec::Vectorized<T>& m1_add,
51+
const at::vec::Vectorized<T>& m2_add,
5252
int64_t& m0,
53-
executorch::vec::Vectorized<T>& m1,
54-
executorch::vec::Vectorized<T>& m2) {
55-
using Vec = executorch::vec::Vectorized<T>;
53+
at::vec::Vectorized<T>& m1,
54+
at::vec::Vectorized<T>& m2) {
55+
using Vec = at::vec::Vectorized<T>;
5656
const int64_t n = m0 + m0_add;
5757
const T c =
5858
n == 0 ? static_cast<T>(0) : static_cast<T>(m0_add) / static_cast<T>(n);
@@ -67,11 +67,11 @@ template <typename T>
6767
inline void UpdateMomentsVec(
6868
int64_t m0,
6969
const T* X_ptr,
70-
const std::array<executorch::vec::Vectorized<acc_t<T>>, kChunkSize>& c_vecs,
70+
const std::array<at::vec::Vectorized<acc_t<T>>, kChunkSize>& c_vecs,
7171
int64_t& m0_stk0,
72-
executorch::vec::Vectorized<acc_t<T>>& m1_stk0,
73-
executorch::vec::Vectorized<acc_t<T>>& m2_stk0) {
74-
using Vec = executorch::vec::Vectorized<acc_t<T>>;
72+
at::vec::Vectorized<acc_t<T>>& m1_stk0,
73+
at::vec::Vectorized<acc_t<T>>& m2_stk0) {
74+
using Vec = at::vec::Vectorized<acc_t<T>>;
7575
Vec m1_vec(0);
7676
Vec m2_vec(0);
7777
for (int64_t j = 0; j < m0; ++j) {
@@ -92,13 +92,13 @@ std::pair<acc_t<T>, acc_t<T>>
9292
RowwiseMomentsImpl(const T* X, int64_t N, int64_t ddof = 0) {
9393
using T_ACC = acc_t<T>;
9494

95-
constexpr int64_t kVecSize = executorch::vec::Vectorized<T>::size();
96-
constexpr int64_t kAccVecSize = executorch::vec::Vectorized<T_ACC>::size();
95+
constexpr int64_t kVecSize = at::vec::Vectorized<T>::size();
96+
constexpr int64_t kAccVecSize = at::vec::Vectorized<T_ACC>::size();
9797
const int64_t n = N / kVecSize;
9898
const int64_t m = executorch::utils::divup(n, kChunkSize);
9999
const int64_t depth = executorch::utils::CeilLog2(m);
100100

101-
using Vec = executorch::vec::Vectorized<T_ACC>;
101+
using Vec = at::vec::Vectorized<T_ACC>;
102102
const Vec kZeroVec(T_ACC(0));
103103
std::array<int64_t, kMaxDepth> m0_stk;
104104
std::array<Vec, kMaxDepth> m1_stk;
@@ -168,7 +168,7 @@ RowwiseMomentsImpl(const T* X, int64_t N, int64_t ddof = 0) {
168168
template <typename T>
169169
std::pair<acc_t<T>, acc_t<T>>
170170
RowwiseMoments(const T* X, int64_t N, int64_t ddof = 0) {
171-
using Vec = executorch::vec::Vectorized<T>;
171+
using Vec = at::vec::Vectorized<T>;
172172
constexpr int64_t kVecSize = Vec::size();
173173
const int64_t n = N / kVecSize;
174174
const int64_t m = executorch::utils::divup(n, kChunkSize);

‎kernels/optimized/cpu/op_add.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,9 @@
66
* LICENSE file in the root directory of this source tree.
77
*/
88

9+
#include <ATen/cpu/vec/functional.h>
10+
#include <ATen/cpu/vec/vec.h>
911
#include <executorch/kernels/optimized/cpu/binary_ops.h>
10-
#include <executorch/kernels/optimized/vec/functional.h>
11-
#include <executorch/kernels/optimized/vec/vec.h>
1212
#include <executorch/kernels/portable/cpu/scalar_utils.h>
1313
#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
1414
#include <executorch/runtime/kernel/kernel_includes.h>
@@ -51,8 +51,8 @@ Tensor& opt_add_out(
5151
CTYPE_B b_val = *b.const_data_ptr<CTYPE_B>();
5252
CTYPE b_casted = static_cast<CTYPE>(b_val);
5353

54-
using Vec = executorch::vec::Vectorized<CTYPE>;
55-
executorch::vec::map<CTYPE>(
54+
using Vec = at::vec::Vectorized<CTYPE>;
55+
at::vec::map<CTYPE>(
5656
[alpha_val, b_casted](Vec x) {
5757
return x + Vec(alpha_val * b_casted);
5858
},
@@ -106,8 +106,8 @@ Tensor& opt_add_scalar_out(
106106
CTYPE alpha_val;
107107
ET_EXTRACT_SCALAR(alpha, alpha_val);
108108

109-
using Vec = executorch::vec::Vectorized<CTYPE>;
110-
executorch::vec::map<CTYPE>(
109+
using Vec = at::vec::Vectorized<CTYPE>;
110+
at::vec::map<CTYPE>(
111111
[alpha_val, b_casted](Vec x) {
112112
return x + Vec(alpha_val * b_casted);
113113
},

‎kernels/optimized/cpu/op_add_sub_impl.h‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,9 @@
66
* LICENSE file in the root directory of this source tree.
77
*/
88

9+
#include <ATen/cpu/vec/functional.h>
10+
#include <ATen/cpu/vec/vec.h>
911
#include <executorch/kernels/optimized/cpu/binary_ops.h>
10-
#include <executorch/kernels/optimized/vec/functional.h>
11-
#include <executorch/kernels/optimized/vec/vec.h>
1212
#include <executorch/kernels/portable/cpu/scalar_utils.h>
1313
#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
1414
#include <executorch/runtime/kernel/kernel_includes.h>
@@ -104,8 +104,8 @@ Tensor& opt_add_sub_out_impl(
104104
if constexpr (is_sub) {
105105
alpha_val = -alpha_val;
106106
}
107-
using Vec = executorch::vec::Vectorized<CTYPE>;
108-
executorch::vec::map2<CTYPE>(
107+
using Vec = at::vec::Vectorized<CTYPE>;
108+
at::vec::map2<CTYPE>(
109109
[alpha_val](Vec x, Vec y) { return x + Vec(alpha_val) * y; },
110110
out.mutable_data_ptr<CTYPE>(),
111111
a.const_data_ptr<CTYPE>(),
@@ -123,7 +123,7 @@ Tensor& opt_add_sub_out_impl(
123123
InvalidArgument,
124124
out,
125125
"Failed to extract scalar alpha.");
126-
using Vec = executorch::vec::Vectorized<CTYPE>;
126+
using Vec = at::vec::Vectorized<CTYPE>;
127127
Vec alpha_val_vec(alpha_val);
128128
if constexpr (is_sub) {
129129
if (selected_optimized_path ==

‎kernels/optimized/cpu/op_div.cpp‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,9 @@
66
* LICENSE file in the root directory of this source tree.
77
*/
88

9+
#include <ATen/cpu/vec/functional.h>
10+
#include <ATen/cpu/vec/vec.h>
911
#include <executorch/kernels/optimized/cpu/binary_ops.h>
10-
#include <executorch/kernels/optimized/vec/functional.h>
11-
#include <executorch/kernels/optimized/vec/vec.h>
1212
#include <executorch/kernels/portable/cpu/scalar_utils.h>
1313
#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
1414
#include <executorch/runtime/kernel/kernel_includes.h>
@@ -76,16 +76,16 @@ Tensor& opt_div_out(
7676
CTYPE_SCALAR scalar_val = *scalar->const_data_ptr<CTYPE_SCALAR>();
7777
CTYPE scalar_casted = static_cast<CTYPE>(scalar_val);
7878

79-
using Vec = executorch::vec::Vectorized<CTYPE>;
79+
using Vec = at::vec::Vectorized<CTYPE>;
8080
if (a.numel() == 1) {
81-
executorch::vec::map<CTYPE>(
81+
at::vec::map<CTYPE>(
8282
[scalar_casted](Vec x) { return Vec(scalar_casted) / x; },
8383
out.mutable_data_ptr<CTYPE>(),
8484
tensor->const_data_ptr<CTYPE>(),
8585
out.numel());
8686
} else {
8787
Vec inv_scalar_casted_vec(CTYPE(1) / scalar_casted);
88-
executorch::vec::map<CTYPE>(
88+
at::vec::map<CTYPE>(
8989
[inv_scalar_casted_vec](Vec x) {
9090
return x * inv_scalar_casted_vec;
9191
},
@@ -111,8 +111,8 @@ Tensor& opt_div_out(
111111
"Failed to resize output tensor.");
112112

113113
ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, "div.out", CTYPE, [&]() {
114-
using Vec = executorch::vec::Vectorized<CTYPE>;
115-
executorch::vec::map2<CTYPE>(
114+
using Vec = at::vec::Vectorized<CTYPE>;
115+
at::vec::map2<CTYPE>(
116116
[](Vec x, Vec y) { return x / y; },
117117
out.mutable_data_ptr<CTYPE>(),
118118
a.const_data_ptr<CTYPE>(),
@@ -198,9 +198,9 @@ Tensor& opt_div_scalar_out(
198198
ET_EXTRACT_SCALAR(b, b_val);
199199
CTYPE b_casted = static_cast<CTYPE>(b_val);
200200

201-
using Vec = executorch::vec::Vectorized<CTYPE>;
201+
using Vec = at::vec::Vectorized<CTYPE>;
202202
Vec inv_b_casted_vec(CTYPE(1) / b_casted);
203-
executorch::vec::map<CTYPE>(
203+
at::vec::map<CTYPE>(
204204
[inv_b_casted_vec](Vec x) { return x * inv_b_casted_vec; },
205205
out.mutable_data_ptr<CTYPE>(),
206206
a.const_data_ptr<CTYPE>(),

‎kernels/optimized/cpu/op_exp.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@
88

99
#include <cmath>
1010

11-
#include <executorch/kernels/optimized/vec/functional.h>
12-
#include <executorch/kernels/optimized/vec/vec.h>
11+
#include <ATen/cpu/vec/functional.h>
12+
#include <ATen/cpu/vec/vec.h>
1313
#include <executorch/runtime/kernel/kernel_includes.h>
1414

1515
namespace torch {
@@ -34,8 +34,8 @@ void exp_data(
3434
const CTYPE_IN* in_data,
3535
const size_t numel,
3636
CTYPE_OUT* out_data) {
37-
using Vec = executorch::vec::Vectorized<CTYPE_IN>;
38-
executorch::vec::map<CTYPE_IN>(
37+
using Vec = at::vec::Vectorized<CTYPE_IN>;
38+
at::vec::map<CTYPE_IN>(
3939
[](Vec x) { return x.exp(); }, out_data, in_data, numel);
4040
}
4141

‎kernels/optimized/cpu/op_le.cpp‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66
* LICENSE file in the root directory of this source tree.
77
*/
88

9-
#include <executorch/kernels/optimized/vec/functional.h>
10-
#include <executorch/kernels/optimized/vec/vec.h>
9+
#include <ATen/cpu/vec/functional.h>
10+
#include <ATen/cpu/vec/vec.h>
1111
#include <executorch/kernels/portable/cpu/scalar_utils.h>
1212
#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
1313
#include <executorch/runtime/kernel/kernel_includes.h>
@@ -60,15 +60,15 @@ Tensor& opt_le_tensor_out(
6060
CTYPE_SCALAR scalar_val = *scalar->const_data_ptr<CTYPE_SCALAR>();
6161
CTYPE scalar_casted = static_cast<CTYPE>(scalar_val);
6262

63-
using Vec = executorch::vec::Vectorized<CTYPE>;
63+
using Vec = at::vec::Vectorized<CTYPE>;
6464
if (a.numel() == 1) {
65-
executorch::vec::map<CTYPE>(
65+
at::vec::map<CTYPE>(
6666
[scalar_casted](Vec x) { return Vec(scalar_casted).le(x); },
6767
out.mutable_data_ptr<CTYPE>(),
6868
tensor->const_data_ptr<CTYPE>(),
6969
out.numel());
7070
} else {
71-
executorch::vec::map<CTYPE>(
71+
at::vec::map<CTYPE>(
7272
[scalar_casted](Vec x) { return x.le(Vec(scalar_casted)); },
7373
out.mutable_data_ptr<CTYPE>(),
7474
tensor->const_data_ptr<CTYPE>(),
@@ -93,8 +93,8 @@ Tensor& opt_le_tensor_out(
9393
if (a_type == b_type && a_type == out_type) {
9494
ET_SWITCH_REAL_TYPES_AND(
9595
Bool, out_type, ctx, "le.Tensor_out", CTYPE, [&]() {
96-
using Vec = executorch::vec::Vectorized<CTYPE>;
97-
executorch::vec::map2<CTYPE>(
96+
using Vec = at::vec::Vectorized<CTYPE>;
97+
at::vec::map2<CTYPE>(
9898
[](Vec x, Vec y) { return x.le(y); },
9999
out.mutable_data_ptr<CTYPE>(),
100100
a.const_data_ptr<CTYPE>(),
@@ -158,8 +158,8 @@ Tensor& opt_le_scalar_out(
158158
CTYPE_B b_val = 0;
159159
ET_EXTRACT_SCALAR(b, b_val);
160160
CTYPE b_casted = static_cast<CTYPE>(b_val);
161-
using Vec = executorch::vec::Vectorized<CTYPE>;
162-
executorch::vec::map<CTYPE>(
161+
using Vec = at::vec::Vectorized<CTYPE>;
162+
at::vec::map<CTYPE>(
163163
[b_casted](Vec x) { return x.le(Vec(b_casted)); },
164164
out.mutable_data_ptr<CTYPE>(),
165165
a.const_data_ptr<CTYPE>(),

‎kernels/optimized/cpu/op_linear.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@
1010

1111
#include <c10/util/irange.h>
1212

13+
#include <ATen/cpu/vec/functional.h>
14+
#include <ATen/cpu/vec/vec.h>
1315
#include <executorch/kernels/optimized/blas/CPUBlas.h>
14-
#include <executorch/kernels/optimized/vec/functional_base.h>
15-
#include <executorch/kernels/optimized/vec/vec_base.h>
1616
#include <executorch/kernels/portable/cpu/util/matmul_ops_util.h>
1717
#include <executorch/runtime/kernel/kernel_includes.h>
1818

@@ -21,12 +21,12 @@ namespace executor {
2121
namespace native {
2222

2323
namespace {
24+
using ::at::vec::map;
25+
using ::at::vec::Vectorized;
2426
using ::executorch::aten::Tensor;
2527
using ::executorch::cpublas::gemm;
2628
using ::executorch::cpublas::TransposeType;
2729
using ::executorch::runtime::toString;
28-
using ::executorch::vec::map;
29-
using ::executorch::vec::Vectorized;
3030

3131
// Use vector store to initialize with scalar bias.
3232
template <typename scalar_t>

0 commit comments

Comments
 (0)