feat: mostly working

Author: avik-pal

CondaPkg.toml

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

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,

ext/ReactantPythonCallExt/ReactantPythonCallExt.jl

@@ -10,6 +10,7 @@ const JAX_TRACING_SUPPORTED = Ref{Bool}(false)
 
 const tfptr = Ref{Py}()
 const tf2xlaptr = Ref{Py}()
+const npptr = Ref{Py}()
 
 const SAVED_MODEL_EXPORT_SUPPORTED = Ref{Bool}(false)
 
@@ -26,9 +27,9 @@ const NUMPY_SIMPLE_TYPES = Dict(
     Float16 => :float16,
     Float32 => :float32,
     Float64 => :float64,
-    ComplexF16 => :complex32,
-    ComplexF32 => :complex64,
-    ComplexF64 => :complex128,
+    ComplexF16 => :complex16,
+    ComplexF32 => :complex32,
+    ComplexF64 => :complex64,
 )
 
 function __init__()
@@ -45,6 +46,7 @@ function __init__()
         tfptr[] = pyimport("tensorflow")
         tfptr[].config.set_visible_devices(pylist(); device_type="GPU")
         tf2xlaptr[] = pyimport("tensorflow.compiler.tf2xla.python.xla")
+        npptr[] = pyimport("numpy")
         SAVED_MODEL_EXPORT_SUPPORTED[] = true
     catch err
         @warn "Failed to import tensorflow. Exporting Reactant compiled functions as \

ext/ReactantPythonCallExt/saved_model.jl

@@ -0,0 +1,94 @@
+# TODO: at some point, we should use the TF C++ API to export the SavedModel
+
+function Reactant.Serialization.serialization_supported(::Val{:SavedModel})
+    return SAVED_MODEL_EXPORT_SUPPORTED[]
+end
+
+function _extract_call_parameters(args::Tuple, input_locations, state_dict)
+    call_args = []
+    for loc in input_locations
+        if loc isa Reactant.Serialization.TFSavedModel.InputArgument
+            push!(call_args, args[loc.position])
+        else
+            push!(call_args, state_dict[loc.name])
+        end
+    end
+    return call_args
+end
+
+function _wrap_as_tf_func(spec::Reactant.Serialization.TFSavedModel.ReactantFunctionSpec)
+    Touts = pylist([string(sig.dtype) for sig in spec.output_signature])
+    Souts = pylist([pylist(sig.shape) for sig in spec.output_signature])
+    return pyfunc(
+        function (args...)
+            return tf2xlaptr[].call_module(
+                pytuple(
+                    _extract_call_parameters(args, spec.input_locations, spec.state_dict)
+                );
+                version=5,
+                Tout=Touts,  # dtype information
+                Sout=Souts,  # Shape information
+                function_list=pylist([]),  # No functions to call
+                :module => spec.bytecode,
+            )
+        end,
+    )
+end
+
+function _make_input_signatures(
+    fn_spec::Reactant.Serialization.TFSavedModel.ReactantFunctionSpec
+)
+    input_pos_to_spec = Dict(
+        loc.position => spec for
+        (loc, spec) in zip(fn_spec.input_locations, fn_spec.input_signature) if
+        loc isa Reactant.Serialization.TFSavedModel.InputArgument
+    )
+
+    sigs = []
+    for i in 1:length(input_pos_to_spec)
+        spec = input_pos_to_spec[i]
+        dtype = getproperty(tfptr[], spec.dtype)
+        push!(
+            sigs,
+            tfptr[].TensorSpec(;
+                shape=pylist(spec.shape), dtype=dtype, name="args_$(i - 1)"
+            ),
+        )
+    end
+    return sigs
+end
+
+function Reactant.Serialization.TFSavedModel.__to_tf_saved_model(
+    fn_spec::Reactant.Serialization.TFSavedModel.ReactantFunctionSpec, path::String
+)
+    tfm = tfptr[].Module()
+
+    state_dict = Dict(
+        k => tfptr[].Variable(
+            npptr[].asarray(permutedims(v, collect(ndims(v):-1:1)));
+            # npptr[].asarray(v);
+            trainable=false,
+            name=k,
+        ) for (k, v) in fn_spec.state_dict
+    )
+
+    @show fn_spec.input_signature
+    @show fn_spec.output_signature
+
+    input_signatures = _make_input_signatures(fn_spec)
+
+    tfm.f = getproperty(tfptr[], :function)(
+        _wrap_as_tf_func(fn_spec); input_signature=pylist(input_signatures)
+    )
+    tfm._variables = pylist(collect(values(state_dict)))
+
+    signatures = Dict(
+        tfptr[].saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY =>
+            tfm.f.get_concrete_function(pylist(input_signatures)...),
+    )
+    save_options = tfptr[].saved_model.SaveOptions(; function_aliases=Dict("" => tfm.f))
+
+    tfptr[].saved_model.save(tfm, path; signatures=signatures, options=save_options)
+
+    return nothing
+end

src/mlir/IR/Operation.jl

@@ -368,3 +368,18 @@ function create_operation_at_front(args...; kwargs...)
     Base.pushfirst!(block(), res)
     return res
 end
