chore: cpp side changes for tf SavedModel export (#1427)

Author: avik-pal

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

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",
     ],

src/mlir/libMLIR_h.jl

@@ -10507,6 +10507,93 @@ function stablehloResultAccuracyAttrGetMode(attr)
     @ccall mlir_c.stablehloResultAccuracyAttrGetMode(attr::MlirAttribute)::MlirAttribute
 end
 
+function mlirGetDialectHandle__stablehlo__()
+    @ccall mlir_c.mlirGetDialectHandle__stablehlo__()::MlirDialectHandle
+end
+
+function stablehloGetApiVersion()
+    @ccall mlir_c.stablehloGetApiVersion()::Cint
+end
+
+@cenum MlirStablehloCompatibilityRequirement::UInt32 begin
+    NONE = 0x0000000000000000
+    WEEK_4 = 0x0000000000000001
+    WEEK_12 = 0x0000000000000002
+    MAX = 0x0000000000000003
+end
+
+function stablehloVersionFromCompatibilityRequirement(requirement, callback, userData)
+    @ccall mlir_c.stablehloVersionFromCompatibilityRequirement(
+        requirement::MlirStablehloCompatibilityRequirement,
+        callback::MlirStringCallback,
+        userData::Ptr{Cvoid},
+    )::Cvoid
+end
+
+function stablehloGetCurrentVersion(callback, userData)
+    @ccall mlir_c.stablehloGetCurrentVersion(
+        callback::MlirStringCallback, userData::Ptr{Cvoid}
+    )::Cvoid
+end
+
+function stablehloGetMinimumVersion(callback, userData)
+    @ccall mlir_c.stablehloGetMinimumVersion(
+        callback::MlirStringCallback, userData::Ptr{Cvoid}
+    )::Cvoid
+end
+
+function stablehloGetSmallerVersion(version1, version2, callback, userData)
+    @ccall mlir_c.stablehloGetSmallerVersion(
+        version1::MlirStringRef,
+        version2::MlirStringRef,
+        callback::MlirStringCallback,
+        userData::Ptr{Cvoid},
+    )::MlirLogicalResult
+end
+
+function stablehloSerializePortableArtifactFromStringRef(
+    moduleStr, targetVersion, callback, userData
+)
+    @ccall mlir_c.stablehloSerializePortableArtifactFromStringRef(
+        moduleStr::MlirStringRef,
+        targetVersion::MlirStringRef,
+        callback::MlirStringCallback,
+        userData::Ptr{Cvoid},
+    )::MlirLogicalResult
+end
+
+function stablehloSerializePortableArtifactFromModule(
+    moduleStr, targetVersion, callback, userData, allowOtherDialects
+)
+    @ccall mlir_c.stablehloSerializePortableArtifactFromModule(
+        moduleStr::MlirModule,
+        targetVersion::MlirStringRef,
+        callback::MlirStringCallback,
+        userData::Ptr{Cvoid},
+        allowOtherDialects::Bool,
+    )::MlirLogicalResult
+end
+
+function stablehloDeserializePortableArtifact(artifactStr, callback, userData)
+    @ccall mlir_c.stablehloDeserializePortableArtifact(
+        artifactStr::MlirStringRef, callback::MlirStringCallback, userData::Ptr{Cvoid}
+    )::MlirLogicalResult
+end
+
+function stablehloDeserializePortableArtifactNoError(artifactStr, ctx)
+    @ccall mlir_c.stablehloDeserializePortableArtifactNoError(
+        artifactStr::MlirStringRef, ctx::MlirContext
+    )::MlirModule
+end
+
+function stablehloTokenTypeGet(ctx)
+    @ccall mlir_c.stablehloTokenTypeGet(ctx::MlirContext)::MlirType
+end
+
+function stablehloTypeIsAToken(type)
+    @ccall mlir_c.stablehloTypeIsAToken(type::MlirType)::Bool
+end
+
 function sdyAttributeIsAMeshAxisAttr(attr)
     @ccall mlir_c.sdyAttributeIsAMeshAxisAttr(attr::MlirAttribute)::Bool
 end