Skip to content

Commit f5bc22a

Browse files
authored
use numel() when appropriate in get_reduced_dim_product (#9142)
No need to re-calculate numel() here.
1 parent 243c885 commit f5bc22a

File tree

1 file changed

+4
-12
lines changed

1 file changed

+4
-12
lines changed

kernels/portable/cpu/util/reduce_util.cpp

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -83,12 +83,8 @@ size_t get_reduced_dim_product(
8383
if (in.dim() == 0) {
8484
return 1;
8585
}
86-
size_t dim_product = 1;
8786
if (!dim.has_value()) {
88-
for (size_t i = 0; i < static_cast<size_t>(in.dim()); ++i) {
89-
dim_product *= in.size(i);
90-
}
91-
return dim_product;
87+
return in.numel();
9288
}
9389
const size_t d = _normalize_non_neg_d(dim.value(), in.dim());
9490
return in.size(d);
@@ -104,16 +100,12 @@ size_t get_reduced_dim_product(
104100
if (in.dim() == 0) {
105101
return 1;
106102
}
107-
size_t dim_product = 1;
108-
const size_t in_dim = in.dim();
109103
if (!dim_list.has_value() || dim_list.value().size() == 0) {
110-
for (size_t i = 0; i < static_cast<size_t>(in.dim()); ++i) {
111-
dim_product *= in.size(i);
112-
}
113-
return dim_product;
104+
return in.numel();
114105
}
106+
size_t dim_product = 1;
115107
for (const auto& d : dim_list.value()) {
116-
const size_t non_neg_d = _normalize_non_neg_d(d, in_dim);
108+
const size_t non_neg_d = _normalize_non_neg_d(d, in.dim());
117109
dim_product *= in.size(non_neg_d);
118110
}
119111
return dim_product;

0 commit comments

Comments
 (0)