+
+function FunctionType(op::Operation)
+    is_function_op = @ccall API.mlir_c.mlirIsFunctionOpInterface(
+        op::API.MlirOperation
+    )::Bool
+    if is_function_op
+        return Type(
+            @ccall API.mlir_c.mlirGetFunctionTypeFromOperation(
+                op::API.MlirOperation
+            )::API.MlirType
+        )
+    else
+        throw("operation is not a function operation")
+    end
+end

src/serialization/Serialization.jl

@@ -7,6 +7,8 @@ module Serialization
 
 using ..Reactant: Reactant, MLIR
 
+serialization_supported(::Val) = false
+
 include("TFSavedModel.jl")
 
 end

src/serialization/TFSavedModel.jl

@@ -1,9 +1,28 @@
 module TFSavedModel
 
+using ..Serialization: serialization_supported
 using ..Reactant: AbstractConcreteArray, AbstractConcreteNumber, Compiler, MLIR
 
 # https://github.com/openxla/stablehlo/blob/955fa7e6e3b0a6411edc8ff6fcce1e644440acbd/stablehlo/integrations/python/stablehlo/savedmodel/stablehlo_to_tf_saved_model.py
 
+const NUMPY_SIMPLE_TYPES = Dict(
+    Bool => :bool,
+    Int8 => :int8,
+    Int16 => :int16,
+    Int32 => :int32,
+    Int64 => :int64,
+    UInt8 => :uint8,
+    UInt16 => :uint16,
+    UInt32 => :uint32,
+    UInt64 => :uint64,
+    Float16 => :float16,
+    Float32 => :float32,
+    Float64 => :float64,
+    ComplexF16 => :complex16,
+    ComplexF32 => :complex32,
+    ComplexF64 => :complex64,
+)
+
 struct VariableSignature
     shape::Vector{Int}
     dtype::Symbol
@@ -19,21 +38,89 @@ struct Parameter <: VariableType
     name::String
 end
 
+struct ReactantFunctionSpec
+    input_signature::Vector{VariableSignature}
+    output_signature::Vector{VariableSignature}
+    input_locations::Vector{<:VariableType}
+    bytecode::Base.CodeUnits{UInt8,String}
+    state_dict::Dict
+end
+
 function export_as_saved_model(
     thunk::Compiler.Thunk,
     saved_model_path::String,
     target_version::VersionNumber,
-    input_locations,
-    state_dict::Dict{String, <:Union{<:AbstractConcreteArray,<:AbstractConcreteNumber}},
+    input_locations::Vector,
+    state_dict::Dict,
 )
     isempty(thunk.module_string) && error(
         "To export a thunk, ensure that it has been compiled with `serializable=true`."
     )
 
-    mlir_mod = parse(MLIR.IR.Module, thunk.module_string)
-    display(mlir_mod)
+    if !serialization_supported(Val(:SavedModel))
+        error("Serialization to SavedModel is not supported. This might happen if \
+               PythonCall hasn't been installed and loaded.")
+    end
+
+    mlir_mod = MLIR.IR.with_context() do ctx
+        parse(MLIR.IR.Module, thunk.module_string)
+    end
+
+    ftype = MLIR.IR.FunctionType(first(MLIR.IR.body(mlir_mod)))
+
+    input_signature = [
+        VariableSignature(
+            reverse(collect(Int64, size(MLIR.IR.input(ftype, i)))),
+            # collect(Int64, size(MLIR.IR.input(ftype, i))),
+            NUMPY_SIMPLE_TYPES[MLIR.IR.julia_type(eltype(MLIR.IR.input(ftype, i)))],
+        ) for i in 1:MLIR.IR.ninputs(ftype)
+    ]
+
+    output_signature = [
+        VariableSignature(
+            reverse(collect(Int64, size(MLIR.IR.result(ftype, i)))),
+            # collect(Int64, size(MLIR.IR.result(ftype, i))),
+            NUMPY_SIMPLE_TYPES[MLIR.IR.julia_type(eltype(MLIR.IR.result(ftype, i)))],
+        ) for i in 1:MLIR.IR.nresults(ftype)
+    ]
+
+    if isempty(input_locations)
+        input_locations = [InputArgument(i) for i in 1:length(input_signature)]
+    end
+
+    c_print_callback = @cfunction(
+        MLIR.IR.print_callback, Cvoid, (MLIR.API.MlirStringRef, Any)
+    )
+    ref = Ref(IOBuffer())
+    result = MLIR.IR.LogicalResult(
+        MLIR.API.stablehloSerializePortableArtifactFromModule(
+            mlir_mod, string(target_version), c_print_callback, ref, true
+        ),
+    )
+    MLIR.IR.isfailure(result) && throw("Couldn't serialize the module")
+    serialized_module = codeunits(String(take!(ref[])))
+
+    return to_tf_saved_model(
+        ReactantFunctionSpec(
+            input_signature,
+            output_signature,
+            input_locations,
+            serialized_module,
+            Dict(k => Array(v) for (k, v) in state_dict),
+        ),
+        saved_model_path,
+    )
+end
 
-    return nothing
+function to_tf_saved_model(fn_spec::ReactantFunctionSpec, path::String)
+    if !serialization_supported(Val(:SavedModel))
+        error("Serialization to SavedModel is not supported. This might happen if \
+               PythonCall hasn't been installed and loaded.")
+    end
+    return __to_tf_saved_model(fn_spec, path)
 end
 
+# Defined in the PythonCallExt module
+function __to_tf_saved_model end
+
 end