Add Julia v1.12 compatibility

.buildkite/pipeline.yml

@@ -22,8 +22,9 @@ steps:
     matrix:
       setup:
         version:
-          - "1.11"
           - "1.10"
+          - "1.11"
+          - "1.12"
         label:
           - "cuda"
           - "nnlib"

.github/workflows/CI.yml

@@ -49,6 +49,7 @@ jobs:
         version:
           - 'lts'
           - '1.11'
+          - '1.12'
         arch:
           - x64
         include:
@@ -108,8 +109,9 @@ jobs:
           {test_type: 'integration_testing', label: 'flux'},
         ]
         version:
-          - '1.11'
           - 'lts'
+          - '1.11'
+          - '1.12'
         arch:
           - x64
     steps:

Project.toml

@@ -58,7 +58,7 @@ Flux = "0.16.3"
 FunctionWrappers = "1.1.3"
 GPUArraysCore = "0.1, 0.2"
 Graphs = "1"
-JET = "0.9, 0.10"
+JET = "0.9, 0.10, 0.11"
 LinearAlgebra = "1"
 Logging = "1"
 LuxLib = "1.11"
@@ -72,7 +72,7 @@ SpecialFunctions = "2"
 StableRNGs = "1"
 Static = "1.1.1"
 Test = "1"
-julia = "~1.10.8, 1.11.6"
+julia = "~1.10.8, ~1.11.6, 1.12.1"
 
 [extras]
 AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"

ext/MooncakeFunctionWrappersExt.jl

@@ -80,7 +80,9 @@ end
 function _function_wrapper_tangent(R, obj::Tobj, A, obj_tangent) where {Tobj}
 
     # Analyse types.
-    _, _, fwd_sig, rvs_sig = _construct_types(R, A)
+    fwd_oc_type, rvs_oc_type, _, _ = _construct_types(R, A)
+    (fwd_sig, fwd_ret) = fwd_oc_type.parameters
+    (rvs_sig, rvs_ret) = rvs_oc_type.parameters
 
     # Construct reference to obj_tangent that we can read / write-to.
     obj_tangent_ref = Ref{tangent_type(Tobj)}(obj_tangent)
@@ -93,17 +95,33 @@ function _function_wrapper_tangent(R, obj::Tobj, A, obj_tangent) where {Tobj}
     pb_stack = Stack{pullback_type(typeof(rule), (Tobj, A.parameters...))}()
 
     # Construct reverse-pass. Note: this closes over `pb_stack`.
-    run_rvs_pass = Base.Experimental.@opaque rvs_sig dy -> begin
-        obj_rdata, dx... = pop!(pb_stack)(dy)
-        obj_tangent_ref[] = increment_rdata!!(obj_tangent_ref[], obj_rdata)
-        return NoRData(), dx...
+    @static if VERSION ≥ v"1.12-"
+        run_rvs_pass = Base.Experimental.@opaque rvs_sig -> rvs_ret dy -> begin
+            obj_rdata, dx... = pop!(pb_stack)(dy)
+            obj_tangent_ref[] = increment_rdata!!(obj_tangent_ref[], obj_rdata)
+            return NoRData(), dx...
+        end
+    else
+        run_rvs_pass = Base.Experimental.@opaque rvs_sig dy -> begin
+            obj_rdata, dx... = pop!(pb_stack)(dy)
+            obj_tangent_ref[] = increment_rdata!!(obj_tangent_ref[], obj_rdata)
+            return NoRData(), dx...
+        end
     end
 
     # Construct fowards-pass. Note: this closes over the reverse-pass and `pb_stack`.
-    run_fwds_pass = Base.Experimental.@opaque fwd_sig (x...) -> begin
-        y, pb = rule(CoDual(obj, fdata(obj_tangent_ref[])), x...)
-        push!(pb_stack, pb)
-        return y, run_rvs_pass
+    @static if VERSION ≥ v"1.12-"
+        run_fwds_pass = Base.Experimental.@opaque fwd_sig -> fwd_ret (x...) -> begin
+            y, pb = rule(CoDual(obj, fdata(obj_tangent_ref[])), x...)
+            push!(pb_stack, pb)
+            return y, run_rvs_pass
+        end
+    else
+        run_fwds_pass = Base.Experimental.@opaque fwd_sig (x...) -> begin
+            y, pb = rule(CoDual(obj, fdata(obj_tangent_ref[])), x...)
+            push!(pb_stack, pb)
+            return y, run_rvs_pass
+        end
     end
 
     t = FunctionWrapperTangent(run_fwds_pass, obj_tangent_ref)

