SpecialFunctions.jl icon indicating copy to clipboard operation
SpecialFunctions.jl copied to clipboard

Allocations in the sum over a product of `besselj(::Int64, ::Float32)`

Open yakovbraver opened this issue 1 year ago • 0 comments

Consider summing a product of three besselj(n, x)s when the type of n is Int64 (the default Int on a 64-bit machine) and type of x is Float32:

import SpecialFunctions.besselj as J
function jtest()
    n = 1
    x = 1.0f0
    sum(J(n+q, x) * J(n+q, x) * J(n+q, x) for q in 1:2)
end

This code allocates:

@time jtest() # 0.000004 seconds (8 allocations: 128 bytes)

The type of the final result is inferred as Union{Float32, Float64} instead of Float32:

Arguments
  #self#::Core.Const(jtest)
Locals
  #19::var"#19#20"{Float32, Int64}
  x::Float32
  n::Int64
Body::Union{Float32, Float64}
1 ─       (n = 1)
│         (x = 1.0f0)
│   %3  = Main.:(var"#19#20")::Core.Const(var"#19#20")
│   %4  = Core.typeof(x::Core.Const(1.0f0))::Core.Const(Float32)
│   %5  = Core.typeof(n::Core.Const(1))::Core.Const(Int64)
│   %6  = Core.apply_type(%3, %4, %5)::Core.Const(var"#19#20"{Float32, Int64})
│   %7  = x::Core.Const(1.0f0)
│         (#19 = %new(%6, %7, n::Core.Const(1)))
│   %9  = #19::Core.Const(var"#19#20"{Float32, Int64}(1.0f0, 1))
│   %10 = (1:2)::Core.Const(1:2)
│   %11 = Base.Generator(%9, %10)::Core.Const(Base.Generator{UnitRange{Int64}, var"#19#20"{Float32, Int64}}(var"#19#20"{Float32, Int64}(1.0f0, 1), 1:2))
│   %12 = Main.sum(%11)::Union{Float32, Float64}
└──       return %12

One workaround is to cast the first argument to Int32, i.e. change J(n+q, x) to J(Int32(n+q), x). When summing products of only two Bessels, the problem does not appear.

BTW, besselj from Bessels.jl does not suffer from this issue 😀

yakovbraver avatar May 01 '24 13:05 yakovbraver