Type issue with captured variables -> let workaround failed

Hi!

I’m aware of #15276 but today I encountered an issue in the proposed workaround.

using FFTW


function f(x)
    p = plan_fft(x)
    x_ft = p * x 

   # given workaround in performance tips
    g = let p = p 
            x_ft = x_ft
        a -> (p*a) * x_ft
    end 

    return g
end


function f2(x)
    p = plan_fft(x)
    x_ft = p * x 

    # difference is a local name which is different to the original x_ft
    g = let p = p 
            x_ft2 = x_ft
        a -> (p*a) * x_ft2
    end 

    return g
end

The first function struggles with boxing, the second not:

julia> g = f(x)
#9 (generic function with 1 method)

julia> @code_warntype g(x)
MethodInstance for (::var"#9#10"{FFTW.cFFTWPlan{ComplexF64, -1, false, 2, UnitRange{Int64}}})(::Matrix{Float64})
  from (::var"#9#10")(a) in Main at /home/fxw/.julia/dev/FourierTools.jl/lel.jl:11
Arguments
  #self#::var"#9#10"{FFTW.cFFTWPlan{ComplexF64, -1, false, 2, UnitRange{Int64}}}
  a::Matrix{Float64}
Locals
  x_ft::Union{}
Body::Any
1 ─ %1  = Core.getfield(#self#, Symbol("#66#p"))::FFTW.cFFTWPlan{ComplexF64, -1, false, 2, UnitRange{Int64}}
β”‚   %2  = (%1 * a)::Matrix{ComplexF64}
β”‚   %3  = Core.getfield(#self#, :x_ft)::Core.Box
β”‚   %4  = Core.isdefined(%3, :contents)::Bool
└──       goto #3 if not %4
2 ─       goto #4
3 ─       Core.NewvarNode(:(x_ft))
└──       x_ft
4 β”„ %9  = Core.getfield(%3, :contents)::Any
β”‚   %10 = (%2 * %9)::Any
└──       return %10


julia> g2 = f2(x)
#15 (generic function with 1 method)

julia> @code_warntype g2(x)
MethodInstance for (::var"#15#16"{Matrix{ComplexF64}, FFTW.cFFTWPlan{ComplexF64, -1, false, 2, UnitRange{Int64}}})(::Matrix{Float64})
  from (::var"#15#16")(a) in Main at /home/fxw/.julia/dev/FourierTools.jl/lel.jl:25
Arguments
  #self#::var"#15#16"{Matrix{ComplexF64}, FFTW.cFFTWPlan{ComplexF64, -1, false, 2, UnitRange{Int64}}}
  a::Matrix{Float64}
Body::Matrix{ComplexF64}
1 ─ %1 = Core.getfield(#self#, Symbol("#76#p"))::FFTW.cFFTWPlan{ComplexF64, -1, false, 2, UnitRange{Int64}}
β”‚   %2 = (%1 * a)::Matrix{ComplexF64}
β”‚   %3 = Core.getfield(#self#, :x_ft2)::Matrix{ComplexF64}
β”‚   %4 = (%2 * %3)::Matrix{ComplexF64}
└──      return %4

What is the reason for this behavior? It is definitely related to the FFTW planning or at least to the FFTW.cFFTWPlan (which is an own type).

Best,

Felix

Does adding a comma at the end of let p = p help? Currently, x_ft = x_ft in f(x) is the first statement inside the let block, not part of its initial comma-separated assignment list.

I don’t understand what let does - or scoping in general - well enough to know why that makes a difference. But in my case, it makes both functions have the same @code_warntype output with no type issues.

3 Likes

You’re totally right!
I totally wasn’t aware of that, but it totally makes sense (need urgently fix a few dozens lines in my code :joy:)

1 Like

You can look at how Julia parses the blocks:


julia> using Base.Meta

julia> @dump let a=a
       b=b
       1
       end
Expr
  head: Symbol let
  args: Array{Any}((2,))
    1: Expr
      head: Symbol =
      args: Array{Any}((2,))
        1: Symbol a
        2: Symbol a
    2: Expr
      head: Symbol block
      args: Array{Any}((4,))
        1: LineNumberNode
          line: Int64 2
          file: Symbol REPL[2]
        2: Expr
          head: Symbol =
          args: Array{Any}((2,))
            1: Symbol b
            2: Symbol b
        3: LineNumberNode
          line: Int64 3
          file: Symbol REPL[2]
        4: Int64 1

julia> @dump let a=a, b=b
       1
       end
Expr
  head: Symbol let
  args: Array{Any}((2,))
    1: Expr
      head: Symbol block
      args: Array{Any}((2,))
        1: Expr
          head: Symbol =
          args: Array{Any}((2,))
            1: Symbol a
            2: Symbol a
        2: Expr
          head: Symbol =
          args: Array{Any}((2,))
            1: Symbol b
            2: Symbol b
    2: Expr
      head: Symbol block
      args: Array{Any}((2,))
        1: LineNumberNode
          line: Int64 2
          file: Symbol REPL[3]
        2: Int64 1

maybe then it’s easier to see what’s wrong. @dump helps when the syntax looks ambiguous.