Improving speed of runtime dispatch detector

I don’t know if module-wise is impossible (not saying it isn’t), but could at calling site work?

@inline used to only work at the definition, and then later also at the caller site. I’m not sure how it works, or why, I think it’s still one (improved) macro… [I suppose it’s really two in one, and it needs to detect where it’s applied.]

@inline is not recursive (nor would you want it to be), but since it’s possible for the calling site, could you do the same, and even recursive…? I mean theoretically, not pushing you do it personally. And I mean @stable f(x) though also intriguing would be: stable using <module>, could such work?

Okay all is good, I fixed it here by moving back to the closure approach:

function _stable(fex::Expr)
    func = splitdef(fex)

    ...
    @gensym closure T
    func[:body] = quote
        let $closure() = $(func[:body]), $T = $(Base).promote_op($closure)
            if !$(Base).isconcretetype($T)
                ...
            end

            return $closure()::$T
        end
    end

    return combinedef(func)
end

It seems like promote_op gets the same results if you put all the function body in a closure. And, as before, the check looks to compile away if everything is type-stable.

I have not tried it, but maybe you can use @stable module A to rewire A.include(x) to make it default to A.include(f, x), where f is a map injecting @stable into function definitions?

Clever! I’ll try that.

You can define something like

struct StableType{T} end
_typeof(x) = typeof(x)
_typeof(::StableType{T}) where {T} = Type{T}

f(a, t::Type{T}) where {T} = sum(a; init=zero(T))
x = rand(4); ST = StableType{Float64}();
Base.promote_op(f, map(_typeof, (x, ST))...)

I get

julia> Base.promote_op(f, map(_typeof, (x, ST))...)
Float64

Of course, also define the standard unwrap that returns just T for passing to f.

Probably

makestable(x) = x
makestable(::Type{T}) where {T} = StableType{T}()
unwrapstable(x) = x
unwrapstable(::StableType{T}) where {T} = T

then call stable_wrap(f, map(makestable, args...)...) and have then have stable_wrap itself call f(map(unwrapsbale, args...)...), along with the Base.promote_op(f, map(_typeof, args)...).

1 Like

Is there any reason to avoid the closure approach above? It seems to work just as well, hugely simplifies the codebase, and also means one can safely use @stable on functions without variable names:

@stable f(x, ::Type{T}) where {T} = ...

whereas with a separate wrapper function, this would not be permissible.

Btw, I got the module-level approach working: Add support for module-level `@stable` by MilesCranmer · Pull Request #4 · MilesCranmer/DispatchDoctor.jl · GitHub

Makes it possible to use on entire codebases:

@stable module A
    using DispatchDoctor: @unstable

    @unstable f1() = rand(Bool) ? 0 : 1.0
    f2(x) = x
    f3(; a=1) = a > 0 ? a : 0.0

    include("other_code.jl")
end
  • @unstable turns it off (in case you have a function you need to be unstable)
  • With @thofma’s idea (thanks!), include(s) is automatically mapped to include(_stable_all_fnc, s) which propagates the @stable through the included code
4 Likes

Wow, this looks incredibly useful! Some thoughts:

  • Currently, neither function-level nor module-level @stable apply recursively to functions defined within the scope of other functions. In other words, g does not receive the @stable treatment in either of the following examples. Is this what you prefer?
    @stable function f(xs)
        g(x) = -2x
        return map(g, xs)
    end
    
    @stable module A
        function f(xs)
            g(x) = -2x
            return map(g, xs)
        end
    end
    
  • Perhaps it would be useful to allow using @stable on a begin-end block as well:
    @stable begin
        # stable defs
    end
    
    I don’t know a lot of metaprogramming, but it looks like this could be achieved simply by replacing _stable_fnc with _stable_all_fnc in the non-module branch here: DispatchDoctor.jl/src/DispatchDoctor.jl at main · MilesCranmer/DispatchDoctor.jl · GitHub

Does it work with nested include?

1 Like

Yeah this was on purpose. There’s even a unittest against such behavior here: DispatchDoctor.jl/test/runtests.jl at 74af05c9f98160c6cdeb7fff246a574604047cf5 · MilesCranmer/DispatchDoctor.jl · GitHub. Only top-level functions are wrapped, no closures.

Since @stable only applies to functions, I think this may be the wrong pattern to use, as it might give the impression that it detects instability in arbitrary code, which it does not. To me module makes more sense because that’s where functions live.

(You should be able to use @stable explicitly on a closure function though)

It should, since all it does is put

function include(path::AbstractString)
    return include(_stable_all_fnc, path)
end

at the top of the module definition. But note it doesn’t recursively apply to submodules.

2 Likes

There seem to be a couple of caveats regarding submodules. Firstly, every module and submodule within @stable requires its own using DispatchDoctor, otherwise you’ll get an error:

julia> @eval @stable module A
           using DispatchDoctor
           foo(x) = x
           module B
               bar(x) = -x
           end
       end
ERROR: syntax: module expression third argument must be a block
[...]

julia> @eval @stable module A
           using DispatchDoctor
           foo(x) = x
           module B
               using DispatchDoctor
               bar(x) = -x
           end
       end
WARNING: replacing module A.
Main.A

