Function IsGpuCompatible

tensorflow/core/grappler/optimizers/remapper.cc:349–385

Checks if we can rewrite a pattern to the `_FusedConv2D` on GPU device.

// Checks if we can rewrite a pattern to the `_FusedConv2D` on GPU device.
bool IsGpuCompatible(const RemapperContext& ctx,
                     const ContractionWithBiasAddAndActivation& matched) {
#if TENSORFLOW_USE_ROCM
  // ROCm does not support _FusedConv2D
  return false;
#endif
  // xla does not cluster _FusedConv2D, and knows how to do this optimization
  if (ctx.xla_on_) return false;

  const GraphDef* graph = ctx.graph_view.graph();
  
  // We rely on cuDNN for fused convolution and cublasLt for fused matmul.
  const NodeDef& activation_node = graph->node(matched.activation);
  if (!IsRelu(activation_node)) return false;
  const NodeDef& contraction_node = graph->node(matched.contraction);
  if (IsConv2D(contraction_node)) {
    const std::vector<OpInfo::TensorProperties>& input_props =
        ctx.graph_properties.GetInputProperties(contraction_node.name());
    const TensorShapeProto& filter_shape =
        input_props.size() >= 2 ? input_props[1].shape() : TensorShapeProto();

    // FusedConv2D on GPU with 1x1 convolution is marginally faster than
    // in-graph computation in micro benchmarks (see kernels/conv_ops_test.cc),
    // and significantly slower in large scale benchmarks.
    bool is_spatial_conv = Rank(filter_shape) == 4 &&          //
                           IsKnown(filter_shape.dim(1)) &&     //
                           IsKnown(filter_shape.dim(2)) &&     //
                           filter_shape.dim(1).size() != 1 &&  //
                           filter_shape.dim(2).size() != 1;

    return is_spatial_conv && IsGpuCompatibleConv2D(ctx, &contraction_node);
  } else if (IsMatMul(contraction_node)) {
    return IsGpuCompatibleMatMul(ctx, &contraction_node);
  }

  return false;
}

// Checks if we can rewrite a pattern to the `_FusedMatMul` on GPU device.
bool IsGpuCompatible(const RemapperContext& ctx,