Updates the solver to accept feasible -- albeit suboptimal -- solutions.

PiperOrigin-RevId: 680092346
Intel-tensorflow · Sep 29, 2024 · defd9fe · defd9fe
1 parent b9fe7ae
commit defd9fe
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc b/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc
@@ -875,11 +875,14 @@ AutoShardingSolverResult SolveAndExtractSolution(
   } else if (status == operations_research::MPSolver::MODEL_INVALID) {
     LOG(FATAL) << "Solver says that the input MIP is invalid. This is most "
                   "likely a bug and should be reported.";
-    return AutoShardingSolverResult(absl::InternalError("Solver timed out."),
+    return AutoShardingSolverResult(absl::InternalError("Invalid MIP."),
                                     /*skip_auto_sharding=*/false);
-  } else if (status != operations_research::MPSolver::OPTIMAL) {
+  } else if (status == operations_research::MPSolver::NOT_SOLVED) {
+    LOG(WARNING) << "Solver timeout; no solution was produced";
     return AutoShardingSolverResult(absl::InternalError("Solver timed out."),
                                     /*skip_auto_sharding=*/true);
+  } else if (status != operations_research::MPSolver::OPTIMAL) {
+    LOG(WARNING) << "Solver timeout; moving forward with a suboptimal solution";
   }
 
   // Fingerprint the model & solution (useful when checking for determinism).