feat: serialization

Author: avik-pal

Project.toml

@@ -13,6 +13,7 @@ EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LLVMOpenMP_jll = "1d63c593-3942-5779-bab2-d838dc0a180e"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -73,6 +74,7 @@ Functors = "0.5"
 GPUArraysCore = "0.2"
 GPUCompiler = "1.3"
 HTTP = "1.10.15"
+JLD2 = "0.5.12"
 KernelAbstractions = "0.9.30"
 LLVM = "9.1"
 LLVMOpenMP_jll = "18.1.7"

src/Compiler.jl

@@ -1300,6 +1300,7 @@ macro compile(args...)
         :raise => false,
         :shardy_passes => :(:to_mhlo_shardings),
         :assert_nonallocating => false,
+        :serializable => true,
     )
     return esc(first(compile_call_expr(__module__, compile, default_options, args...)))
 end
@@ -1963,7 +1964,7 @@ function __resolve_device_and_client(client, seen_args, linear_args, is_sharded)
     return (client, device)
 end
 
-function compile_xla(f, args; client=nothing, kwargs...)
+function compile_xla(f, args; client=nothing, serializable::Bool=false, kwargs...)
     # register MLIR dialects
     ctx = MLIR.IR.Context(Reactant.registry[], false)
     context_gc_vector[ctx] = Vector{Union{TracedRArray,TracedRNumber}}(undef, 0)
@@ -2002,6 +2003,15 @@ function compile_xla(f, args; client=nothing, kwargs...)
         global_device_ids = collect(Int64, mlir_fn_res.global_device_ids)
         mlir_fn_res.is_sharded && (device = nothing)
 
+        # XLA.compile mutates the module, for serialization we need to keep a copy
+        if serializable
+            mod_pre_xla = MLIR.IR.Module(
+                MLIR.API.mlirModuleFromOperation(copy(MLIR.IR.Operation(mod)))
+            )
+        else
+            mod_pre_xla = mod
+        end
+
         exec = XLA.compile(
             client,
             device,
@@ -2015,7 +2025,7 @@ function compile_xla(f, args; client=nothing, kwargs...)
             mlir_fn_res.use_shardy_partitioner,
         )
 
-        return mod, exec, mlir_fn_res, device, client
+        return mod_pre_xla, exec, mlir_fn_res, device, client
     finally
         MLIR.IR.deactivate!(ctx)
     end
@@ -2158,16 +2168,24 @@ function compile(f, args; sync=false, kwargs...)
         mlir_fn_res.fnwrapped,
         exec,
         mlir_fn_res.is_sharded ? nothing : device,
+        serializable ? mod : nothing,
     )
 end
 
 # inspired by RuntimeGeneratedFunction.jl
 const __thunk_body_cache = Dict{Symbol,Expr}()
 
-struct Thunk{FTy,tag,IsClosure,ArgTypes,ExecTy,DeviceTy}
+struct Thunk{FTy,tag,IsClosure,ArgTypes,ExecTy,DeviceTy,M<:Union{Nothing,MLIR.IR.Module}}
     f::FTy
     exec::ExecTy
     device::DeviceTy
+    mod::M
+end
+
+function Base.show(
+    io::IO, thunk::Thunk{FTy,tag,IsClosure,ArgTypes,ExecTy,DeviceTy}
+) where {FTy,tag,IsClosure,ArgTypes,ExecTy,DeviceTy}
+    return print(io, "Reactant compiled function $(thunk.f) (with tag $(tag))")
 end
 
 XLA.cost_analysis(thunk::Thunk) = XLA.cost_analysis(thunk.exec)
