Skip to content

Commit 8895573

Browse files
pytorchbotswolchok
pytorchbot
and
swolchok
authored
Remove ExecuTorch copy of Vectorized (#11613)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #7042 by @swolchok ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/swolchok/121/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/swolchok/121/head Merge bot PR base: https://github.com/pytorch/executorch/tree/main Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/swolchok/121/orig @diff-train-skip-merge Co-authored-by: Scott Wolchok <[email protected]>
1 parent 402f421 commit 8895573

34 files changed

+241
-5572
lines changed

‎.ci/scripts/test_llava.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ run_and_verify() {
147147

148148
# verify result.txt
149149
RESULT=$(cat result.txt)
150-
EXPECTED_PREFIX="ASSISTANT: image captures a basketball game in progress, with several players on the court. "
150+
EXPECTED_PREFIX="ASSISTANT: image captures a basketball game in progress, with"
151151

152152
if [[ "${RESULT}" == *"${EXPECTED_PREFIX}"* ]]; then
153153
echo "Expected result prefix: ${EXPECTED_PREFIX}"

‎.ci/scripts/unittest-buck2.sh

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,10 @@ buck2 query "//backends/apple/... + //backends/example/... + \
1515
//kernels/optimized/... + //kernels/portable/... + //kernels/quantized/... + \
1616
//kernels/test/... + //runtime/... + //schema/... + //test/... + //util/..."
1717

18+
# TODO: optimized ops are unbuildable because they now use ATen; put
19+
# them back after we can use PyTorch in OSS buck.
1820
UNBUILDABLE_OPTIMIZED_OPS_REGEX="_elu|gelu|fft|log_softmax"
19-
BUILDABLE_OPTIMIZED_OPS=$(buck2 query //kernels/optimized/cpu/... | grep -E -v $UNBUILDABLE_OPTIMIZED_OPS_REGEX)
21+
BUILDABLE_OPTIMIZED_OPS= #$(buck2 query //kernels/optimized/cpu/... | grep -E -v $UNBUILDABLE_OPTIMIZED_OPS_REGEX)
2022

2123
# TODO: build prim_ops_test_cpp again once supported_features works in
2224
# OSS buck.

‎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: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,10 @@
88

99
#pragma once
1010

11+
#include <ATen/cpu/vec/vec.h>
12+
#include <ATen/cpu/vec/vec_n.h>
1113
#include <executorch/kernels/optimized/blas/CPUBlas.h>
1214
#include <executorch/kernels/optimized/vec/functional.h>
13-
#include <executorch/kernels/optimized/vec/vec.h>
1415
#include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
1516
// @lint-ignore CLANGTIDY facebook-unused-include-check
1617
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
@@ -319,7 +320,7 @@ void _qk_at_v_gemm(
319320
constexpr size_t kKVDim = 4;
320321

321322
template <typename T>
322-
inline void _store(T* dst, ::executorch::vec::Vectorized<T> src) {
323+
inline void _store(T* dst, ::at::vec::Vectorized<T> src) {
323324
src.store(dst);
324325
}
325326

@@ -356,30 +357,45 @@ inline double calculate_scale(
356357
return softmax_scale;
357358
}
358359

359-
namespace vec = ::executorch::vec;
360+
namespace vec = ::at::vec;
360361
using Tensor = ::executorch::aten::Tensor;
361362

362363
// 1) out = exp(a - val)
363364
// 2) val = sum(out)
364365
template <typename T1, typename T2>
365366
inline void
366367
_exp_reduce_sum_fusion_kernel(T1* a, const int& size, T2* out, T1& val) {
367-
auto vec_size = vec::Vectorized<T1>::size();
368-
auto vec_max = vec::Vectorized<T1>(val);
368+
// NOTE: we observed numerics issues with this function when
369+
// deleting the old executorch::vec and replacing with at::vec
370+
// here. The major known difference is that executorch::vec was 256
371+
// bits wide vs 128 bits for at::vec (and the hardware). Preserving
372+
// this function's execution width at 256 bits and avoiding
373+
// vec_reduce_all below removed the issues.
374+
constexpr auto vec_size = vec::Vectorized<T1>::size() * 2;
375+
auto vec_max = vec::VectorizedN<T1, 2>(val);
369376
T1 tmp_sum = 0;
370-
auto vec_tmp_sum = vec::Vectorized<T1>(tmp_sum);
377+
auto vec_tmp_sum = vec::VectorizedN<T1, 2>(tmp_sum);
371378
for (int i = 0; i < vec_size * (size / vec_size); i += vec_size) {
372-
auto tmp0 = vec::Vectorized<T1>::loadu(a + i);
379+
auto tmp0 = vec::VectorizedN<T1, 2>::loadu(a + i);
373380
auto tmp1 = tmp0 - vec_max;
374381
// Replace with exp_u20 later
375382
// auto tmp2 = tmp1.exp_u20();
376383
auto tmp2 = tmp1.exp();
377-
vec_tmp_sum += tmp2;
378-
_store(out + i, tmp2);
384+
vec_tmp_sum = vec_tmp_sum + tmp2;
385+
tmp2.store(out + i);
379386
}
380-
tmp_sum = vec::vec_reduce_all<T1>(
381-
[](vec::Vectorized<T1>& x, vec::Vectorized<T1>& y) { return x + y; },
382-
vec_tmp_sum);
387+
388+
__at_align__ T1 vec_tmp_sum_array[vec_size];
389+
vec_tmp_sum.store(vec_tmp_sum_array);
390+
for (const auto i : c10::irange(vec_size)) {
391+
tmp_sum += vec_tmp_sum_array[i];
392+
}
393+
// See NOTE above; we should replace the scalar reduction above with
394+
// this reduction (which uses vaddvq_f32 internally), but it changes
395+
// numerics.
396+
// tmp_sum = vec::vec_reduce_all<T1>(
397+
// [](vec::Vectorized<T1>& x, vec::Vectorized<T1>& y) { return x + y; },
398+
// vec_tmp_sum);
383399
for (int i = vec_size * (size / vec_size); i < size; i++) {
384400
auto tmp0 = a[i];
385401
auto tmp1 = tmp0 - 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>(),

0 commit comments

Comments
 (0)