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

Author: sbrantq

.github/workflows/format-pr-bazel.yml

@@ -24,9 +24,9 @@ jobs:
       uses: peter-evans/create-pull-request@v7
       with:
         token: ${{ secrets.GITHUB_TOKEN }}
-        commit-message: Format code
-        title: 'Format code of branch "main"'
-        branch: format-main
+        commit-message: Format Bazel code
+        title: 'Format Bazel code of branch "main"'
+        branch: format-main-bazel
         delete-branch: true
         labels: format
         author: enzyme-ci-bot[bot] <78882869+enzyme-ci-bot[bot]@users.noreply.github.com>

.github/workflows/format-pr.yml

@@ -1,4 +1,4 @@
-name: Format 'main'
+name: Format Julia code on 'main'
 
 on:
   schedule:
@@ -34,9 +34,9 @@ jobs:
         uses: peter-evans/create-pull-request@v7
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
-          commit-message: Format code
-          title: 'Format code of branch "main"'
-          branch: format-main
+          commit-message: Format Julia code
+          title: 'Format Julia code of branch "main"'
+          branch: format-main-julia
           delete-branch: true
           labels: format
           author: enzyme-ci-bot[bot] <78882869+enzyme-ci-bot[bot]@users.noreply.github.com>

.github/workflows/regenerate-mlir-bindings.yml

@@ -16,9 +16,17 @@ jobs:
         with:
           version: '1.10'
       - uses: julia-actions/cache@v2
-      - uses: actions/checkout@v4
+      - uses: bazel-contrib/[email protected]
+        name: Set up Bazel
         with:
-          ref: main
+          # Avoid downloading Bazel every time.
+          bazelisk-cache: true
+          # Store build cache per workflow.
+          disk-cache: ${{ github.workflow }}-${{ matrix.os }}-${{ matrix.version }}
+          # Share repository cache between workflows.
+          repository-cache: true
+          bazelisk-version: 1.x
+      - uses: actions/checkout@v4
       - name: Install JuliaFormatter.jl
         shell: julia --color=yes {0}
         run: |

CondaPkg.toml

@@ -1,2 +1,4 @@
 [pip.deps]
-jax = ">=0.4"
+jax = ">= 0.6"
+tensorflow = ">= 2.17"
+numpy = ">= 2"

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.139"
+version = "0.2.143"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -86,11 +86,11 @@ OneHotArrays = "0.2.10"
 OrderedCollections = "1"
 PrecompileTools = "1.2"
 Preferences = "1.4"
-PythonCall = "0.9"
+PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.15"
-Reactant_jll = "0.0.213"
+Reactant_jll = "0.0.214"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

deps/ReactantExtra/API.cpp

@@ -270,6 +270,17 @@ extern "C" MlirOperation mlirOperationParse(MlirContext ctx, MlirBlock block,
           .release()};
 }
 
+extern "C" MlirType mlirGetFunctionTypeFromOperation(MlirOperation op) {
+  if (auto funcOp = dyn_cast<mlir::FunctionOpInterface>(unwrap(op))) {
+    return wrap(funcOp.getFunctionType());
+  }
+  ReactantThrowError("Not a function op");
+}
+
+extern "C" bool mlirIsFunctionOpInterface(MlirOperation op) {
+  return llvm::isa<mlir::FunctionOpInterface>(unwrap(op));
+}
+
 // TODO mlirComplexAttrGetnValue
 // TODO extern "C" MlirTypeID mlirComplexAttrGetTypeID(void) { return
 // wrap(complex::NumberAttr::getTypeID()); }

deps/ReactantExtra/BUILD

@@ -888,6 +888,8 @@ cc_library(
             "-Wl,-exported_symbol,_hlo_sharding_*",
             "-Wl,-exported_symbol,_free_ifrt_sharding",
             "-Wl,-exported_symbol,_addSdyPropagationPipeline",
+            "-Wl,-exported_symbol,_mlirGetFunctionTypeFromOperation",
+            "-Wl,-exported_symbol,_mlirIsFunctionOpInterface",
         ],
     }),
     linkstatic = True,
