implement new macro @consistent_overlay instead of @assume_effects

Author: aviatesk

base/experimental.jl

@@ -318,7 +318,7 @@ function show_error_hints(io, ex, args...)
     isnothing(hinters) && return
     for handler in hinters
         try
-            Base.invokelatest(handler, io, ex, args...)
+            @invokelatest handler(io, ex, args...)
         catch err
             tn = typeof(handler).name
             @error "Hint-handler $handler for $(typeof(ex)) in $(tn.module) caused an error"
@@ -330,17 +330,97 @@ end
 include("opaque_closure.jl")
 
 """
-    Experimental.@overlay mt [function def]
+    Base.Experimental.@overlay mt [function def]
 
 Define a method and add it to the method table `mt` instead of to the global method table.
 This can be used to implement a method override mechanism. Regular compilation will not
 consider these methods, and you should customize the compilation flow to look in these
 method tables (e.g., using [`Core.Compiler.OverlayMethodTable`](@ref)).
+
+!!! note
+    Please be aware that when defining overlay methods using `@overlay`, it is not necessary
+    to have an original method that corresponds exactly in terms of how the method dispatches.
+    This means that the method overlay mechanism enabled by `@overlay` is not implemented by
+    replacing the methods themselves, but through an additional and prioritized method
+    lookup during the method dispatch.
+
+    Considering this, it is important to understand that in compilations using an overlay
+    method table like the following, the method dispatched by `callx(x)` is not the regular
+    method `callx(::Float64)`, but the overlay method `callx(x::Real)`:
+    ```julia
+    callx(::Real) = :real
+    @overlay SOME_OVERLAY_MT callx(::Real) = :overlay_real
+    callx(::Float64) = :float64
+
+    # some overlay callsite
+    let x::Float64
+        callx(x) #> :overlay_real
+    end
+    ```
 """
 macro overlay(mt, def)
-    def = macroexpand(__module__, def) # to expand @inline, @generated, etc
-    is_function_def(def) || error("@overlay requires a function definition")
-    return esc(overlay_def!(mt, def))
+    inner = Base.unwrap_macrocalls(def)
+    is_function_def(inner) || error("@overlay requires a function definition")
+    overlay_def!(mt, inner)
+    return esc(def)
+end
+
+"""
+    Base.Experimental.@consistent_overlay mt [function def]
+
+This macro operates almost identically to [`Base.Experimental.@overlay`](@ref), defining a
+new overlay method. The key difference with this macro is that it informs the compiler that
+the invocation of the overlay method it defines is `:consistent` with a regular,
+non-overlayed method call.
+
+More formally, when evaluating a generic function call ``f(x)`` at a specific world age
+``i``, if a regular method call ``fᵢ(x)`` is redirected to an overlay method call ``fᵢ′(x)``
+defined by this macro, it must be ensured that ``fᵢ(x) ≡ fᵢ′(x)``.
+
+For a detailed definition of `:consistent`-cy, consult the corresponding section in
+[`Base.@assume_effects`](@ref).
+
+!!! note
+    Note that the requirements for `:consistent`-cy include not only that the return values
+    are egal, but also that the manner of termination is the same.
+    However, it's important to aware that when they throw exceptions, the exceptions
+    themselves don't necessarily have to be egal as explained in the note of `:consistent`.
+    In other words, if ``fᵢ(x)`` throws an exception, ``fᵢ′(x)`` is required to also throw
+    one, but the exact exceptions may differ.
+
+!!! note
+    Please note that the `:consistent`-cy requirement applies not to method itself but to
+    _method invocation_. This means that for the use of `@consistent_overlay`, it is
+    necessary for method invocations with the native regular compilation and those with
+    a compilation with overlay method table to be `:consistent`.
+
+    For example, it is important to understand that, `@consistent_overlay` can be used like
+    the following:
+    ```julia
+    callsin(x::Real) = x < 0 ? error(x) : sin(x)
+    @consistent_overlay SOME_OVERLAY_MT callsin(x::Float64) =
+        x < 0 ? error_somehow(x) : sin(x)
+    ```
+    However, be aware that this `@consistent_overlay` will immediately become invalid if a
+    new method for `callsin` is defined subsequently, such as:
+    ```julia
+    callsin(x::Float64) = cos(x)
+    ```
+
+    This specifically implies that the use of `@consistent_overlay` should be restricted as
+    much as possible to cases where a regular method with a concrete signature is replaced
+    by an overlay method with the same concrete signature.
+
+    This constraint is closely related to the note in [`Base.Experimental.@overlay`](@ref);
+    you are advised to consult that as well.
+"""
+macro consistent_overlay(mt, def)
+    inner = Base.unwrap_macrocalls(def)
+    is_function_def(inner) || error("@consistent_overlay requires a function definition")
+    overlay_def!(mt, inner)
+    override = Core.Compiler.EffectsOverride(; consistent_overlay=true)
+    Base.pushmeta!(def::Expr, Base.form_purity_expr(override))
+    return esc(def)
 end
 
 function overlay_def!(mt, @nospecialize ex)