Secondly, @stable will apply to functions defined in the same file as the module statement for the submodule (it just won’t propagate through include within the submodule). In the code above, A.B.bar gets the treatment:

julia> @code_lowered A.B.bar(1.0)
CodeInfo(
1 ─ %1  = Main.A.B.:(var"##bar_closure#235#1")
│   %2  = Core.typeof(x)
│   %3  = Core.apply_type(%1, %2)
│         bar_closure#235 = %new(%3, x)
│   %5  = Base.getproperty(Base, :promote_op)
│   %6  = (%5)(bar_closure#235)
│         bar_return_type#236 = %6
│   %8  = Base.getproperty(Base, :isconcretetype)
│   %9  = (%8)(bar_return_type#236)
│   %10 = !%9
└──       goto #3 if not %10
2 ─ %12 = Main.A.B.bar
│   %13 = Core.tuple(x)
│   %14 = Core.NamedTuple()
│   %15 = (TypeInstabilityError)(%12, %13, %14, bar_return_type#236)
└──       Main.A.B.throw(%15)
3 ┄ %17 = (bar_closure#235)()
│   %18 = Core.typeassert(%17, bar_return_type#236)
└──       return %18
)

Fair enough! I think applying to begin-end could be practically useful in many cases, say if you want only a subset of functions within a module to be @stable, perhaps every method of a function or every constructor for a type. But I get where you’re coming from. However, @unstable can be applied to a begin-end block as-is, so there’s a slight asymmetry there.

The macro could add using DispatchDoctor: @stable.

Actually, it looks like import DispatchDoctor is sufficient. No names needed.

Oh yeah I spoke incorrectly. There’s actually no logic to skip submodules, so I guess it will work. It’s only that include within submodules won’t apply the transformations. The only time the recursion stops is for @stable (to avoid repeats) or @unstable. But other than that it will just apply to all function definitions it can find.

Should we add functionality to recursively overload include maybe? It should be pretty simple – can just add that overload on any module inside the wrapped expression.

I’m confused why that error shows up for you, but not in any unit tests… See DispatchDoctor.jl/test/runtests.jl at 74af05c9f98160c6cdeb7fff246a574604047cf5 · MilesCranmer/DispatchDoctor.jl · GitHub

Are you using the latest commit? Could have been a bug from an earlier revision.

Yes. In a fresh Julia session:

(@v1.10) pkg> activate --temp
  Activating new project at `/tmp/jl_AK8zpN`

(jl_AK8zpN) pkg> add https://github.com/MilesCranmer/DispatchDoctor.jl
    Updating git-repo `https://github.com/MilesCranmer/DispatchDoctor.jl`
   Resolving package versions...
    Updating `/tmp/jl_AK8zpN/Project.toml`
  [8d63f2c5] + DispatchDoctor v0.1.0 `https://github.com/MilesCranmer/DispatchDoctor.jl#main`
    Updating `/tmp/jl_AK8zpN/Manifest.toml`
  [8d63f2c5] + DispatchDoctor v0.1.0 `https://github.com/MilesCranmer/DispatchDoctor.jl#main`
  [1914dd2f] + MacroTools v0.5.13
  [1c621080] + TestItems v0.1.1
  [2a0f44e3] + Base64
  [d6f4376e] + Markdown
  [9a3f8284] + Random
  [ea8e919c] + SHA v0.7.0

julia> using DispatchDoctor

julia> @eval @stable module A
           import DispatchDoctor
           foo(x) = x
           module B
               bar(x) = -x
           end
       end
ERROR: syntax: module expression third argument must be a block
Stacktrace:
 [1] eval(m::Module, e::Any)
   @ Core ./boot.jl:385
 [2] top-level scope
   @ REPL[4]:1

Same if I put the code, sans @eval, in a file and use include.


julia> versioninfo()
Julia Version 1.10.3
Commit 0b4590a5507 (2024-04-30 10:59 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 8 × Intel(R) Core(TM) i7-7700HQ CPU @ 2.80GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, skylake)
Threads: 1 default, 0 interactive, 1 GC (on 8 virtual cores)

Oh, you just mean submodules? But the top-level module doesn’t need that, right?

I’m afraid I need it at top level too:

julia> @eval @stable module C
           foo(x) = x
       end
ERROR: LoadError: AssertionError: module_body.head == :block
Stacktrace:
 [1] _stable_module(ex::Expr)
   @ DispatchDoctor ~/.julia/packages/DispatchDoctor/J9g5M/src/DispatchDoctor.jl:47
 [2] _stable(ex::Expr)
   @ DispatchDoctor ~/.julia/packages/DispatchDoctor/J9g5M/src/DispatchDoctor.jl:37
 [3] var"@stable"(__source__::LineNumberNode, __module__::Module, fex::Any)
   @ DispatchDoctor ~/.julia/packages/DispatchDoctor/J9g5M/src/DispatchDoctor.jl:136
 [4] eval(m::Module, e::Any)
   @ Core ./boot.jl:385
 [5] top-level scope
   @ REPL[9]:1
in expression starting at REPL[9]:1

julia> @eval @stable module C
           import DispatchDoctor
           foo(x) = x
       end
Main.C

Oh I think it’s not the import, but rather the fact that your module is a single line of code — the macro expects it to be a block. If you have any import there it should work. So it’s a bug but unrelated to the import (I think)

2 Likes