feat: jll changes for mhlo/xla passes (#1464)

* feat: jll changes for mhlo/xla passes

* ci: free disk space

* Update .github/workflows/CI-localjll.yml

Author: avik-pal

.github/workflows/CI-localjll.yml

@@ -39,6 +39,13 @@ jobs:
     env:
       TMPDIR: ${{ github.workspace }}/tmp
     steps:
+      - name: Free Disk Space
+        uses: jlumbroso/free-disk-space@main
+        with:
+          tool-cache: false
+        if: ${{ startsWith(matrix.os, 'ubuntu-') }}
+      - name: Clean `/opt`
+        run: sudo rm -rf /opt/*
       - uses: actions/checkout@v4
       - name: Create TMPDIR
         run: |

deps/ReactantExtra/API.cpp

@@ -81,6 +81,7 @@
 #include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
 
 #include "xla/hlo/translate/hlo_to_mhlo/hlo_utils.h"
+#include "xla/hlo/translate/stablehlo.h"
 
 // CPU collectives
 #include "xla/backends/cpu/collectives/mpi_collectives.h"
@@ -2941,3 +2942,10 @@ extern "C" void reactantXLAExec(LinkableRuntime **__restrict__ lrtP,
     }
   }
 }
+
+extern "C" HeldHloModule *convertMlirModuleToHloModule(MlirModule mod) {
+  mlir::ModuleOp cmod_op = cast<ModuleOp>(*unwrap(mod));
+  std::shared_ptr<xla::HloModule> hlo_module =
+      std::move(MyValueOrThrow(xla::ConvertStablehloToHlo(cmod_op)));
+  return reactant::capture(hlo_module);
+}

deps/ReactantExtra/BUILD

@@ -890,6 +890,7 @@ cc_library(
             "-Wl,-exported_symbol,_addSdyPropagationPipeline",
             "-Wl,-exported_symbol,_mlirGetFunctionTypeFromOperation",
             "-Wl,-exported_symbol,_mlirIsFunctionOpInterface",
+            "-Wl,-exported_symbol,_convertMlirModuleToHloModule",
         ],
     }),
     linkstatic = True,
@@ -960,6 +961,7 @@ cc_library(
         "@xla//xla/pjrt:status_casters",
         "@xla//xla/python/ifrt",
         "@xla//xla/python/pjrt_ifrt",
+        "@xla//xla/hlo/translate:stablehlo",
         "@xla//xla/python/ifrt_proxy/server:grpc_server",
         "@xla//xla/python/ifrt_proxy/client:grpc_client",
         "@xla//xla/python/ifrt_proxy/client:registry",

deps/ReactantExtra/WORKSPACE

@@ -11,7 +11,7 @@ http_archive(
     urls = ["https://github.com/wsmoses/nsync/archive/{commit}.tar.gz".format(commit = NSYNC_COMMIT)],
 )
 
-ENZYMEXLA_COMMIT = "85adc0201b0e605ce1bc89cf235ac0bfc23e77ab"
+ENZYMEXLA_COMMIT = "6774f1afb90c377bbf234a7a7dbfab4f7b726481"
 
 ENZYMEXLA_SHA256 = ""