PR #22454: [Shardy] Create dump directory.

Imported from GitHub PR #22454 Copybara import of the project: -- ad713d8 by Yunlong Liu <[email protected]>: Update shardy_xla_pass.cc -- d5bb158 by Yunlong Liu <[email protected]>: Update BUILD Merging this change closes #22454 COPYBARA_INTEGRATE_REVIEW=#22454 from yliu120:create_dump_dir d5bb158 PiperOrigin-RevId: 724355015
openxla · Feb 7, 2025 · a2dcb60 · a2dcb60
1 parent 0c474b0
commit a2dcb60
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/xla/service/spmd/shardy/BUILD b/xla/service/spmd/shardy/BUILD
@@ -62,6 +62,7 @@ cc_library(
         "@llvm-project//mlir:Support",
         "@shardy//shardy/common:file_utils",
         "@shardy//shardy/dialect/sdy/transforms/propagation:passes",
+        "@tsl//tsl/platform:env",
         "@tsl//tsl/platform:errors",
         "@tsl//tsl/platform:path",
         "@tsl//tsl/platform:statusor",

diff --git a/xla/service/spmd/shardy/shardy_xla_pass.cc b/xla/service/spmd/shardy/shardy_xla_pass.cc
@@ -67,6 +67,7 @@ limitations under the License.
 #include "xla/tsl/framework/mlir/status_scoped_diagnostic_handler.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
+#include "tsl/platform/env.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/path.h"
 #include "tsl/platform/statusor.h"
@@ -329,7 +330,7 @@ absl::StatusOr<bool> ShardyXLA::Run(
         tsl::io::JoinPath(shardyDir, "shardy", uniqueModuleName(*hloModule));
     LOG(INFO) << "Using Shardy output directory: " << shardyDir;
   }
-
+  TF_RETURN_IF_ERROR(tsl::Env::Default()->RecursivelyCreateDir(shardyDir));
   // MLIR pipeline: (1) import, (2) Shardy, and (3) export.
 
   bool enableVerifier = false;