Function IsMklNameChangeOp

tensorflow/core/graph/mkl_graph_util.h:203–243

Check whether opname with type T is registered as an OneDNN operator that will go through name change. @input: name of the op @input: T datatype to be used for checking op @return: true if opname is registered as OneDNN op that will go through name change; false otherwise

// @return: true if opname is registered as OneDNN op that will go through name
// change; false otherwise
static inline bool IsMklNameChangeOp(const string& op_name, DataType T) {
  string kernel = KernelsRegisteredForOp(op_name);
  // String returned by KernelsRegisteredForOp looks like below:
  //
  // Op = _MklMatMul, kernels =
  // device='CPU'; label='MklNameChangeOp'; T in [DT_COMPLEX128]
  // device='CPU'; label='MklNameChangeOp'; T in [DT_COMPLEX64]
  // device='CPU'; label='MklNameChangeOp'; T in [DT_DOUBLE]
  // device='CPU'; label='MklNameChangeOp'; T in [DT_FLOAT]

  // Now we just construct a search string to match what we are looking for.
  string search_string = kMklNameChangeOpLabelPattern;
  search_string += string(";") + string(" T in [");
  search_string += DataType_Name(T) + string("]");

  // Temporarily replacing earlier check by adding a type-specific check so
  // that we can selectively decide which type is supported by OneDNN operators.
  // That way kernel registration does not decide which operators we support.
  // We are using this change to temporarily disable BFLOAT16 support. Once
  // we want to enable it, we will go back to earlier check.
  bool isTypeAllowed = false;
  if (kernel.find(search_string) != string::npos) {
    isTypeAllowed = (T == DT_COMPLEX128 || T == DT_COMPLEX64 ||
                     T == DT_DOUBLE || T == DT_FLOAT);
    if (!isTypeAllowed) {
      if (T == DT_BFLOAT16) {
        if (IsBF16SupportedByOneDNNOnThisCPU()) {
          isTypeAllowed = true;
        } else {
          // Restrict bfloat16 ops to platforms with at least AVX512 support,
          // fall back to Eigen implementation otherwise.
          BF16UnsupportedWarning();
          isTypeAllowed = false;
        }
      }
    }
    return isTypeAllowed;
  }

  return false;
}

// Check if the operator with 'op_name' and type 'T' is an OneDNN operator that
// will either understand input tensors in OneDNN layout or will go through name