Skip to content

Commit 4ac2d6a

Browse files
Gasoonjiafacebook-github-bot
Gasoonjia
authored and
facebook-github-bot
committed
a-start ops | add dim order regulation (#4330)
Summary: Pull Request resolved: #4330 This diff updates the sanity checks on operators starting with a for the dim order regulation. tracking table https://docs.google.com/spreadsheets/d/1Gttxkur8H6QnNfiCGfSAKwtBqdL6MSxn9eJ62bVYS_w/edit?gid=0#gid=0 Differential Revision: D59824508
1 parent 50ab717 commit 4ac2d6a

20 files changed

+206
-84
lines changed

‎kernels/portable/cpu/op_add.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,8 @@ Tensor& add_out(
7979
out);
8080

8181
ET_KERNEL_CHECK(ctx, tensor_is_realhb_type(out), InvalidArgument, out);
82+
ET_KERNEL_CHECK(
83+
ctx, tensors_have_same_dim_order(a, b, out), InvalidArgument, out);
8284

8385
ScalarType a_type = a.scalar_type();
8486
ScalarType b_type = b.scalar_type();
@@ -131,6 +133,8 @@ Tensor& add_scalar_out(
131133
"Failed to resize output tensor.");
132134

133135
ET_KERNEL_CHECK(ctx, tensor_is_realhb_type(out), InvalidArgument, out);
136+
ET_KERNEL_CHECK(
137+
ctx, tensors_have_same_dim_order(a, out), InvalidArgument, out);
134138

135139
ScalarType a_type = a.scalar_type();
136140
ScalarType b_type = utils::get_scalar_dtype(b);

‎kernels/portable/cpu/op_addmm.cpp

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,14 @@ Tensor& addmm_out(
4545
ET_KERNEL_CHECK(
4646
ctx, tensor_is_broadcastable_to(in, out), InvalidArgument, out);
4747

48+
ET_KERNEL_CHECK(
49+
ctx,
50+
tensors_have_same_dim_order(in, mat1, mat2, out),
51+
InvalidArgument,
52+
out);
53+
54+
ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
55+
4856
ScalarType alpha_dtype = utils::get_scalar_dtype(alpha);
4957
ScalarType beta_dtype = utils::get_scalar_dtype(beta);
5058
ET_SWITCH_REAL_TYPES_AND(

‎kernels/portable/cpu/op_alias_copy.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ Tensor& alias_copy_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
2828
"Failed to resize output tensor.");
2929

3030
ET_KERNEL_CHECK(ctx, tensors_have_same_dtype(in, out), InvalidArgument, out);
31+
ET_KERNEL_CHECK(
32+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
3133

3234
if (in.nbytes() > 0) {
3335
// Note that this check is important. It's valid for a tensor with numel 0

‎kernels/portable/cpu/op_allclose.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,9 @@ Tensor& allclose_out(
104104
out.scalar_type() == ScalarType::Bool,
105105
"Out tensor must be type Bool; saw type %" PRId8,
106106
static_cast<int8_t>(out.scalar_type()));
107+
ET_CHECK_MSG(
108+
tensors_have_same_dim_order(self, other, out),
109+
"self, other and out tensors should have same dim order");
107110
ET_CHECK_MSG(
108111
out.numel() == 1,
109112
"Out tensor must be a single element; saw %zu elements",

‎kernels/portable/cpu/op_amax.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,9 @@ Tensor& amax_out(
3939
InvalidArgument,
4040
out);
4141

42+
ET_KERNEL_CHECK(
43+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
44+
4245
ET_SWITCH_REAL_TYPES_AND(
4346
Bool, in.scalar_type(), ctx, "amax.out", CTYPE, [&]() {
4447
CTYPE* out_data = out.mutable_data_ptr<CTYPE>();

‎kernels/portable/cpu/op_amin.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,9 @@ Tensor& amin_out(
3939
InvalidArgument,
4040
out);
4141

42+
ET_KERNEL_CHECK(
43+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
44+
4245
ET_SWITCH_REAL_TYPES_AND(
4346
Bool, in.scalar_type(), ctx, "amin.out", CTYPE, [&]() {
4447
CTYPE* out_data = out.mutable_data_ptr<CTYPE>();

‎kernels/portable/cpu/op_any.cpp

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@ Tensor& any_all_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
2222
ET_KERNEL_CHECK(
2323
ctx, resize_tensor(out, {}) == Error::Ok, InvalidArgument, out);
2424

25+
ET_KERNEL_CHECK(
26+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
27+
2528
ScalarType in_type = in.scalar_type();
2629
ScalarType out_type = out.scalar_type();
2730
constexpr auto name = "any.all_out";
@@ -68,6 +71,9 @@ Tensor& any_dims_out(
6871
out);
6972
}
7073

74+
ET_KERNEL_CHECK(
75+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
76+
7177
ScalarType in_type = in.scalar_type();
7278
ScalarType out_type = out.scalar_type();
7379
constexpr auto name = "any.dims_out";
@@ -122,6 +128,9 @@ Tensor& any_out(
122128
InvalidArgument,
123129
out);
124130

131+
ET_KERNEL_CHECK(
132+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
133+
125134
ScalarType in_type = in.scalar_type();
126135
ScalarType out_type = out.scalar_type();
127136
constexpr auto name = "any.out";

‎kernels/portable/cpu/op_arange.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ Tensor& arange_out(RuntimeContext& ctx, const Scalar& end, Tensor& out) {
2727
ET_KERNEL_CHECK(
2828
ctx, check_arange_args(0.0, end_val, 1.0, out), InvalidArgument, out);
2929

30+
ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(out), InvalidArgument, out);
31+
3032
size_t size = static_cast<size_t>(std::ceil(end_val));
3133

3234
Tensor::SizesType out_length = static_cast<Tensor::SizesType>(size);
@@ -73,6 +75,8 @@ Tensor& arange_start_out(
7375
InvalidArgument,
7476
out);
7577

78+
ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(out), InvalidArgument, out);
79+
7680
double size_d = (d_end - d_start) / d_step;
7781
size_t size = static_cast<size_t>(std::ceil(size_d));
7882

‎kernels/portable/cpu/op_argmax.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@ Tensor& argmax_out(
4040
InvalidArgument,
4141
out);
4242

43+
ET_KERNEL_CHECK(
44+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
45+
4346
ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "argmax.out", CTYPE, [&] {
4447
long* out_data = out.mutable_data_ptr<long>();
4548

‎kernels/portable/cpu/op_argmin.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@ Tensor& argmin_out(
4040
InvalidArgument,
4141
out);
4242

43+
ET_KERNEL_CHECK(
44+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
45+
4346
ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "argmin.out", CTYPE, [&] {
4447
long* out_data = out.mutable_data_ptr<long>();
4548

‎kernels/portable/cpu/op_as_strided_copy.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@ Tensor& as_strided_copy_out(
3737
InvalidArgument,
3838
out);
3939

40+
ET_KERNEL_CHECK(
41+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
42+
43+
ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
44+
4045
if (in.numel() == 0) {
4146
return out;
4247
}

‎kernels/portable/cpu/op_atan2.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@ atan2_out(RuntimeContext& ctx, const Tensor& a, const Tensor& b, Tensor& out) {
2626
InvalidArgument,
2727
out);
2828

29+
ET_KERNEL_CHECK(
30+
ctx, tensors_have_same_dim_order(a, b, out), InvalidArgument, out);
31+
2932
ScalarType a_type = a.scalar_type();
3033
ScalarType b_type = b.scalar_type();
3134
ScalarType out_type = out.scalar_type();

‎kernels/portable/cpu/op_avg_pool2d.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,11 @@ Tensor& avg_pool2d_out(
4444
InvalidArgument,
4545
out);
4646

47+
ET_KERNEL_CHECK(
48+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
49+
50+
ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
51+
4752
size_t output_ndim = 0;
4853
exec_aten::SizesType output_sizes[kTensorDimensionLimit];
4954
get_avg_pool2d_out_target_size(

‎kernels/portable/cpu/pattern/binary_ufunc_realb_realb_to_realb_logical.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ Tensor& binary_ufunc_realb_realb_to_realb_logical(
2727
InvalidArgument,
2828
out);
2929

30+
ET_KERNEL_CHECK(
31+
ctx, tensors_have_same_dim_order(a, b, out), InvalidArgument, out);
32+
3033
ScalarType a_type = a.scalar_type();
3134
ScalarType b_type = b.scalar_type();
3235
ScalarType out_type = out.scalar_type();

‎kernels/portable/cpu/pattern/unary_ufunc_realh.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@ Tensor& unary_ufunc_realh(
3333
ET_KERNEL_CHECK(
3434
ctx, tensors_have_same_shape_and_dtype(in, out), InvalidArgument, out);
3535

36+
ET_KERNEL_CHECK(
37+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
38+
3639
ET_SWITCH_REALH_TYPES(in.scalar_type(), ctx, __func__, CTYPE, [&] {
3740
apply_unary_map_fn(
3841
[fn](const CTYPE val_in) { return static_cast<CTYPE>(fn(val_in)); },

‎kernels/portable/cpu/pattern/unary_ufunc_realhb_to_bool.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@ Tensor& unary_ufunc_realhb_to_bool(
3838
"Expected out tensor to have dtype Bool, but got %" PRId8 " instead.",
3939
static_cast<int8_t>(out.scalar_type()));
4040

41+
ET_KERNEL_CHECK(
42+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
43+
4144
const auto in_type = in.scalar_type();
4245

4346
ET_SWITCH_REALHB_TYPES(in_type, ctx, __func__, CTYPE_IN, [&] {

‎kernels/portable/cpu/pattern/unary_ufunc_realhb_to_floath.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@ Tensor& unary_ufunc_realhb_to_floath(
3232
out,
3333
"Failed to resize output tensor.");
3434

35+
ET_KERNEL_CHECK(
36+
ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
37+
3538
const auto in_type = in.scalar_type();
3639
const auto out_type = out.scalar_type();
3740

‎runtime/core/exec_aten/util/tensor_util.h

Lines changed: 58 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1155,33 +1155,80 @@ bool tensor_has_valid_dim_order(exec_aten::Tensor t);
11551155
bool tensor_is_default_or_channels_last_dim_order(exec_aten::Tensor t);
11561156

11571157
/**
1158-
* Asserts that two tensors have the same dim_order
1158+
* Checks whether a tensor has the default dimension order.
1159+
* Logs an error message if the tensor does not meet the expected criteria.
11591160
*
1160-
* Note that this macro only tests dim order, but not others like actual data,
1161-
* sizes, etc. Also this macro does not support ATen mode since we do not
1162-
* support dim order in ATen mode.
1161+
* @param t The tensor to check the dimension order of.
1162+
* @return True if the tensor has the default dimension order, false otherwise.
1163+
*/
1164+
bool tensor_is_default_dim_order(exec_aten::Tensor t);
1165+
1166+
/**
1167+
* Checks whether a tensor has the channels last dimension order.
1168+
* Logs an error message if the tensor does not meet the expected criteria.
11631169
*
1164-
* TODO(T183094318): Add dim order and related function support for ATen mode.
1170+
* @param t The tensor to check the dimension order of.
1171+
* @return True if the tensor has the channels last dimension order, false
1172+
* otherwise.
11651173
*/
1174+
bool tensor_is_channels_last_dim_order(exec_aten::Tensor t);
11661175

1176+
/**
1177+
* Asserts that four tensors have the same dim_order
1178+
*
1179+
* Note that this macro only tests dim order, but not others like actual data,
1180+
* sizes, etc.
1181+
*
1182+
*/
11671183
bool tensors_have_same_dim_order(
1184+
const exec_aten::ArrayRef<exec_aten::Tensor> tensor_list);
1185+
1186+
/**
1187+
* Asserts that two tensors have the same dim_order
1188+
*
1189+
* Note that this macro only tests dim order, but not others like actual data,
1190+
* sizes, etc.
1191+
*/
1192+
1193+
inline bool tensors_have_same_dim_order(
11681194
const exec_aten::Tensor& a,
1169-
const exec_aten::Tensor& b);
1195+
const exec_aten::Tensor& b) {
1196+
exec_aten::Tensor tensor_list[2] = {a, b};
1197+
return tensors_have_same_dim_order(tensor_list);
1198+
}
11701199

11711200
/**
11721201
* Asserts that three tensors have the same dim_order
11731202
*
11741203
* Note that this macro only tests dim order, but not others like actual data,
1175-
* sizes, etc. Also this macro does not support ATen mode since we do not
1176-
* support dim order in ATen mode.
1204+
* sizes, etc.
11771205
*
1178-
* TODO(T183094318): Add dim order and related function support for ATen mode.
11791206
*/
11801207

1181-
bool tensors_have_same_dim_order(
1208+
inline bool tensors_have_same_dim_order(
11821209
const exec_aten::Tensor& a,
11831210
const exec_aten::Tensor& b,
1184-
const exec_aten::Tensor& c);
1211+
const exec_aten::Tensor& c) {
1212+
exec_aten::Tensor tensor_list[3] = {a, b, c};
1213+
return tensors_have_same_dim_order(tensor_list);
1214+
}
1215+
1216+
/**
1217+
* Asserts that four tensors have the same dim_order
1218+
*
1219+
* Note that this macro only tests dim order, but not others like actual data,
1220+
* sizes, etc.
1221+
*
1222+
*/
1223+
1224+
inline bool tensors_have_same_dim_order(
1225+
const exec_aten::Tensor& a,
1226+
const exec_aten::Tensor& b,
1227+
const exec_aten::Tensor& c,
1228+
const exec_aten::Tensor& d) {
1229+
exec_aten::Tensor tensor_list[4] = {a, b, c, d};
1230+
return tensors_have_same_dim_order(tensor_list);
1231+
}
11851232

11861233
/**
11871234
* Given an n-dimensional coordinate array and an array of tensor strides,

‎runtime/core/exec_aten/util/tensor_util_aten.cpp

Lines changed: 30 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -78,61 +78,43 @@ inline bool tensor_is_default_or_channels_last_dim_order(at::Tensor t) {
7878
}
7979

8080
bool tensors_have_same_dim_order(
81-
const exec_aten::Tensor& a,
82-
const exec_aten::Tensor& b) {
83-
exec_aten::DimOrderType a_dim_order[kTensorDimensionLimit];
84-
exec_aten::DimOrderType b_dim_order[kTensorDimensionLimit];
85-
86-
ET_LOG_MSG_AND_RETURN_IF_FALSE(
87-
get_dim_order(a, a_dim_order, a.dim()) == Error::Ok,
88-
"Failed to retrieve dim order from first input tensor!");
89-
ET_LOG_MSG_AND_RETURN_IF_FALSE(
90-
get_dim_order(b, b_dim_order, b.dim()) == Error::Ok,
91-
"Failed to retrieve dim order from second input tensor!");
92-
93-
bool all_contiguous = is_contiguous_dim_order(a_dim_order, a.dim()) &&
94-
is_contiguous_dim_order(b_dim_order, b.dim());
95-
96-
bool all_channels_last = is_channels_last_dim_order(a_dim_order, a.dim()) &&
97-
is_channels_last_dim_order(b_dim_order, b.dim());
98-
99-
ET_LOG_MSG_AND_RETURN_IF_FALSE(
100-
all_contiguous || all_channels_last,
101-
"Two input tensors have different dim orders");
81+
const exec_aten::ArrayRef<exec_aten::Tensor> tensor_list) {
82+
if (tensor_list.size() < 2) {
83+
return true;
84+
}
10285

103-
return true;
104-
}
86+
exec_aten::DimOrderType first_dim_order[kTensorDimensionLimit];
87+
exec_aten::DimOrderType other_dim_order[kTensorDimensionLimit];
10588

106-
bool tensors_have_same_dim_order(
107-
const exec_aten::Tensor& a,
108-
const exec_aten::Tensor& b,
109-
const exec_aten::Tensor& c) {
110-
exec_aten::DimOrderType a_dim_order[kTensorDimensionLimit];
111-
exec_aten::DimOrderType b_dim_order[kTensorDimensionLimit];
112-
exec_aten::DimOrderType c_dim_order[kTensorDimensionLimit];
113-
ET_LOG_MSG_AND_RETURN_IF_FALSE(
114-
get_dim_order(a, a_dim_order, a.dim()) == Error::Ok,
115-
"Failed to retrieve dim order from first input tensor!");
116-
ET_LOG_MSG_AND_RETURN_IF_FALSE(
117-
get_dim_order(b, b_dim_order, b.dim()) == Error::Ok,
118-
"Failed to retrieve dim order from second input tensor!");
11989
ET_LOG_MSG_AND_RETURN_IF_FALSE(
120-
get_dim_order(c, c_dim_order, c.dim()) == Error::Ok,
121-
"Failed to retrieve dim order from third input tensor!");
122-
123-
bool all_contiguous = is_contiguous_dim_order(a_dim_order, a.dim()) &&
124-
is_contiguous_dim_order(b_dim_order, b.dim()) &&
125-
is_contiguous_dim_order(c_dim_order, c.dim());
126-
127-
bool all_channels_last = is_channels_last_dim_order(a_dim_order, a.dim()) &&
128-
is_channels_last_dim_order(b_dim_order, b.dim()) &&
129-
is_channels_last_dim_order(c_dim_order, c.dim());
90+
get_dim_order(tensor_list[0], first_dim_order, tensor_list[0].dim()) ==
91+
Error::Ok,
92+
"Failed to retrieve dim order from 1st input tensor!");
93+
94+
bool all_contiguous =
95+
is_contiguous_dim_order(first_dim_order, tensor_list[0].dim());
96+
bool all_channels_last =
97+
is_channels_last_dim_order(first_dim_order, tensor_list[0].dim());
98+
99+
for (size_t i = 1; i < tensor_list.size(); ++i) {
100+
ET_LOG_MSG_AND_RETURN_IF_FALSE(
101+
get_dim_order(tensor_list[i], other_dim_order, tensor_list[i].dim()) ==
102+
Error::Ok,
103+
"Failed to retrieve dim order from %zd-th input tensor!",
104+
i);
105+
106+
all_contiguous = all_contiguous &&
107+
is_contiguous_dim_order(other_dim_order, tensor_list[i].dim());
108+
all_channels_last = all_channels_last &&
109+
is_channels_last_dim_order(other_dim_order, tensor_list[i].dim());
110+
}
130111

131112
ET_LOG_MSG_AND_RETURN_IF_FALSE(
132113
all_contiguous || all_channels_last,
133-
"Three input tensors have different dim orders");
114+
"%zd input tensors have different dim orders",
115+
tensor_list.size());
134116

135-
return true;
117+
return all_contiguous || all_channels_last;
136118
}
137119

138120
namespace internal {

0 commit comments

Comments
 (0)