| 1249 | } |
| 1250 | |
| 1251 | Status AutoMixedPrecisionImpl::Optimize() { |
| 1252 | string optimization_level; |
| 1253 | TF_RETURN_IF_ERROR(ReadStringFromEnvVar( |
| 1254 | "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_LEVEL", "", &optimization_level)); |
| 1255 | optimization_level = absl::AsciiStrToUpper(optimization_level); |
| 1256 | force_all_fp16_ = optimization_level == "UNSAFE_FORCE_ALL"; |
| 1257 | if (force_all_fp16_ && mode_ == AutoMixedPrecisionMode::MKL) { |
| 1258 | // Many ops do not support bfloat16 on the CPU so we disallowing forcing to |
| 1259 | // bfloat16. |
| 1260 | return errors::InvalidArgument( |
| 1261 | "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_LEVEL cannot be set to " |
| 1262 | "UNSAFE_FORCE_ALL when MKL is used"); |
| 1263 | } |
| 1264 | |
| 1265 | std::unique_ptr<AutoMixedPrecisionLists> mp_lists = |
| 1266 | get_mixed_precision_lists(); |
| 1267 | f16_allowlist_ = mp_lists->AllowList(); |
| 1268 | f16_denylist_ = mp_lists->DenyList(); |
| 1269 | f16_inferlist_ = mp_lists->InferList(); |
| 1270 | f16_clearlist_ = mp_lists->ClearList(); |
| 1271 | TF_RETURN_IF_ERROR(ValidateLists(f16_allowlist_, f16_denylist_, |
| 1272 | f16_inferlist_, f16_clearlist_)); |
| 1273 | |
| 1274 | size_t timestamp = Env::Default()->NowMicros() / 1000; |
| 1275 | TF_RETURN_IF_ERROR(PrintDebugLogs(/* preop = */ true, timestamp)); |
| 1276 | |
| 1277 | VLOG(2) << "Identifying nodes that should be processed"; |
| 1278 | for (const NodeDef& node : graph_->node()) { |
| 1279 | bool should_process; |
| 1280 | switch (mode_) { |
| 1281 | case AutoMixedPrecisionMode::CUDA: |
| 1282 | should_process = |
| 1283 | !MustPreserve(node) && IsOnDevice(node, DEVICE_GPU) && |
| 1284 | (ShouldIgnorePerformance() || IsOnSuitableGPUArch(node)); |
| 1285 | break; |
| 1286 | case AutoMixedPrecisionMode::MKL: |
| 1287 | should_process = !MustPreserve(node) && IsOnDevice(node, DEVICE_CPU); |
| 1288 | break; |
| 1289 | } |
| 1290 | if (should_process) { |
| 1291 | should_process_nodes_.insert(&node); |
| 1292 | } else { |
| 1293 | LogSkippedNode(node); |
| 1294 | } |
| 1295 | } |
| 1296 | |
| 1297 | VLOG(2) << "Converting FusedBatchNorm* ops to V2"; |
| 1298 | ConvertBatchNormOpsToV2(); |
| 1299 | |
| 1300 | VLOG(2) << "Building node type map for graph"; |
| 1301 | TF_RETURN_IF_ERROR(node_type_map_.Init(*graph_)); |
| 1302 | |
| 1303 | VLOG(2) << "Constructing graph type attribute topology view"; |
| 1304 | TF_RETURN_IF_ERROR( |
| 1305 | graph_type_view_.InitializeFromGraph(*graph_, node_type_map_)); |
| 1306 | |
| 1307 | absl::flat_hash_set<int> deny_set; |
| 1308 |
nothing calls this directly
no test coverage detected