base/expr.jl

@@ -505,7 +505,6 @@ The following `setting`s are supported.
 - `:inaccessiblememonly`
 - `:noub`
 - `:noub_if_noinbounds`
-- `:consistent_overlay`
 - `:foldable`
 - `:removable`
 - `:total`
@@ -674,29 +673,6 @@ The `:noub` setting asserts that the method will not execute any undefined behav
 any other effect assertions (such as `:consistent` or `:effect_free`) as well, but we do
 not model this, and they assume the absence of undefined behavior.
 
----
-## `:consistent_overlay`
-
-The `:consistent_overlay` setting asserts that any overlayed methods potentially called by
-the method are `:consistent` with their original, non-overlayed counterparts. For the exact
-definition of `:consistent`, refer to the earlier explanation.
-
-More formally, when evaluating a generic function call ``f(x)`` at a specific world-age ``i``,
-and the regular method call ``fᵢ(x)`` is redirected to an overlay method ``fᵢ′(x)``, this
-setting requires that ``fᵢ(x) ≡ fᵢ′(x)``.
-
-!!! note
-    Note that the requirements for `:consistent`-cy include not only that the return values
-    are egal, but also that the manner of termination is the same.
-    However, it's important to aware that when they throw exceptions, the exceptions
-    themselves don't necessarily have to be egal as explained in the note of `:consistent`.
-    In other words, if ``fᵢ(x)`` throws an exception, this settings requires ``fᵢ′(x)`` to
-    also raise one, but the exact exceptions may differ.
-
-!!! note
-    This setting isn't supported at the callsite; it has to be applied at the definition
-    site. Also, given its nature, it's expected to be used together with `Base.Experimental.@overlay`.
-
 ---
 ## `:foldable`
 
@@ -761,7 +737,7 @@ macro assume_effects(args...)
     lastex = args[end]
     override = compute_assumed_settings(args[begin:end-1])
     if is_function_def(unwrap_macrocalls(lastex))
-        return esc(pushmeta!(lastex, form_purity_expr(override)))
+        return esc(pushmeta!(lastex::Expr, form_purity_expr(override)))
     elseif isexpr(lastex, :macrocall) && lastex.args[1] === Symbol("@ccall")
         lastex.args[1] = GlobalRef(Base, Symbol("@ccall_effects"))
         insert!(lastex.args, 3, Core.Compiler.encode_effects_override(override))
@@ -773,8 +749,6 @@ macro assume_effects(args...)
         return Expr(:meta, form_purity_expr(override′))
     else
         # call site annotation case
