MCPcopy Create free account
hub / github.com/PaddlePaddle/Paddle / HasOptimizedOneDNNKernel

Function HasOptimizedOneDNNKernel

paddle/phi/kernels/reduce_kernel_impl.cc:23–55  ·  view source on GitHub ↗

oneDNN's reduction kernel is optimized only for reducing throughout the most outer dims, so in case of another type of reduction, it would be better to fallback to native implementation

Source from the content-addressed store, hash-verified

21// most outer dims, so in case of another type of reduction, it would be
22// better to fallback to native implementation
23inline bool HasOptimizedOneDNNKernel(const KernelContext* dev_ctx) {
24 const DenseTensor& x = dev_ctx->InputAt<DenseTensor>(0);
25 IntArray dims_array;
26 const TensorRef& dims_tmp = dev_ctx->AttrAt<TensorRef>(0);
27 dims_array = IntArray(*dims_tmp.Get());
28 int ndims = x.dims().size();
29 const bool reduce_all = recompute_reduce_all(x, dims_array);
30 auto dims = dims_array.GetData();
31
32 // native reduce kernels don't support bf16
33 // so oneDNN kernel is enforced in that case
34 if (x.dtype() == phi::DataType::BFLOAT16) return true;
35
36 if (reduce_all) {
37 return true;
38 }
39
40 for (auto& dim : dims) {
41 if (dim < 0) {
42 dim += ndims;
43 }
44 }
45
46 sort(dims.begin(), dims.end());
47
48 for (size_t i = 0; i < dims.size(); ++i) {
49 if (dims[dims.size() - i - 1] != static_cast<int>(ndims - i - 1)) {
50 return false;
51 }
52 }
53
54 return true;
55}
56
57bool ReduceCheckIfOneDNNSupport(const KernelContext* dev_ctx) {
58 if (dev_ctx->InputAt<DenseTensor>(0).dims().size() > 5 ||

Callers 3

GetExpectedKernelTypeMethod · 0.85

Calls 10

IntArrayClass · 0.85
recompute_reduce_allFunction · 0.85
sortFunction · 0.50
GetMethod · 0.45
sizeMethod · 0.45
dimsMethod · 0.45
GetDataMethod · 0.45
dtypeMethod · 0.45
beginMethod · 0.45
endMethod · 0.45

Tested by

no test coverage detected