Check whether opname with type T is registered as OneDNN operator that can accept input tensors in OneDNN layout. @input: name of the op @input: T datatype to be used for checking op @return: true if opname is registered as OneDNN-layout dependent op; false otherwise
| 155 | // @return: true if opname is registered as OneDNN-layout dependent op; |
| 156 | // false otherwise |
| 157 | static inline bool IsMklLayoutDependentOp(const string& op_name, DataType T) { |
| 158 | string kernel = KernelsRegisteredForOp(op_name); |
| 159 | |
| 160 | // Restrict quantized ops to QUINT8 and QINT8 for now |
| 161 | if (kernel.find(kMklQuantizedOpLabelPattern) != string::npos) { |
| 162 | return (T == DT_QUINT8 || T == DT_QINT8 || T == DT_QINT32); |
| 163 | } |
| 164 | // Restrict regular ops to FLOAT and BFLOAT16 |
| 165 | if (kernel.find(kMklLayoutDependentOpLabelPattern) != string::npos) { |
| 166 | if (T == DT_FLOAT) return true; |
| 167 | if (T == DT_BFLOAT16) { |
| 168 | if (IsBF16SupportedByOneDNNOnThisCPU()) { |
| 169 | return true; |
| 170 | } else { |
| 171 | // Restrict bfloat16 ops to platforms with at least AVX512 support, fall |
| 172 | // back to Eigen implementation otherwise. |
| 173 | BF16UnsupportedWarning(); |
| 174 | return false; |
| 175 | } |
| 176 | } |
| 177 | return false; |
| 178 | } |
| 179 | return false; |
| 180 | } |
| 181 | |
| 182 | // TODO(mdfaijul): QuantizedConv2D is registered with input: QUINT8 |
| 183 | // filter:QINT8 for dnnl integration. First a dummy kernel is created |
no test coverage detected