-        override.consistent_overlay &&
-            throw(ArgumentError("Callsite `@assume_effects :consistent_overlay` is not supported"))
         return Expr(:block,
                     form_purity_expr(override),
                     Expr(:local, Expr(:(=), :val, esc(lastex))),
@@ -819,8 +793,6 @@ function compute_assumed_setting(override::EffectsOverride, @nospecialize(settin
         return EffectsOverride(override; noub = val)
     elseif setting === :noub_if_noinbounds
         return EffectsOverride(override; noub_if_noinbounds = val)
-    elseif setting === :consistent_overlay
-        return EffectsOverride(override; consistent_overlay = val)
     elseif setting === :foldable
         consistent = effect_free = terminates_globally = noub = val
         return EffectsOverride(override; consistent, effect_free, terminates_globally, noub)

src/method.c

@@ -470,16 +470,27 @@ jl_code_info_t *jl_new_code_info_from_ir(jl_expr_t *ir)
                     li->constprop = 2;
                 else if (jl_is_expr(ma) && ((jl_expr_t*)ma)->head == jl_purity_sym) {
                     if (jl_expr_nargs(ma) == NUM_EFFECTS_OVERRIDES) {
-                        li->purity.overrides.ipo_consistent = jl_unbox_bool(jl_exprarg(ma, 0));
-                        li->purity.overrides.ipo_effect_free = jl_unbox_bool(jl_exprarg(ma, 1));
-                        li->purity.overrides.ipo_nothrow = jl_unbox_bool(jl_exprarg(ma, 2));
-                        li->purity.overrides.ipo_terminates_globally = jl_unbox_bool(jl_exprarg(ma, 3));
-                        li->purity.overrides.ipo_terminates_locally = jl_unbox_bool(jl_exprarg(ma, 4));
-                        li->purity.overrides.ipo_notaskstate = jl_unbox_bool(jl_exprarg(ma, 5));
-                        li->purity.overrides.ipo_inaccessiblememonly = jl_unbox_bool(jl_exprarg(ma, 6));
-                        li->purity.overrides.ipo_noub = jl_unbox_bool(jl_exprarg(ma, 7));
-                        li->purity.overrides.ipo_noub_if_noinbounds = jl_unbox_bool(jl_exprarg(ma, 8));
-                        li->purity.overrides.ipo_consistent_overlay = jl_unbox_bool(jl_exprarg(ma, 9));
+                        // N.B. this code allows multiple :purity expressions to be present in a single `:meta` node
+                        int8_t consistent = jl_unbox_bool(jl_exprarg(ma, 0));
+                        if (consistent) li->purity.overrides.ipo_consistent = consistent;
+                        int8_t effect_free = jl_unbox_bool(jl_exprarg(ma, 1));
+                        if (effect_free) li->purity.overrides.ipo_effect_free = effect_free;
+                        int8_t nothrow = jl_unbox_bool(jl_exprarg(ma, 2));
+                        if (nothrow) li->purity.overrides.ipo_nothrow = nothrow;
+                        int8_t terminates_globally = jl_unbox_bool(jl_exprarg(ma, 3));
+                        if (terminates_globally) li->purity.overrides.ipo_terminates_globally = terminates_globally;
+                        int8_t terminates_locally = jl_unbox_bool(jl_exprarg(ma, 4));
+                        if (terminates_locally) li->purity.overrides.ipo_terminates_locally = terminates_locally;
+                        int8_t notaskstate = jl_unbox_bool(jl_exprarg(ma, 5));
+                        if (notaskstate) li->purity.overrides.ipo_notaskstate = notaskstate;
+                        int8_t inaccessiblememonly = jl_unbox_bool(jl_exprarg(ma, 6));
+                        if (inaccessiblememonly) li->purity.overrides.ipo_inaccessiblememonly = inaccessiblememonly;
+                        int8_t noub = jl_unbox_bool(jl_exprarg(ma, 7));
+                        if (noub) li->purity.overrides.ipo_noub = noub;
+                        int8_t noub_if_noinbounds = jl_unbox_bool(jl_exprarg(ma, 8));
+                        if (noub_if_noinbounds) li->purity.overrides.ipo_noub_if_noinbounds = noub_if_noinbounds;
+                        int8_t consistent_overlay = jl_unbox_bool(jl_exprarg(ma, 9));
+                        if (consistent_overlay) li->purity.overrides.ipo_consistent_overlay = consistent_overlay;
                     }
                 }
                 else

test/compiler/AbstractInterpreter.jl

@@ -21,7 +21,7 @@ CC.transform_result_for_cache(::AbsIntOnlyInterp2, ::Core.MethodInstance, ::CC.W
 # OverlayMethodTable
 # ==================
 
-using Base.Experimental: @MethodTable, @overlay
+using Base.Experimental: @MethodTable, @overlay, @consistent_overlay
 
 # @overlay method with return type annotation
 @MethodTable RT_METHOD_DEF
@@ -147,17 +147,29 @@ end
 raise_on_gpu1(x) = error(x)
 @overlay OVERLAY_MT @noinline raise_on_gpu1(x) = #=do something with GPU=# error(x)
 raise_on_gpu2(x) = error(x)
-@overlay OVERLAY_MT @noinline Base.@assume_effects :consistent_overlay raise_on_gpu2(x) = #=do something with GPU=# error(x)
+@consistent_overlay OVERLAY_MT @noinline raise_on_gpu2(x) = #=do something with GPU=# error(x)
+raise_on_gpu3(x) = error(x)
+@consistent_overlay OVERLAY_MT @noinline Base.@assume_effects :foldable raise_on_gpu3(x) = #=do something with GPU=# error_on_gpu(x)
 cpu_factorial(x::Int) = myfactorial(x, error)
 gpu_factorial1(x::Int) = myfactorial(x, raise_on_gpu1)
 gpu_factorial2(x::Int) = myfactorial(x, raise_on_gpu2)
+gpu_factorial3(x::Int) = myfactorial(x, raise_on_gpu3)
 
 @test Base.infer_effects(cpu_factorial, (Int,); interp=MTOverlayInterp()) |> Core.Compiler.is_nonoverlayed
 @test Base.infer_effects(gpu_factorial1, (Int,); interp=MTOverlayInterp()) |> !Core.Compiler.is_nonoverlayed
 @test Base.infer_effects(gpu_factorial2, (Int,); interp=MTOverlayInterp()) |> Core.Compiler.is_consistent_overlay
+let effects = Base.infer_effects(gpu_factorial3, (Int,); interp=MTOverlayInterp())
+    # check if `@consistent_overlay` together works with `@assume_effects`
+    # N.B. the overlaid `raise_on_gpu3` is not :foldable otherwise since `error_on_gpu` is (intetionally) undefined.
+    @test Core.Compiler.is_consistent_overlay(effects)
+    @test Core.Compiler.is_foldable(effects)
+end
 @test Base.infer_return_type(; interp=MTOverlayInterp()) do
     Val(gpu_factorial2(3))
 end == Val{6}
+@test Base.infer_return_type(; interp=MTOverlayInterp()) do
+    Val(gpu_factorial3(3))
+end == Val{6}
 
 # GPUCompiler needs accurate inference through kwfunc with the overlay of `Core.throw_inexacterror`
 # https://github.com/JuliaLang/julia/issues/48097