Skip to content

Commit 07cd102

Browse files
committed
Update on "[Exutorch] Add broadcast support for le op"
For refactored hf repro requires this to support mask generation Differential Revision: [D76456398](https://our.internmc.facebook.com/intern/diff/D76456398/) [ghstack-poisoned]
1 parent d4c8f29 commit 07cd102

File tree

2 files changed

+41
-42
lines changed

2 files changed

+41
-42
lines changed

kernels/optimized/cpu/op_le.cpp

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,11 @@
99
#include <executorch/kernels/optimized/cpu/binary_ops.h>
1010
#include <executorch/kernels/optimized/vec/functional.h>
1111
#include <executorch/kernels/optimized/vec/vec.h>
12+
#include <executorch/kernels/portable/cpu/pattern/comparison_op.h>
1213
#include <executorch/kernels/portable/cpu/scalar_utils.h>
1314
#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
1415
#include <executorch/runtime/kernel/kernel_includes.h>
1516
#include <executorch/runtime/platform/assert.h>
16-
#include <executorch/kernels/portable/cpu/pattern/comparison_op.h>
1717

1818
namespace torch {
1919
namespace executor {
@@ -83,7 +83,9 @@ Tensor& opt_le_tensor_out(
8383

8484
// Check for optimized broadcast paths
8585
auto selected_optimized_path = select_optimized_path(a, b, out);
86-
printf("selected_optimized_path: %d\n", static_cast<int>(selected_optimized_path));
86+
printf(
87+
"selected_optimized_path: %d\n",
88+
static_cast<int>(selected_optimized_path));
8789
if (selected_optimized_path == ElementwiseOptimizedPath::kTreatAs1d) {
8890
// Resize for dynamic shape
8991
auto error = resize_tensor(out, a.sizes());

kernels/test/op_le_test.cpp

Lines changed: 37 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -204,7 +204,9 @@ TEST_F(OpLeTensorOutTest, Broadcast2DTest) {
204204
// Row 0: b[0]=2, so [1<=2, 2<=2, 3<=2, ...] = [true, true, false, false, ...]
205205
// Row 1: b[1]=4, so [1<=4, 2<=4, 3<=4, 4<=4, 5<=4, ...] = [true, true, true,
206206
// true, false, ...]
207-
using ctype = executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<ScalarType::Bool>::ctype;
207+
using ctype =
208+
executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<
209+
ScalarType::Bool>::ctype;
208210
std::vector<ctype> expected_data = {
209211
// Row 0 (b=2): 1<=2, 2<=2, 3<=2, 4<=2, 5<=2, 6<=2, 7<=2, 8<=2, 9<=2,
210212
// 10<=2
@@ -279,10 +281,7 @@ TEST_F(OpLeTensorOutTest, Broadcast2DTest) {
279281
true,
280282
true};
281283

282-
EXPECT_TENSOR_EQ(
283-
out,
284-
tf_bool.make(
285-
{6, 10}, expected_data));
284+
EXPECT_TENSOR_EQ(out, tf_bool.make({6, 10}, expected_data));
286285
}
287286

288287
TEST_F(OpLeTensorOutTest, Broadcast1DTo2DTest) {
@@ -298,7 +297,9 @@ TEST_F(OpLeTensorOutTest, Broadcast1DTo2DTest) {
298297

299298
op_le_tensor_out(a, b, out);
300299

301-
using ctype = executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<ScalarType::Bool>::ctype;
300+
using ctype =
301+
executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<
302+
ScalarType::Bool>::ctype;
302303
std::vector<ctype> expected_data = {
303304
// Row 0 (a=2): 2<=1, 2<=2, 2<=3, 2<=4, 2<=5, 2<=6, 2<=7, 2<=8, 2<=9,
304305
// 2<=10
@@ -373,10 +374,7 @@ TEST_F(OpLeTensorOutTest, Broadcast1DTo2DTest) {
373374
false,
374375
false};
375376

376-
EXPECT_TENSOR_EQ(
377-
out,
378-
tf_bool.make(
379-
{6, 10}, expected_data));
377+
EXPECT_TENSOR_EQ(out, tf_bool.make({6, 10}, expected_data));
380378
}
381379

382380
TEST_F(OpLeTensorOutTest, BroadcastReverseTest) {
@@ -393,7 +391,9 @@ TEST_F(OpLeTensorOutTest, BroadcastReverseTest) {
393391
op_le_tensor_out(a, b, out);
394392

395393
// Expected: each row i should be [a[i]<=1, a[i]<=2, ..., a[i]<=10]
396-
using ctype = executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<ScalarType::Bool>::ctype;
394+
using ctype =
395+
executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<
396+
ScalarType::Bool>::ctype;
397397
std::vector<ctype> expected_data = {
398398
// Row 0 (a=2): 2<=1, 2<=2, 2<=3, 2<=4, 2<=5, 2<=6, 2<=7, 2<=8, 2<=9,
399399
// 2<=10
@@ -468,10 +468,7 @@ TEST_F(OpLeTensorOutTest, BroadcastReverseTest) {
468468
false,
469469
false};
470470

471-
EXPECT_TENSOR_EQ(
472-
out,
473-
tf_bool.make(
474-
{6, 10}, expected_data));
471+
EXPECT_TENSOR_EQ(out, tf_bool.make({6, 10}, expected_data));
475472
}
476473

477474
TEST_F(OpLeTensorOutTest, BroadcastLastDimTest) {
@@ -553,7 +550,9 @@ TEST_F(OpLeTensorOutTest, BroadcastLastDimTest) {
553550

554551
op_le_tensor_out(a, b, out);
555552

556-
using ctype = executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<ScalarType::Bool>::ctype;
553+
using ctype =
554+
executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<
555+
ScalarType::Bool>::ctype;
557556
std::vector<ctype> expected_data = {
558557
// First slice: a values are 1,2,3,4
559558
true,
@@ -617,12 +616,10 @@ TEST_F(OpLeTensorOutTest, BroadcastLastDimTest) {
617616
true,
618617
true,
619618
true,
620-
true// 12 <= [12,13,14,15,16]
619+
true // 12 <= [12,13,14,15,16]
621620
};
622621

623-
EXPECT_TENSOR_EQ(
624-
out,
625-
tf_bool.make({3, 4, 5}, expected_data));
622+
EXPECT_TENSOR_EQ(out, tf_bool.make({3, 4, 5}, expected_data));
626623
}
627624

628625
TEST_F(OpLeTensorOutTest, BroadcastLastDimReverseTest) {
@@ -706,7 +703,9 @@ TEST_F(OpLeTensorOutTest, BroadcastLastDimReverseTest) {
706703

707704
op_le_tensor_out(a, b, out);
708705

709-
using ctype = executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<ScalarType::Bool>::ctype;
706+
using ctype =
707+
executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<
708+
ScalarType::Bool>::ctype;
710709
std::vector<ctype> expected_data = {
711710
// First slice: b values are all 5
712711
true,
@@ -773,9 +772,7 @@ TEST_F(OpLeTensorOutTest, BroadcastLastDimReverseTest) {
773772
false // [12,13,14,15,16] <= 15
774773
};
775774

776-
EXPECT_TENSOR_EQ(
777-
out,
778-
tf_bool.make({3, 4, 5}, expected_data));
775+
EXPECT_TENSOR_EQ(out, tf_bool.make({3, 4, 5}, expected_data));
779776
}
780777

781778
TEST_F(OpLeTensorOutTest, BroadcastNdByNdTest) {
@@ -819,7 +816,9 @@ TEST_F(OpLeTensorOutTest, BroadcastNdByNdTest) {
819816

820817
op_le_tensor_out(a, b, out);
821818

822-
using ctype = executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<ScalarType::Bool>::ctype;
819+
using ctype =
820+
executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<
821+
ScalarType::Bool>::ctype;
823822
std::vector<ctype> expected_data = {
824823
// First slice: a[0,0,:] = [1,2,3,4]
825824
true,
@@ -849,9 +848,7 @@ TEST_F(OpLeTensorOutTest, BroadcastNdByNdTest) {
849848
true // [5,6,7,8] <= [7,8,9,10]
850849
};
851850

852-
EXPECT_TENSOR_EQ(
853-
out,
854-
tf_bool.make({2, 3, 4}, expected_data));
851+
EXPECT_TENSOR_EQ(out, tf_bool.make({2, 3, 4}, expected_data));
855852
}
856853

857854
TEST_F(OpLeTensorOutTest, BroadcastNdByNdReverseTest) {
@@ -896,7 +893,9 @@ TEST_F(OpLeTensorOutTest, BroadcastNdByNdReverseTest) {
896893

897894
op_le_tensor_out(a, b, out);
898895

899-
using ctype = executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<ScalarType::Bool>::ctype;
896+
using ctype =
897+
executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<
898+
ScalarType::Bool>::ctype;
900899
std::vector<ctype> expected_data = {
901900
// First slice: b[0,0,:] = [2,3,4,5]
902901
true,
@@ -926,9 +925,7 @@ TEST_F(OpLeTensorOutTest, BroadcastNdByNdReverseTest) {
926925
false // [7,8,9,10] <= [6,7,8,9]
927926
};
928927

929-
EXPECT_TENSOR_EQ(
930-
out,
931-
tf_bool.make({2, 3, 4}, expected_data));
928+
EXPECT_TENSOR_EQ(out, tf_bool.make({2, 3, 4}, expected_data));
932929
}
933930

934931
TEST_F(OpLeTensorOutTest, Broadcast2dBy1dTest) {
@@ -944,7 +941,9 @@ TEST_F(OpLeTensorOutTest, Broadcast2dBy1dTest) {
944941

945942
op_le_tensor_out(a, b, out);
946943

947-
using ctype = executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<ScalarType::Bool>::ctype;
944+
using ctype =
945+
executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<
946+
ScalarType::Bool>::ctype;
948947
std::vector<ctype> expected_data = {
949948
true,
950949
true,
@@ -960,9 +959,7 @@ TEST_F(OpLeTensorOutTest, Broadcast2dBy1dTest) {
960959
false // [9,10,11,12] <= [2,4,6,8]
961960
};
962961

963-
EXPECT_TENSOR_EQ(
964-
out,
965-
tf_bool.make({3, 4}, expected_data));
962+
EXPECT_TENSOR_EQ(out, tf_bool.make({3, 4}, expected_data));
966963
}
967964

968965
TEST_F(OpLeTensorOutTest, Broadcast2dBy1dReverseTest) {
@@ -978,7 +975,9 @@ TEST_F(OpLeTensorOutTest, Broadcast2dBy1dReverseTest) {
978975

979976
op_le_tensor_out(a, b, out);
980977

981-
using ctype = executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<ScalarType::Bool>::ctype;
978+
using ctype =
979+
executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<
980+
ScalarType::Bool>::ctype;
982981
std::vector<ctype> expected_data = {
983982
false,
984983
false,
@@ -994,7 +993,5 @@ TEST_F(OpLeTensorOutTest, Broadcast2dBy1dReverseTest) {
994993
true // [2,4,6,8] <= [9,10,11,12]
995994
};
996995

997-
EXPECT_TENSOR_EQ(
998-
out,
999-
tf_bool.make({3, 4}, expected_data));
996+
EXPECT_TENSOR_EQ(out, tf_bool.make({3, 4}, expected_data));
1000997
}

0 commit comments

Comments
 (0)