Method CheckForNodeRewrite

tensorflow/core/graph/mkl_layout_pass.cc:3802–3860

}

const MklLayoutRewritePass::RewriteInfo*
MklLayoutRewritePass::CheckForNodeRewrite(const Node* n) const {
  CHECK_NOTNULL(n);

  // QuantizedOps may have attributes other than "T", so decoupled the check
  // with a function, CheckForQuantizedNodeRewrite(const Node*).
  const RewriteInfo* ri = CheckForQuantizedNodeRewrite(n);
  if (ri != nullptr) return ri;

  // First check if node along with its type is supported by OneDNN layer.
  // We do not want to rewrite an op into OneDNN op if types are not supported.
  // E.g., OneDNNRelu does not support INT32. So we cannot rewrite Relu to
  // OneDNNRelu if type is INT32.
  DataType T;
  if (!TryGetNodeAttr(n->def(), "T", &T)) {
    return nullptr;
  }

  // We make an exception for Conv2D, as the corresponding OneDNN ops
  // currently do not support the case of padding == EXPLICIT yet.
  if (n->type_string() == csinfo_.conv2d ||
      n->type_string() == csinfo_.conv2d_grad_input ||
      n->type_string() == csinfo_.conv2d_grad_filter) {
    string padding;
    TF_CHECK_OK(GetNodeAttr(n->def(), "padding", &padding));
    if (padding == "EXPLICIT") return nullptr;
  }

  // We make an exception for __MklDummyConv2DWithBias,
  // __MklConv2DBackpropFilterWithBias, and __MklDummyPadWithConv2D since their
  // names do not match Mkl node names.
  if (n->type_string() != csinfo_.conv2d_with_bias &&
      n->type_string() != csinfo_.pad_with_conv2d &&
      n->type_string() != csinfo_.pad_with_fused_conv2d &&
      n->type_string() != csinfo_.conv2d_grad_filter_with_bias &&
      n->type_string() != csinfo_.fused_batch_norm_ex &&
      n->type_string() != csinfo_.fused_conv2d &&
      n->type_string() != csinfo_.fused_depthwise_conv2d &&
      n->type_string() != csinfo_.fused_matmul &&
      n->type_string() != csinfo_.fused_batch_matmul &&
      n->type_string() != csinfo_.fused_batch_matmul_v2 &&
      !mkl_op_registry::IsMklOp(mkl_op_registry::GetMklOpName(n->type_string()),
                                T)) {
    return nullptr;
  }

  // We now check if rewrite rule applies for this op. If rewrite rule passes
  // for this op, then we rewrite it to OneDNN op.
  // Find matching RewriteInfo and then check that rewrite rule applies.
  if (!CheckForRecoOpsListNodeRewrite(n)) return nullptr;
  for (auto ri = rinfo_.cbegin(); ri != rinfo_.cend(); ++ri) {
    if (n->type_string().compare(ri->name) == 0 && ri->rewrite_rule(n)) {
      return &*ri;
    }
  }

  // Else return not found.
  return nullptr;