Add specialized numba codegen for `as_while=False` scans

[jessegrabowski]

Description

tl;dr: The numba codegen for scan currently always emits a while loop. This can lead to up to 4x slowdown on long sequences vs a native for-loop. Here's graph:

numba_funcify_Scan unconditionally creates code of the form:

i = 0
cond = np.array(False)
while i < n_steps and not cond.item():
   inner_fn()
   i += 1

while scans are extremely rare relative to for scans. In the for case, we never update the cond variable, so the code is logically equivalent to a for loop. It turns out, however, that LLVM can't prove that. The LLM tells me that this is because np.array(False) creates a heap-resident numpy array via numba's memory allocator. LLVM cannot prove that: 1) the cond array doesn't alias the buffers being written to inside the loop, nor 2) the .item() call reads from a location that is loop-invariant.

LLVM conservatively assumes that the memory could possibly be modified by the stores inside the loop body, so it executes the following IR:

B140.endif:                                       ; the loop header
  %lsr.iv10 = phi ptr [ %scevgep, %B140.endif.preheader ], [ %scevgep11, %B172 ]
  %lsr.iv = phi i64 [ %0, %B140.endif.preheader ], [ %lsr.iv.next, %B172 ]

  ;; ↓ MUST load cond from heap memory every iteration (can't prove it's invariant)
  %.193 = load i8, ptr %.6.i.i, align 1
  %.196 = icmp eq i8 %.193, 0                     ; cond == False?
  br i1 %.196, label %B172, label %B174            ; if cond, exit loop

B172:                                             ; the loop body
  %scevgep12 = getelementptr i8, ptr %lsr.iv10, i64 -8
  %.131 = load double, ptr %scevgep12, align 8    ; ← MUST load buf[i] from memory
  %.132 = fadd double %.131, 1.000000e+00         ; buf[i+1] = buf[i] + 1.0
  store double %.132, ptr %lsr.iv10, align 8       ; write to buf
  %lsr.iv.next = add i64 %lsr.iv, -1              ; decrement
  %scevgep11 = getelementptr i8, ptr %lsr.iv10, i64 8  ; advance
  %exitcond.not = icmp eq i64 %lsr.iv.next, 0     ; counter == 0?
  br i1 %exitcond.not, label %B174, label %B140.endif  ; loop back to cond check

Ideally, we should be getting this:

B30:                                              ; the loop body
  %lsr.iv1 = phi i64 [ %arg.n_steps, %B30.preheader ], [ %lsr.iv.next, %B30 ]
  %lsr.iv = phi ptr [ %invariant.gep, %B30.preheader ], [ %scevgep, %B30 ]
  %store_forwarded = phi double [ %load_initial, %B30.preheader ], [ %.168, %B30 ]
  ;; ↑ LLVM forwards the stored value directly via a phi node — no memory load needed!

  %.168 = fadd double %store_forwarded, 1.000000e+00    ; buf[i+1] = buf[i] + 1.0
  store double %.168, ptr %lsr.iv, align 8               ; write to buf
  %scevgep = getelementptr i8, ptr %lsr.iv, i64 8        ; advance pointer
  %lsr.iv.next = add i64 %lsr.iv1, -1                    ; decrement counter
  %exitcond.not = icmp eq i64 %lsr.iv.next, 0            ; counter == 0?
  br i1 %exitcond.not, label %B62, label %B30             ; loop or exit

The fix is pretty straight-forward. We currently always generate loops that look like this:

def scan(n_steps, outer_in_1):

    outer_in_1_len = outer_in_1.shape[0]
    outer_in_1_sitsot_storage = outer_in_1

    outer_in_1_sitsot_storage_temp_scalar_0 = np.empty((), dtype=np.float64)

    i = 0
    cond = np.array(False)                                         # ← heap allocation
    while i < n_steps and not cond.item():                         # ← checked every iter
        outer_in_1_sitsot_storage_temp_scalar_0[()] = outer_in_1_sitsot_storage[(i) % outer_in_1_len]

        (inner_out_0,) = scan_inner_func(
            outer_in_1_sitsot_storage_temp_scalar_0
        )

        outer_in_1_sitsot_storage[(i + 1) % outer_in_1_len] = inner_out_0
        i += 1                                                     # ← cond never updated

    ...
    return outer_in_1_sitsot_storage

Instead, we can check the as_while flag at compile time and, if False, specialize to something like this:

def scan(n_steps, outer_in_1):

    outer_in_1_len = outer_in_1.shape[0]
    outer_in_1_sitsot_storage = outer_in_1

    outer_in_1_sitsot_storage_temp_scalar_0 = np.empty((), dtype=np.float64)

    for i in range(n_steps):                                       # ← simple counted loop
        outer_in_1_sitsot_storage_temp_scalar_0[()] = outer_in_1_sitsot_storage[(i) % outer_in_1_len]

        (inner_out_0,) = scan_inner_func(
            outer_in_1_sitsot_storage_temp_scalar_0
        )

        outer_in_1_sitsot_storage[(i + 1) % outer_in_1_len] = inner_out_0

    ...
    return outer_in_1_sitsot_storage

