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
| 21 | // most outer dims, so in case of another type of reduction, it would be |
| 22 | // better to fallback to native implementation |
| 23 | inline 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 | |
| 57 | bool ReduceCheckIfOneDNNSupport(const KernelContext* dev_ctx) { |
| 58 | if (dev_ctx->InputAt<DenseTensor>(0).dims().size() > 5 || |
no test coverage detected