Merge branch 'main' of https://github.com/EnzymeAD/Reactant.jl into probprog-trace-operand

Author: sbrantq

.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: |

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.145"
+version = "0.2.146"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -90,7 +90,7 @@ PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.15"
-Reactant_jll = "0.0.217"
+Reactant_jll = "0.0.219"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

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"
@@ -2956,3 +2957,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",
@@ -1310,6 +1312,24 @@ gentbl_cc_library(
     ],
 )
 
+gentbl_cc_library(
+    name = "MosaicGPUJLIncGen",
+    tbl_outs = [
+        (
+            [
+                "--generator=jl-op-defs",
+                "--disable-module-wrap=0",
+            ],
+            "MosaicGPU.jl",
+        ),
+    ],
+    tblgen = "//:mlir-jl-tblgen",
+    td_file = "@jax//jaxlib/mosaic/dialect/gpu:mosaic_gpu.td",
+    deps = [
+        "@jax//jaxlib/mosaic/dialect/gpu:mosaic_gpu_td_files",
+    ],
+)
+
 gentbl_cc_library(
     name = "TritonJLIncGen",
     tbl_outs = [
@@ -1440,14 +1460,19 @@ genrule(
         "@stablehlo//:stablehlo/integrations/c/StablehloDialectApi.h",
         "@stablehlo//:stablehlo/integrations/c/StablehloTypes.h",
         "@shardy//shardy/integrations/c:attributes.h",
+        "@jax//jaxlib/triton:triton_dialect_capi.h",
+        "@jax//jaxlib/mosaic:dialect/tpu/integrations/c/tpu_dialect.h",
+        "@jax//jaxlib/mosaic:dialect/tpu/integrations/c/tpu_passes.capi.h.inc",
+        "@jax//jaxlib/mosaic/dialect/gpu:integrations/c/attributes.h",
+        "@jax//jaxlib/mosaic/dialect/gpu:integrations/c/gpu_dialect.h",
         "//:Project.toml",
         "//:Manifest.toml",
         "//:wrap.toml",
         "//:missing_defs.jl",
         "//:make.jl",
     ],
     outs = ["libMLIR_h.jl"],
-    cmd = "$$JULIA \"--color=yes\" \"--project=$(location //:Project.toml)\" \"$(location //:make.jl)\" \"$(location @llvm-project//mlir:include/mlir-c/Bindings/Python/Interop.h)\" \"$(location @llvm-project//llvm:include/llvm-c/Support.h)\" \"$(locations @llvm-project//mlir:ConversionPassIncGen_filegroup)\" \"$(location @stablehlo//:stablehlo/integrations/c/StablehloAttributes.h)\" \"$(location @shardy//shardy/integrations/c:attributes.h)\" \"$@\"",
+    cmd = "$$JULIA \"--color=yes\" \"--project=$(location //:Project.toml)\" \"$(location //:make.jl)\" \"$(location @llvm-project//mlir:include/mlir-c/Bindings/Python/Interop.h)\" \"$(location @llvm-project//llvm:include/llvm-c/Support.h)\" \"$(locations @llvm-project//mlir:ConversionPassIncGen_filegroup)\" \"$(location @stablehlo//:stablehlo/integrations/c/StablehloAttributes.h)\" \"$(location @shardy//shardy/integrations/c:attributes.h)\" \"$(location @jax//jaxlib/triton:triton_dialect_capi.h)\" \"$(location @jax//jaxlib/mosaic:dialect/tpu/integrations/c/tpu_dialect.h)\" \"$(location @jax//jaxlib/mosaic/dialect/gpu:integrations/c/gpu_dialect.h)\" \"$@\"",
     tags = [
         "jlrule",
     ],

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 = "20aff0cfd430339ca01c9febb96675d62a4a7995"
+ENZYMEXLA_COMMIT = "6774f1afb90c377bbf234a7a7dbfab4f7b726481"
 
 ENZYMEXLA_SHA256 = ""
 

deps/ReactantExtra/make-bindings.jl

@@ -36,6 +36,7 @@ for file in [
     "Gpu.jl",
     "Affine.jl",
     "TPU.jl",
+    "MosaicGPU.jl",
     "Triton.jl",
     "Shardy.jl",
     "MPI.jl",

deps/ReactantExtra/make.jl

@@ -19,10 +19,11 @@ let options = deepcopy(options)
     genarg = first(eachsplit(ARGS[3], " "))
 
     gen_include_dir = joinpath(splitpath(genarg)[1:(end - 4)]...)
-
-    hlo_include_dir = joinpath(splitpath(ARGS[end - 2])[1:(end - 1)]...)
-
-    sdy_include_dir = joinpath(splitpath(ARGS[end - 1])[1:(end - 1)]...)
+    hlo_include_dir = joinpath(splitpath(ARGS[end - 5])[1:(end - 1)]...)
+    sdy_include_dir = joinpath(splitpath(ARGS[end - 4])[1:(end - 1)]...)
+    triton_include_dir = joinpath(splitpath(ARGS[end - 3])[1:(end - 1)]...)
+    mosaic_tpu_include_dir = joinpath(splitpath(ARGS[end - 2])[1:(end - 1)]...)
+    mosaic_gpu_include_dir = joinpath(splitpath(ARGS[end - 1])[1:(end - 1)]...)
 
     append!(
         args,
@@ -37,6 +38,12 @@ let options = deepcopy(options)
             hlo_include_dir,
             "-I",
             sdy_include_dir,
+            "-I",
+            triton_include_dir,
+            "-I",
+            mosaic_tpu_include_dir,
+            "-I",
+            mosaic_gpu_include_dir,
             "-x",
             "c++",
         ],
@@ -46,6 +53,9 @@ let options = deepcopy(options)
         detect_headers(include_dir, args, Dict(), endswith("Python/Interop.h"))...,
         detect_headers(hlo_include_dir, args, Dict())...,
         detect_headers(sdy_include_dir, args, Dict())...,
+        detect_headers(triton_include_dir, args, Dict())...,
+        detect_headers(mosaic_tpu_include_dir, args, Dict())...,
+        detect_headers(mosaic_gpu_include_dir, args, Dict())...,
     ]
 
     ctx = create_context(headers, args, options)

docs/src/.vitepress/config.mts

@@ -112,6 +112,7 @@ export default defineConfig({
               { text: "LLVM", link: "/api/dialects/llvm" },
               { text: "MPI", link: "/api/dialects/mpi" },
               { text: "MemRef", link: "/api/dialects/memref" },
+              { text: "Mosaic GPU", link: "/api/dialects/mosaicgpu" },
               { text: "NVVM", link: "/api/dialects/nvvm" },
               { text: "Shardy", link: "/api/dialects/shardy" },
               { text: "SparseTensor", link: "/api/dialects/sparsetensor" },
@@ -188,6 +189,7 @@ export default defineConfig({
               { text: "LLVM", link: "/api/dialects/llvm" },
               { text: "MPI", link: "/api/dialects/mpi" },
               { text: "MemRef", link: "/api/dialects/memref" },
+              { text: "Mosaic GPU", link: "/api/dialects/mosaicgpu" },
               { text: "NVVM", link: "/api/dialects/nvvm" },
               { text: "Shardy", link: "/api/dialects/shardy" },
               { text: "SparseTensor", link: "/api/dialects/sparsetensor" },

docs/src/api/dialects/mosaicgpu.md

@@ -0,0 +1,9 @@
+```@meta
+CollapsedDocStrings = true
+```
+
+# Mosaic GPU Dialect
+
+```@autodocs
+Modules = [Reactant.MLIR.Dialects.mosaic_gpu]
+```