Basically something that looks like this (untested LLM code):

if op.info.as_while:
    loop_src = f"""
    i = 0
    while i < n_steps:
{indent(inner_scalar_in_args_to_temp_storage, " " * 8)}

        {inner_outputs} = scan_inner_func(
            {inner_in_args}
        )
{indent(inner_out_post_processing_block, " " * 8)}
{indent(inner_out_to_outer_out_stmts, " " * 8)}
        if cond:
            i += 1
            break
        i += 1
"""
else:
    loop_src = f"""
    for i in range(n_steps):
{indent(inner_scalar_in_args_to_temp_storage, " " * 8)}

        {inner_outputs} = scan_inner_func(
            {inner_in_args}
        )
{indent(inner_out_post_processing_block, " " * 8)}
{indent(inner_out_to_outer_out_stmts, " " * 8)}
"""

[jessegrabowski]

Admittedly we're only talking about ~4 extra instructions (~2 ns per iteration), so this difference will converge to 0 as the inner function dominates the compute time. But it still seems like a free win for simple loops that are common in models: GRW, AR(1), small gradient accumulators.

[ricardoV94]

those graphs disagree with each other? constant overhead of 2ns but the "impact" ratio grows for large sequences?

Anyway yes it's easy to tweak. We should probably make the cond be a scalar. We should make many things in scan be a scalar when they can (like vector sit-sot that indexes to scalar), it was one of the largest impacts in here #1632

[ricardoV94]

Also I would like a pytensor MWE (plots without MWE are lame btw). I suspect you won't find one because our overhead of working with arrays dominates

[ricardoV94]

So indeed I can't find a MWE with PyTensor.

• If we optimize away the sit-sot to untraced-sit-sot Convert unit buffer sit sot to untraced sit sot #2005, the overhead isn't seen
  • Bot thinks the reason LLVM didn't optimize in your example was due to concern of alias between cond and the arrays going in and out of the loop function
• If we keep a sit-sot, the cost of copying to the temporary 0d scalar array completely dominates

But to be honest the bot can't yet reproduce the way we build scans / elemwise reliably to reproduce anything meaningful.

Imo this is a non-concern for us right now, we are too far from trivial loop overhead... otoh no reason not to cleanup. I think this would be sufficient:

  cond = False
  while i < n_steps and not cond:
      ...
      if as_while:
          *args, cond_arr = inner_fn(*args, *constants)
          cond = cond_arr.item()  # extract from inner output to scalar

And ideally we start defining the condition to be scalar inside, so we don't even need the cond = cond_arr.item(), have to see if cython scan has any limitation against this. We also need to start optimizing scalar_from_tensor(Elemwise(core_op, x)) as scalar_op(scalar_from_tensor(x)) in case the condition is an Elemwise.

That's where we lose the most performance vs these trivial numba loops people come up with.

[jessegrabowski]

Yeah I got a bit too excited late at night staring at assembly with the stupid robot : )

This issue is dumb and we can close it, but I do think there room to think about places where we can specialize the actual numba codegen. I think in-lining the inner scan function is a place where we would see actual gains, not late-night hallucinations.

[ricardoV94]

Much more than you ever wanted to know: https://github.com/ricardoV94/bench-scan-numba

TLDR, my opinion on actionable items:

• We need to work on untraced_sit_sot with dynamic n_steps, because the gains are big. We see that here and we saw that with the LKJCorr. No reason to penalize dynamic n_steps, we always (should) know when the user chose x[-1]
  • I already started, eagerly destroying local_subtensor_merge can complicate graphs  #112 in the process, which makes some scan graphs compile 8s->ms)
• Allow scans to have scalar graphs inside. Maybe as a late rewrite because all our rewrites tackle elemwise not scalar graphs.
• Allow numba inner function to hold on to temporary allocated variables in the inner graph, those allocs add a lot in a hot loop. Not sure how we can get there. Scalars are mostly a way to avoid those allocs, that we CAN achieve, but are limited to such graphs.
• Pass the output storage for numba to write to directly, instead of copying from an inner allocated buffer. This is half of the complexity of cython Scan, the worst part being that cython Scan can't trust the output was actually written to, it's just a suggestion not a demand. My useless work on Replace join by direct writes in ancestor producers #2014 found the same issue. We need a strict (out=) API so we can trust an Op will store into a pre-defined buffer.
  • We could add an inplace write at the end, but what we really want is for the final Op (if there's one), to compute it's result directly on the pre-allocated buffer.