Skip to content

Commit 785066e

Browse files
committed
MKLDNN elementwise_mul: Check if AVX512 is available
test=develop
1 parent 08f63c4 commit 785066e

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

paddle/fluid/operators/elementwise_mul_mkldnn_op.cc

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,10 +136,13 @@ class ElementwiseMulMKLDNNKernel : public framework::OpKernel<T> {
136136
UpdateDataFormat(ctx, (Tensor*)x, "x_data_format");
137137
UpdateDataFormat(ctx, (Tensor*)y, "y_data_format");
138138

139+
Xbyak::util::Cpu cpu;
140+
const bool is_avx512_enabled = cpu.has(Xbyak::util::Cpu::tAVX512F);
139141
const bool are_dims_divisable = !(x_int_dims[1] % 16);
140142
const bool is_x_format_correct = x->format() == memory::format::nChw16c;
141143
const bool is_y_format_correct = y->format() == memory::format::nc;
142-
if (is_x_format_correct && is_y_format_correct && are_dims_divisable) {
144+
if (is_x_format_correct && is_y_format_correct && are_dims_divisable &&
145+
is_avx512_enabled) {
143146
int pre, n, post;
144147
get_mid_dims(x_dims, y_dims_untrimmed, axis, &pre, &n, &post);
145148

0 commit comments

Comments
 (0)