u"A^(3/2)" fails to infer

[ggggggggg]

I need to work with some values with non-integer powers, and ran into some issue getting code to run without allocations. I think it's due to an inference failure. In the following example (with Unitful 1.3) I call Base.literal_pow explicitly to try to make sure the exponent is known at compile time, and despite this the output of @code_warntype shows the return value as ::Any and @btime shows 3 allocation.

Related discourse post: https://discourse.julialang.org/t/calculate-i-a-2-3-with-no-allocations-with-unitful/42998/2

julia> I = 1u"A"
1 A

julia> A = 1.0u"A/K^(3/2)"
1.0 A K⁻³ᐟ²

julia> f(I, A) = Base.literal_pow(^,I/A, Val(2//3))
f (generic function with 1 method)

julia> @code_warntype f(I, A)
Variables
  #self#::Core.Compiler.Const(f, false)
  I::Quantity{Int64,𝐈,Unitful.FreeUnits{(A,),𝐈,nothing}}
  A::Quantity{Float64,𝐈 𝚯⁻³ᐟ²,Unitful.FreeUnits{(A, K⁻³ᐟ²),𝐈 𝚯⁻³ᐟ²,nothing}}

Body::Any
1 ─ %1 = Base.literal_pow::Core.Compiler.Const(Base.literal_pow, false)
│   %2 = (I / A)::Quantity{Float64,𝚯³ᐟ²,Unitful.FreeUnits{(K³ᐟ²,),𝚯³ᐟ²,nothing}}
│   %3 = (2 // 3)::Rational{Int64}
│   %4 = Main.Val(%3)::Val{_A} where _A
│   %5 = (%1)(Main.:^, %2, %4)::Any
└──      return %5

julia> @btime f($I, $A)
  7.500 μs (3 allocations: 64 bytes)
1.0 K

julia> @which Base.literal_pow(^, A, Val(3//2))
literal_pow(::typeof(^), x::Unitful.AbstractQuantity, ::Val{v}) where v in Unitful at /Users/user/.julia/packages/Unitful/MOEUx/src/quantities.jl:461

julia> versioninfo()
Julia Version 1.4.0
Commit b8e9a9ecc6 (2020-03-21 16:36 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin18.6.0)
  CPU: Intel(R) Core(TM) i7-6700HQ CPU @ 2.60GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-8.0.1 (ORCJIT, skylake)

[ggggggggg]

Here is a much simpler example that doesn't infer for some reason. This just returns a fully defined unit, but it allocates and @code_warntype says it returns type Any.

julia> using Unitful, BenchmarkTools

julia> f() = u"A/K^(3//2)"
f (generic function with 1 method)

julia> @btime f()

  4.104 μs (12 allocations: 768 bytes)
A K⁻³ᐟ²

julia> @code_warntype f()
Variables
  #self#::Core.Compiler.Const(f, false)

Body::Any
1 ─ %1 = (3 // 2)::Rational{Int64}
│   %2 = (K ^ %1)::Any
│   %3 = (A / %2)::Any
└──      return %3

[ggggggggg]

And even simpler, a function returning the units A^2 is inferred, while a function returning A^(3/2) is not. It appear that A is not identified as a compiler const in the second function, though that may be reading too much into it.

julia> f() = u"A^2"
f (generic function with 1 method)

julia> @code_warntype f()
Variables
  #self#::Core.Compiler.Const(f, false)

Body::Unitful.FreeUnits{(A²,),𝐈²,nothing}
1 ─ %1 = Core.apply_type(Base.Val, 2)::Core.Compiler.Const(Val{2}, false)
│   %2 = (%1)()::Core.Compiler.Const(Val{2}(), false)
│   %3 = Base.literal_pow(Main.:^, A, %2)::Core.Compiler.Const(A², false)
└──      return %3

julia> g() = u"A^(3/2)"
g (generic function with 1 method)

julia> @code_warntype g()
Variables
  #self#::Core.Compiler.Const(g, false)

Body::Any
1 ─ %1 = (3 / 2)::Core.Compiler.Const(1.5, false)
│   %2 = (A ^ %1)::Any
└──      return %2

[giordano]

For the record, the involved methods are https://github.com/PainterQubits/Unitful.jl/blob/0d7decd6ca8a36c880b9977f7a7e88e4442dea59/src/units.jl#L115-L116 u"A"^(2) hits the first method, u"A"^(3/2) the second one.

[ggggggggg]

I would expect since the exponent is known at compile time that it should hit Base.literal_pow here https://github.com/PainterQubits/Unitful.jl/blob/0d7decd6ca8a36c880b9977f7a7e88e4442dea59/src/units.jl#L132 And then the methods you point to. Then I stepper through calls with Debugger.jl and found it goes to: https://github.com/PainterQubits/Unitful.jl/blob/0d7decd6ca8a36c880b9977f7a7e88e4442dea59/src/units.jl#L104-L105

[ggggggggg]

This reproduced on 1.6

julia> @code_warntype g()
Variables
  #self#::Core.Const(g)

Body::Any
1 ─ %1 = (3 / 2)::Core.Const(1.5)
│   %2 = (A ^ %1)::Any
└──      return %2

julia> versioninfo()
Julia Version 1.6.1
Commit 6aaedecc44 (2021-04-23 05:59 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: Intel(R) Xeon(R) W-2123 CPU @ 3.60GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-11.0.1 (ORCJIT, skylake-avx512)

[ggggggggg]

Also here is a workaround, it avoid using a float or rational power.

julia> g()= sqrt(u"A"^3)
g (generic function with 1 method)

julia> @code_warntype g()
Variables
  #self#::Core.Const(g)

Body::Unitful.FreeUnits{(A^3/2,), 𝐈^3/2, nothing}
1 ─ %1 = Core.apply_type(Base.Val, 3)::Core.Const(Val{3})
│   %2 = (%1)()::Core.Const(Val{3}())
│   %3 = Base.literal_pow(Main.:^, A, %2)::Core.Const(A^3)
│   %4 = Main.sqrt(%3)::Core.Const(A^3/2)
└──      return %4

[sostock]

The first two examples (f(I, A) = Base.literal_pow(^, I/A, Val(2//3)) and f() = u"A/K^(3//2)") are now inferred correctly (in Julia 1.8).

The case g() = u"A^(3/2)" (using a Float64 instead of a Rational) still isn’t inferred.