Type instability in code with branches and QuadGK uses

General Usage
1

I have a code which calculates complicated integral and which one should be used depends on the parameters. The code is chronically type-unstable (union with Core.Box appears) which makes it few times slower than it should be. The MWE is like this

using QuadGK

function f1(x,b)
    x < 0 && ( b = -b )
    I, _err = quadgk(t->exp(-t*b),0,1)
    return I
end

You can also try b = sign(x)*b, result is the same. In this case I can fix it by switching to

function f2(x,b)
    I, _err = if x < 0
         quadgk(t->exp(-t*b),0,1)
    else
        quadgk(t->exp(t*b),0,1)
    end
    return I
end

But it is hard to understand why this behaviour occurs, and in the full code I have multiple conditions to deal with.

2

Reassigning a captured variable like b at any point currently makes the lowerer implement it as a Core.Box field in the closure’s type. In your fix, you don’t reassign b, you just switch between different closures. Alternatively, you could capture a new local variable b2 = if x < 0; -b else b end or the branchless b2 = ifelse(x<0, -b, b).

3

The code

function f3(x,b)
   b2 = ifelse(x < 0, -b, b)
   I, _err = quadgk(t->exp(-t*b2),0,1)
   return I
end

is still type-unstable.

4

I think this is the usual problem with captured mutated variables in closures (julia#15276). A simple workaround is to re-assign b to a local variable in a let block before capturing it:

I, _err = let b=b; quadgk(t->exp(-t*b),0,1); end
5

Another option is to use function composition and Base.Fix:

function f3(x,b)
     x < 0 && ( b = -b )
     I, _err = quadgk(exp∘Base.Fix1(*,-b),0,1)
     return I
 end
6

I can’t see it in a cursory @code_warntype, what are you using for type inference?

julia> @code_warntype f3(-1.0, 2.0)
MethodInstance for f3(::Float64, ::Float64)
  from f3(x, b) @ Main REPL[36]:1
Arguments
  #self#::Core.Const(Main.f3)
  x::Float64
  b::Float64
Locals
  #2::var"#f3##2#f3##3"{Float64}
  @_5::Int64
  _err::Float64
  I::Float64
  b2::Float64
Body::Float64
1 ─ %1  = Main.ifelse::Core.Const(ifelse)
│   %2  = Main.:<::Core.Const(<)
│   %3  = (%2)(x, 0)::Bool
│   %4  = Main.:-::Core.Const(-)
│   %5  = (%4)(b)::Float64
│         (b2 = (%1)(%3, %5, b))
│   %7  = Main.quadgk::Core.Const(QuadGK.quadgk)
│   %8  = Main.:(var"#f3##2#f3##3")::Core.Const(var"#f3##2#f3##3")
│   %9  = b2::Float64
│   %10 = Core._typeof_captured_variable(%9)::Core.Const(Float64)
│   %11 = Core.apply_type(%8, %10)::Core.Const(var"#f3##2#f3##3"{Float64})
│   %12 = b2::Float64
│         (#2 = %new(%11, %12))
│   %14 = #2::var"#f3##2#f3##3"{Float64}
│   %15 = (%7)(%14, 0, 1)::Tuple{Float64, Float64}
│   %16 = Base.indexed_iterate(%15, 1)::Core.PartialStruct(Tuple{Float64, Int64}, Any[Float64, Core.Const(2)])
│         (I = Core.getfield(%16, 1))
│         (@_5 = Core.getfield(%16, 2))
│   %19 = @_5::Core.Const(2)
│   %20 = Base.indexed_iterate(%15, 2, %19)::Core.PartialStruct(Tuple{Float64, Int64}, Any[Float64, Core.Const(3)])
│         (_err = Core.getfield(%20, 1))
│   %22 = I::Float64
└──       return %22
7

Now this is just bizarre. Function

function f3(x,b)
          b2 = ifelse(x < 0, -b, b)
          I, _err = quadgk(t->exp(-t*b2),0,1)
          return I
       end

gives

@code_warntype f3(-1.,1.)
MethodInstance for f3(::Float64, ::Float64)
  from f3(x, b) @ Main REPL[4]:1
Arguments
  #self#::Core.Const(Main.f3)
  x::Float64
  b::Float64
Locals
  #1::var"#f3##0#f3##1"{Float64}
  @_5::Any
  _err::Any
  I::Any
  b2::Float64
Body::Any
1 ─ %1  = Main.ifelse::Core.Const(ifelse)
│   %2  = Main.:<::Core.Const(<)
│   %3  = (%2)(x, 0)::Bool
│   %4  = Main.:-::Core.Const(-)
│   %5  = (%4)(b)::Float64
│         (b2 = (%1)(%3, %5, b))
│   %7  = Main.quadgk::Any
│   %8  = Main.:(var"#f3##0#f3##1")::Core.Const(var"#f3##0#f3##1")
│   %9  = b2::Float64
│   %10 = Core._typeof_captured_variable(%9)::Core.Const(Float64)
│   %11 = Core.apply_type(%8, %10)::Core.Const(var"#f3##0#f3##1"{Float64})
│   %12 = b2::Float64
│         (#1 = %new(%11, %12))
│   %14 = #1::var"#f3##0#f3##1"{Float64}
│   %15 = (%7)(%14, 0, 1)::Any
│   %16 = Base.indexed_iterate(%15, 1)::Any
│         (I = Core.getfield(%16, 1))
│         (@_5 = Core.getfield(%16, 2))
│   %19 = @_5::Any
│   %20 = Base.indexed_iterate(%15, 2, %19)::Any
│         (_err = Core.getfield(%20, 1))
│   %22 = I::Any
└──       return %22

Is it version dependent? I use Julia 1.12.5.

8

The part of your @code_warntype f3(-1.,1.) printout that pertains to the closure is concretely typed. The type instability starts with Main.quadgk::Any, compared with my Main.quadgk::Core.Const(QuadGK.quadgk). I can replicate that if I did using QuadGK; quadgk = QuadGK.quadgk in a module, though it’s hard for me to imagine you doing that. The const-ness of global variables is irreversible, but you can still try using QuadGK: quadgk and a duplicate f3 in a fresh module or session.

I doubt it’s version dependent, but I’m using QuadGK v2.11.3, Julia v1.12.6.

9

All right, sorry for the confusion. I can’t say for sure what happened before but likely at some point I switched to Julia instance with QuadGK unloaded, which did not change the result much for other variants but obviously did for ifelse one. Thanks for checking thoroughly!

10

Thank you for all the responses. I managed to remove all instabilities. I defined internal functions too call them out in main if else statements, and then I treated smaller branches inside these functions with ifelse. The function now runs 3x faster.