@@ -2179,11 +2197,16 @@ XLA.get_parameter_shardings(thunk::Thunk) = XLA.get_parameter_shardings(thunk.ex
 struct MisMatchedThunkTypeError{ThunkTy,FoundTypes} <: Base.Exception end
 
 function Base.showerror(
-    io::IO, ece::MisMatchedThunkTypeError{Thunk{FTy,tag,ArgTypes,IsClosure},FoundTypes}
-) where {FTy,tag,ArgTypes,FoundTypes,IsClosure}
+    io::IO,
+    ::MisMatchedThunkTypeError{
+        Thunk{FTy,tag,ArgTypes,IsClosure,ExecTy,DeviceTy},FoundTypes
+    },
+) where {FTy,tag,ArgTypes,FoundTypes,IsClosure,ExecTy,DeviceTy}
     print(
         io,
-        "\nThe Reactant-compiled function `$(Thunk{FTy, tag, ArgTypes, IsClosure})` exists, but no method is defined for this combination of argument types.",
+        "\nThe Reactant-compiled function \
+         `$(Thunk{FTy, tag, ArgTypes, IsClosure, ExecTy, DeviceTy})` exists, but no method \
+         is defined for this combination of argument types.",
     )
     print(
         io,
@@ -2231,10 +2254,19 @@ function register_thunk(
     isclosure::Bool,
     exec,
     device,
+    mod,
 )
     __thunk_body_cache[tag] = body
-    return Thunk{Core.Typeof(f),tag,argtys,isclosure,Core.Typeof(exec),Core.Typeof(device)}(
-        f, exec, device
+    return Thunk{
+        Core.Typeof(f),
+        tag,
+        argtys,
+        isclosure,
+        Core.Typeof(exec),
+        Core.Typeof(device),
+        Core.Typeof(mod),
+    }(
+        f, exec, device, mod
     )
 end
 

src/Reactant.jl

@@ -181,6 +181,8 @@ include("Compiler.jl")
 
 include("Overlay.jl")
 
+include("Serialize.jl")
+
 function Enzyme.make_zero(
     ::Type{RT}, seen::IdDict, prev::RT, ::Val{copy_if_inactive}=Val(false)
 )::RT where {copy_if_inactive,RT<:Union{RArray,RNumber}}

src/Serialize.jl

@@ -0,0 +1,52 @@
+module Serialize
+
+# TODO: move these deps into an extension
+# TODO: Deal with sharding/global devices
+
+using JLD2
+using Reactant: Reactant, MLIR
+
+struct SerializedThunk{FTy,tag,ArgTypes,IsClosure}
+    f::FTy
+    body::Expr
+end
+
+# function JLD2.writeas(
+#     ::Type{<:Reactant.Compiler.Thunk{FTy,tag,ArgTypes,IsClosure}}
+# ) where {FTy,tag,ArgTypes,IsClosure}
+#     return SerializedThunk{FTy,tag,ArgTypes,IsClosure}
+# end
+
+# function JLD2.wconvert(
+#     ::Type{SerializedThunk{FTy,tag,ArgTypes,IsClosure}},
+#     thunk::Reactant.Compiler.Thunk{FTy,tag,ArgTypes,IsClosure},
+# ) where {FTy,tag,ArgTypes,IsClosure}
+#     if thunk.mod === nothing
+#         throw("To serialize a compiled thunk, ensure it is called with `serializable=true`")
+#     end
+
+#     return error("TODO")
+# end
+
+# function JLD2.rconvert(
+#     ::Type{Reactant.Compiler.Thunk{FTy,tag,ArgTypes,IsClosure}},
+#     serialized::SerializedThunk{FTy,tag,ArgTypes,IsClosure},
+# ) where {FTy,tag,ArgTypes,IsClosure}
+#     return error("TODO")
+# end
+
+function serialize(
+    thunk::Reactant.Compiler.Thunk{FTy,tag,ArgTypes,IsClosure}
+) where {FTy,tag,ArgTypes,IsClosure}
+    if thunk.mod === nothing
+        throw("To serialize a compiled thunk, ensure it is called with `serializable=true`")
+    end
+
+    serializable_thunk = SerializedThunk{FTy,tag,ArgTypes,IsClosure}(
+        thunk.f, Reactant.Compiler.__thunk_body_cache[tag]
+    )
+end
+
+function deserialize() end
+
+end