Index inside a comprehension is boxed

Hi

I got a bug report for ResumableFunctions and after some analysis the problem appeared to be the boxing of an integer used as an index in a comprehension.

Minimal example showing the boxing of the variable i

function test0()
  a = [[1, 2],[3, 4]]
  for i in 1:2 
    u = [v[i] for v in a]
  end
end
code_warntype(test0,())
Variables:
  #self# <optimized out>
  #69::##69#70
  i::Core.Box
  u <optimized out>
  #temp#::Int64
  a::Array{Array{Int64,1},1}

Body:
  begin
      a::Array{Array{Int64,1},1} = $(Expr(:invoke, MethodInstance for vect(::Array{Int64,1}, ::Vararg{Array{Int64,1},N} where N), :(Base.vect), :($(Expr(:invoke, MethodInstance for vect(::Int64, ::Vararg{Int64,N} where N), :(Base.vect), 1, 2))), :($(Expr(:invoke, MethodInstance for vect(::Int64, ::Vararg{Int64,N} where N), :(Base.vect), 3, 4))))) # line 3:
      SSAValue(5) = (Base.select_value)((Base.sle_int)(1, 2)::Bool, 2, (Base.sub_int)(1, 1)::Int64)::Int64
      #temp#::Int64 = 1
      5:
      unless (Base.not_int)((#temp#::Int64 === (Base.add_int)(SSAValue(5), 1)::Int64)::Bool)::Bool goto 19
      i::Core.Box = $(Expr(:new, :(Core.Box)))
      SSAValue(6) = #temp#::Int64
      SSAValue(7) = (Base.add_int)(#temp#::Int64, 1)::Int64
      SSAValue(2) = SSAValue(6)
      (Core.setfield!)(i::Core.Box, :contents, SSAValue(2))::Int64
      #temp#::Int64 = SSAValue(7) # line 3:
      #69::##69#70 = $(Expr(:new, :(Main.##69#70), :(i)))
      SSAValue(4) = $(Expr(:new, Base.Generator{Array{Array{Int64,1},1},##69#70}, :(#69), :(a)))
      $(Expr(:invoke, MethodInstance for collect(::Base.Generator{Array{Array{Int64,1},1},##69#70}), :(Base.collect), SSAValue(4)))
      17:
      goto 5
      19:
      return
  end::Void

Is this normal? I don’t see why the boxing is needed.
If no for loop is used, there is no boxing and if i is not used as an index in a comprehension, there is also no boxing…

This is a known (very annoying) issue: performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub

3 Likes

The let block workaround works here too:

function test0()
    a = [[1, 2],[3, 4]]
    for i in 1:2
        u = let i = i
            [v[i] for v in a]
        end
    end
end

Thanks!
The package ResumableFunctions does not support let blocks yet. So I have to implement the use of let blocks or find another workaround;)

You could try using https://github.com/c42f/FastClosures.jl as well.

I implemented let blocks in ResumableFunctions.

@resumable function test_let()
  for u in [[(1,2),(3,4)], [(5,6),(7,8)]]
    for i in 1:2
      let i=i
        val = [a[i] for a in u]
      end
      @yield val
    end
  end
end
collect(test_let())

[[1,3],[2,4],[5,7],[6,8]]

The variables in a let block are not stored between function calls. This is a nice way of improving performance;)
Thanks for the hints!

Ben

1 Like