src/Mooncake.jl

@@ -161,8 +161,8 @@ include("public.jl")
 end
 #! format: on
 
-@public Config, value_and_pullback!!, prepare_pullback_cache, value_and_derivative!!
-@public prepare_derivative_cache, Dual
+@public Config, value_and_pullback!!, prepare_pullback_cache
+@public Dual
 
 # Public, exported
 export value_and_gradient!!, prepare_gradient_cache, value_and_derivative!!

src/interpreter/abstract_interpretation.jl

@@ -2,7 +2,7 @@
 # in conjunction with the contexts above to decide what should be inlined and what should
 # not be inlined. Similar strategies are employed by Enzyme and Diffractor.
 
-# The most important bit of this code is `inlining_policy` -- the rest is copy + pasted
+# The most important bit of this code is `inlining_policy` (renamed to `src_inlining_policy` in Julia v1.12+) -- the rest is copy + pasted
 # boiler plate, largely taken from https://github.com/JuliaLang/julia/blob/2fe4190b3d26b4eee52b2b1b1054ddd6e38a941e/test/compiler/newinterp.jl#L11
 #
 # Credit: much of the code in here is copied over from the main Julia repo, and from
@@ -121,6 +121,11 @@ end
 CC.nsplit_impl(info::NoInlineCallInfo) = CC.nsplit(info.info)
 CC.getsplit_impl(info::NoInlineCallInfo, idx::Int) = CC.getsplit(info.info, idx)
 CC.getresult_impl(info::NoInlineCallInfo, idx::Int) = CC.getresult(info.info, idx)
