| 2101 | } // namespace |
| 2102 | |
| 2103 | Status Remapper::Optimize(Cluster* cluster, const GrapplerItem& item, |
| 2104 | GraphDef* optimized_graph) { |
| 2105 | GrapplerItem mutable_item = item; |
| 2106 | Status status; |
| 2107 | RemapperContext ctx(&mutable_item, &status, xla_on_); |
| 2108 | TF_RETURN_IF_ERROR(status); |
| 2109 | // Processing graph in reverse-topological sorted order allows to remap |
| 2110 | // longer chains of dependent ops in one pass. |
| 2111 | TF_RETURN_IF_ERROR( |
| 2112 | ctx.graph_view.SortTopologically(/*ignore_cycles=*/false, {})); |
| 2113 | const int num_nodes = item.graph.node_size(); |
| 2114 | // Skip nodes that were invalidated by a remapper, e.g. do not process BiasAdd |
| 2115 | // and Activation nodes that were fused into a Conv2D node. |
| 2116 | std::vector<bool> invalidated_nodes(num_nodes); |
| 2117 | std::vector<bool> nodes_to_delete(num_nodes); |
| 2118 | |
| 2119 | // _Fused{...} kernels do not have registered gradient function, so we must |
| 2120 | // not perform rewrite if the graph will be differentiated later. |
| 2121 | bool allow_non_differentiable_rewrites = |
| 2122 | item.optimization_options().allow_non_differentiable_rewrites; |
| 2123 | |
| 2124 | for (int i = num_nodes - 1; i >= 0; --i) { |
| 2125 | // Check if node was invalidated by one of the previous remaps. |
| 2126 | if (invalidated_nodes[i] || nodes_to_delete[i]) { |
| 2127 | continue; |
| 2128 | } |
| 2129 | |
| 2130 | // Infer properties lazily in case they are not needed. |
| 2131 | if (!ctx.inferred_graph_properties && RequiresInferredShapes(ctx, i)) { |
| 2132 | const bool assume_valid_feeds = opt_level_ == RewriterConfig::AGGRESSIVE; |
| 2133 | TF_RETURN_IF_ERROR(ctx.graph_properties.InferStatically( |
| 2134 | assume_valid_feeds, |
| 2135 | /*aggressive_shape_inference=*/false, |
| 2136 | /*include_input_tensor_values=*/true, |
| 2137 | /*include_output_tensor_values=*/false)); |
| 2138 | ctx.inferred_graph_properties = true; |
| 2139 | } |
| 2140 | |
| 2141 | #ifdef INTEL_MKL |
| 2142 | if (!DisableMKL()) { |
| 2143 | ContractionWithBiasAddAndAdd contract_with_bias_and_add; |
| 2144 | ContractionWithBiasAndAddActivation contract_with_bias_and_add_activation; |
| 2145 | ContractionWithMul contract_with_mul; |
| 2146 | |
| 2147 | if (!item.optimization_options().is_eager_mode) { |
| 2148 | // Remap Conv2D+BiasAdd+Add+relu into the _FusedConv2D. |
| 2149 | if (MklLayoutPassLists::FindFusedMatMul() && |
| 2150 | FindContractionWithBiasAndAddActivation( |
| 2151 | ctx, i, &contract_with_bias_and_add_activation)) { |
| 2152 | TF_RETURN_IF_ERROR(AddFusedContractionNode( |
| 2153 | &ctx, contract_with_bias_and_add_activation, &invalidated_nodes, |
| 2154 | &nodes_to_delete)); |
| 2155 | continue; |
| 2156 | } |
| 2157 | |
| 2158 | // // Remap Conv2D+BiasAdd+Add into the _FusedConv2D. |
| 2159 | if (MklLayoutPassLists::FindFusedMatMul() && |
| 2160 | FindContractionWithBiasAddAndAdd(ctx, i, |