@@ -1434,6 +1436,9 @@ genrule(
         "@llvm-project//mlir:AsyncPassIncGen_filegroup",
         "@llvm-project//mlir:GPUPassIncGen_filegroup",
         "@stablehlo//:stablehlo/integrations/c/StablehloAttributes.h",
+        "@stablehlo//:stablehlo/integrations/c/StablehloDialect.h",
+        "@stablehlo//:stablehlo/integrations/c/StablehloDialectApi.h",
+        "@stablehlo//:stablehlo/integrations/c/StablehloTypes.h",
         "@shardy//shardy/integrations/c:attributes.h",
         "//:Project.toml",
         "//:Manifest.toml",
@@ -1442,7 +1447,7 @@ genrule(
         "//:make.jl",
     ],
     outs = ["libMLIR_h.jl"],
-    cmd = "$$JULIA \"--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)\" \"$@\"",
     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 = "e52dc6b096213e21940693cda860c99ada2de247"
+ENZYMEXLA_COMMIT = "2527ca4bb8fa4499cd10ffb42ce4c2cda3738e91"
 
 ENZYMEXLA_SHA256 = ""
 
@@ -52,7 +52,24 @@ hedron_compile_commands_setup_transitive_transitive_transitive()
 
 load("@enzyme_ad//:workspace.bzl", "ENZYME_COMMIT", "ENZYME_SHA256", "JAX_COMMIT", "JAX_SHA256", "XLA_PATCHES")
 
-XLA_PATCHES = XLA_PATCHES + [
+CUPTI_OLD = [
+    """
+sed -i.bak0 "s/cupti_driver_cbid/cupti/g" xla/backends/profiler/gpu/cupti_tracer.cc
+""",
+    """
+sed -i.bak0 "/CUPTI_DRIVER_TRACE_CBID_cuGraphAddNode/d" xla/backends/profiler/gpu/cupti_tracer.cc
+""",
+    """
+sed -i.bak0 "/CUPTI_DRIVER_TRACE_CBID_cuGraphAddNode_v2/d" xla/backends/profiler/gpu/cupti_tracer.cc
+""",
+]
+
+CUPTI_NEW = []
+
+XLA_PATCHES = XLA_PATCHES + CUPTI_NEW + [
+    """
+sed -i.bak0 "s/cupti_driver_cbid/cupti/g" xla/backends/profiler/gpu/cupti_tracer.cc
+""",
     """
 sed -i.bak0 "s/patch_cmds = \\[/patch_cmds = \\[\\\"find . -type f -name config.bzl -exec sed -i.bak0 's\\/HAVE_LINK_H=1\\/HAVE_LINK_H=0\\/g' {} +\\\",/g" third_party/llvm/workspace.bzl
 """,

docs/src/.vitepress/config.mts

@@ -95,6 +95,7 @@ export default defineConfig({
         items: [
           { text: "Core Reactant API", link: "/api/api" },
           { text: "Sharding", link: "/api/sharding" },
+          { text: "Serialization", link: "/api/serialization" },
           { text: "Ops", link: "/api/ops" },
           { text: "Configuration", link: "/api/config" },
           {
@@ -169,6 +170,7 @@ export default defineConfig({
             link: "/api/api",
           },
           { text: "Sharding", link: "/api/sharding" },
+          { text: "Serialization", link: "/api/serialization" },
           { text: "Ops", link: "/api/ops" },
           { text: "Configuration", link: "/api/config" },
           {

docs/src/api/serialization.md

@@ -0,0 +1,29 @@
+```@meta
+CollapsedDocStrings = true
+```
+
+# Serialization
+
+```@docs
+Reactant.Serialization
+```
+
+## Exporting to TensorFlow SavedModel
+
+!!! note "Load PythonCall"
+
+    Serialization to TensorFlow SavedModel requires PythonCall to be loaded. Loading
+    PythonCall will automatically install tensorflow. If tensorflow installation fails,
+    we won't be able to export to SavedModel.
+
+A SavedModel contains a complete TensorFlow program, including trained parameters (i.e,
+tf.Variables) and computation. It does not require the original model building code to run,
+which makes it useful for sharing or deploying with [TFLite](https://tensorflow.org/lite),
+[TensorFlow.js](https://js.tensorflow.org/),
+[TensorFlow Serving](https://www.tensorflow.org/tfx/serving/tutorials/Serving_REST_simple),
+or [TensorFlow Hub](https://tensorflow.org/hub). Refer to the
+[official documentation](https://www.tensorflow.org/guide/saved_model) for more details.
+
+```@docs
+Reactant.Serialization.export_as_tf_saved_model
+```