Surprising slowdown from using array comprehension

Checking @code_warntype, the root difference seems to be ncols getting boxed instead of inferred as Int, while the paired nrows is inferred as Int. Comprehensions are built on closures, which risk that sort of thing, but normally this wouldn’t happen if you never reassign the captured variables. As you’ve pointed out, the boxing doesn’t happen if you comment out false && nrows * ncols. More tweaks:

  1. switching the assignment order to ncols, nrows = 7, 7 or separating them in 2 lines doesn’t change which gets boxed
  2. switching the order to false && ncols * nrows doesn’t change which gets boxed
  3. switching the comprehension loop levels for col = 1:ncols for row = nrows:-1:1 (which also changes the problem so I didn’t bother running it) does switch to boxing nrows.
  4. switching the comprehension to 2D for row = nrows:-1:1, col = 1:ncols infers both captured variables but is such a drastic change it makes the whole thing fail

So it’s just a hole in the lowerer’s ability to identify when captured variables in nested comprehensions don’t need boxing; I have no idea why the false && nrows * ncols is necessary because it’s not a reassignment, but that’s probably a factor why this doesn’t happen often. ncols::Int (and that alone) indeed restores good type inference in the rest of the algorithm, even though it doesn’t remove the box. You should open an issue.

EDIT: Okay I did manage to make a MWE, and it suggests @label is the more critical factor. Commenting the label (and its paired @goto) out in the original example also undoes the boxing (but ruins the algorithm, don’t run it). Just writing the captured variable after the label is enough for boxing. Comprehension for-headers also seem to have it easier, as the body (or a more typical closure ()->(x,y)) boxes all captured variables.

# Reduced example doesn't do anything, only for @code_warntype
function solvemwe()
    nrows, ncols, x, y = 7, 7, 1, 1
              # body boxes x,y     # for boxes latter variable, swap to check
    rowcol = [(x, y, row, col + 1) for row = nrows:-1:1 for col = 1:ncols]
    @label next_iter # comment out for good inference
    nrows, ncols, x, y # comment out boxed variables for good inference
end
5 Likes