+@static if VERSION > v"1.12-"
+    CC.add_edges_impl(edges::Vector{Any}, info::NoInlineCallInfo) = CC.add_edges!(
+        edges, info.info
+    )
+end
 
 function Core.Compiler.abstract_call_gf_by_type(
     interp::MooncakeInterpreter{C,M},
@@ -133,7 +138,7 @@ function Core.Compiler.abstract_call_gf_by_type(
 ) where {C,M}
 
     # invoke the default abstract call to get the default CC.CallMeta.
-    cm = @invoke CC.abstract_call_gf_by_type(
+    ret = @invoke CC.abstract_call_gf_by_type(
         interp::CC.AbstractInterpreter,
         f::Any,
         arginfo::CC.ArgInfo,
@@ -142,16 +147,25 @@ function Core.Compiler.abstract_call_gf_by_type(
         sv::CC.AbsIntState,
         max_methods::Int,
     )
+    is_primitive(C, M, atype) || return ret
+    # Insert a `NoInlineCallInfo` to prevent any potential inlining.
+    @static if VERSION < v"1.12-"
+        call = ret::CC.CallMeta
+        info = NoInlineCallInfo(call.info, atype)
+        return rewrap_callmeta(call, info)
+    else
+        return CC.Future{CC.CallMeta}(ret::CC.Future, interp, sv) do call, interp, sv
+            info = NoInlineCallInfo(call.info, atype)
+            return rewrap_callmeta(call, info)
+        end
+    end
+end
 
-    # Check to see whether the call in question is a Mooncake primitive. If it is, set its
-    # call info such that in the `CC.inlining_policy` it is not inlined away.
-    callinfo = is_primitive(C, M, atype) ? NoInlineCallInfo(cm.info, atype) : cm.info
-
-    # Construct a CallMeta correctly depending on the version of Julia.
+function rewrap_callmeta(call::CC.CallMeta, info::CC.CallInfo)
     @static if VERSION ≥ v"1.11-"
-        return CC.CallMeta(cm.rt, cm.exct, cm.effects, callinfo)
+        return CC.CallMeta(call.rt, call.exct, call.effects, info)
     else
-        return CC.CallMeta(cm.rt, cm.effects, callinfo)
+        return CC.CallMeta(call.rt, call.effects, info)
     end
 end
 
@@ -179,7 +193,7 @@ end
         )
     end
 
-else # 1.11 and up.
+elseif VERSION < v"1.12-" # 1.11
     function CC.inlining_policy(
         interp::MooncakeInterpreter,
         @nospecialize(src),
@@ -194,6 +208,22 @@ else # 1.11 and up.
             interp::CC.AbstractInterpreter, src::Any, info::CC.CallInfo, stmt_flag::UInt32
         )
     end
+
+else # 1.12 and up.
+    function CC.src_inlining_policy(
+        interp::MooncakeInterpreter,
+        @nospecialize(src),
+        @nospecialize(info::CC.CallInfo),
+        stmt_flag::UInt32,
+    )
+        # Do not inline away primitives.
+        info isa NoInlineCallInfo && return false
+
+        # If not a primitive, AD doesn't care about it. Use the usual inlining strategy.
+        return @invoke CC.src_inlining_policy(
+            interp::CC.AbstractInterpreter, src::Any, info::CC.CallInfo, stmt_flag::UInt32
+        )
+    end
 end
 
 """

src/interpreter/bbcode.jl

@@ -272,45 +272,57 @@ These should be assumed to be ordered.
 """
 collect_stmts(bb::BBlock)::Vector{IDInstPair} = collect(zip(bb.inst_ids, bb.insts))
 
-"""
-    BBCode(
+@eval begin
+    """
+        BBCode(
+            blocks::Vector{BBlock}
+            argtypes::Vector{Any}
+            sptypes::Vector{CC.VarState}
+            linetable::Vector{Core.LineInfoNode} (v1.11 and lower)
+            debuginfo::CC.DebugInfoStream (v1.12+)
+            meta::Vector{Expr}
+            valid_worlds::CC.WorldRange (v1.12+)
+        )
+
+    A `BBCode` is a data structure which is similar to `IRCode`, but adds additional structure.
+
+    In particular, a `BBCode` comprises a sequence of basic blocks (`BBlock`s), each of which
+    comprise a sequence of statements. Moreover, each `BBlock` has its own unique `ID`, as does
+    each statment.
+
+    The consequence of this is that new basic blocks can be inserted into a `BBCode`. This is
+    distinct from `IRCode`, in which to create a new basic block, one must insert additional
+    statments which you know will create a new basic block -- this is generally quite an
+    unreliable process, while inserting a new `BBlock` into `BBCode` is entirely predictable.
+    Furthermore, inserting a new `BBlock` does not change the `ID` associated to the other
+    blocks, meaning that you can safely assume that references from existing basic block
+    terminators / phi nodes to other blocks will not be modified by inserting a new basic block.
+
+    Additionally, since each statment in each basic block has its own unique `ID`, new
+    statments can be inserted without changing references between other blocks. `IRCode` also
+    has some support for this via its `new_nodes` field, but eventually all statements will be
+    renamed upon `compact!`ing the `IRCode`, meaning that the name of any given statement will
+    eventually change.
+
+    Finally, note that the basic blocks in a `BBCode` support the custom `Switch` statement.
+    This statement is not valid in `IRCode`, and is therefore lowered into a collection of
+    `GotoIfNot`s and `GotoNode`s when a `BBCode` is converted back into an `IRCode`.
+    """
+    struct BBCode
         blocks::Vector{BBlock}
         argtypes::Vector{Any}
         sptypes::Vector{CC.VarState}
-        linetable::Vector{Core.LineInfoNode}
+        $(
+            if VERSION > v"1.12-"
+                :(debuginfo::CC.DebugInfoStream)
+            else
+                :(linetable::Vector{Core.LineInfoNode})
+            end
+        )
         meta::Vector{Expr}
-    )
-
-A `BBCode` is a data structure which is similar to `IRCode`, but adds additional structure.
-
-In particular, a `BBCode` comprises a sequence of basic blocks (`BBlock`s), each of which
-comprise a sequence of statements. Moreover, each `BBlock` has its own unique `ID`, as does
-each statment.
-
-The consequence of this is that new basic blocks can be inserted into a `BBCode`. This is
-distinct from `IRCode`, in which to create a new basic block, one must insert additional
-statments which you know will create a new basic block -- this is generally quite an
-unreliable process, while inserting a new `BBlock` into `BBCode` is entirely predictable.
-Furthermore, inserting a new `BBlock` does not change the `ID` associated to the other
-blocks, meaning that you can safely assume that references from existing basic block
-terminators / phi nodes to other blocks will not be modified by inserting a new basic block.
-
-Additionally, since each statment in each basic block has its own unique `ID`, new
-statments can be inserted without changing references between other blocks. `IRCode` also
-has some support for this via its `new_nodes` field, but eventually all statements will be
-renamed upon `compact!`ing the `IRCode`, meaning that the name of any given statement will
-eventually change.
-
-Finally, note that the basic blocks in a `BBCode` support the custom `Switch` statement.
-This statement is not valid in `IRCode`, and is therefore lowered into a collection of
-`GotoIfNot`s and `GotoNode`s when a `BBCode` is converted back into an `IRCode`.
-"""
-struct BBCode
-    blocks::Vector{BBlock}
-    argtypes::Vector{Any}
-    sptypes::Vector{CC.VarState}
-    linetable::Vector{Core.LineInfoNode}
-    meta::Vector{Expr}
+        $(VERSION > v"1.12-" ? :(valid_worlds::CC.WorldRange) : nothing)
+    end
+    export BBCode
 end
 
 """
@@ -319,14 +331,27 @@ end
 Make a new `BBCode` whose `blocks` is given by `new_blocks`, and fresh copies are made of
 all other fields from `ir`.
 """
-function BBCode(ir::Union{IRCode,BBCode}, new_blocks::Vector{BBlock})
-    return BBCode(
-        new_blocks,
-        CC.copy(ir.argtypes),
-        CC.copy(ir.sptypes),
-        CC.copy(ir.linetable),
-        CC.copy(ir.meta),
-    )
+@static if VERSION > v"1.12-"
+    function BBCode(ir::Union{IRCode,BBCode}, new_blocks::Vector{BBlock})
+        return BBCode(
+            new_blocks,
+            CC.copy(ir.argtypes),
+            CC.copy(ir.sptypes),
+            CC.copy(ir.debuginfo),
+            CC.copy(ir.meta),
+            ir.valid_worlds,
+        )
+    end
+else
+    function BBCode(ir::Union{IRCode,BBCode}, new_blocks::Vector{BBlock})
+        return BBCode(
+            new_blocks,
+            CC.copy(ir.argtypes),
+            CC.copy(ir.sptypes),
+            CC.copy(ir.linetable),
+            CC.copy(ir.meta),
+        )
+    end
 end
 
 # Makes use of the above outer constructor for `BBCode`.
@@ -640,20 +665,47 @@ function CC.IRCode(bb_code::BBCode)
     insts = _ids_to_line_numbers(bb_code)
     cfg = control_flow_graph(bb_code)
     insts = _lines_to_blocks(insts, cfg)
-    return IRCode(
-        CC.InstructionStream(
-            map(x -> x.stmt, insts),
-            map(x -> x.type, insts),
-            map(x -> x.info, insts),
-            map(x -> x.line, insts),
-            map(x -> x.flag, insts),
-        ),
-        cfg,
-        CC.copy(bb_code.linetable),
-        CC.copy(bb_code.argtypes),
-        CC.copy(bb_code.meta),
-        CC.copy(bb_code.sptypes),
-    )
+    @static if VERSION > v"1.12-"
+        lines = CC.copy(bb_code.debuginfo.codelocs)
+        n = length(insts)
+        if length(lines) > 3n
+            resize!(lines, 3n)
+        elseif length(lines) < 3n
+            for _ in (length(lines) + 1):3n
+                push!(lines, 0)
+            end
+        end
+        return IRCode(
+            CC.InstructionStream(
+                Any[x.stmt for x in insts],
+                Any[x.type for x in insts],
+                CC.CallInfo[x.info for x in insts],
+                lines,
+                UInt32[x.flag for x in insts],
+            ),
+            cfg,
+            CC.copy(bb_code.debuginfo),
+            CC.copy(bb_code.argtypes),
+            CC.copy(bb_code.meta),
+            CC.copy(bb_code.sptypes),
+            bb_code.valid_worlds,
+        )
+    else
+        return IRCode(
+            CC.InstructionStream(
+                Any[x.stmt for x in insts],
+                Any[x.type for x in insts],
+                CC.CallInfo[x.info for x in insts],
+                Int32[x.line for x in insts],
+                UInt32[x.flag for x in insts],
+            ),
+            cfg,
+            CC.copy(bb_code.linetable),
+            CC.copy(bb_code.argtypes),
+            CC.copy(bb_code.meta),
+            CC.copy(bb_code.sptypes),
+        )
+    end
 end
 
 """