v slow evaluation with particular combo of host_callback.id_tap + cond ?

[phc27x]

I've been playnig around with Jax/jit for implementing certain stream-processing patterns (typical in async timeseries) by using the lax while_loop to run an event loop, and trying to extract the results with the experimental id_tap features in host_callback. It was working really well, and I thought I groked how it all worked, and then suddenly it was 1000x slower after i changed something, and it then took me ages to find a repro.

So apologies this repro is not smaller, but it was a real pain to find which scenario had triggered the huge slow down.

My questions:

• how can I make sense of the x1000 slowdown here, just because i use a cond and elsewhere use id_print ? To be that slow, I feel it must be copying a massive array around constantly, or reallocating mem constantly, but the jit-compiler should have realised that everything can be updated in place.
• is there a good way to 'see' what its spending all its time on?

Grateful for any advice about this - I can file it as an 'issue' if appropriate, but am new around here so starting with a discussion.

import datetime as dt
import jax
from jax import jit
import jax.numpy as jnp
from jax.lax import while_loop, cond
import jax.experimental.host_callback as hcb # NB: experimental
partial = jax.tree_util.Partial
from collections import namedtuple

def a_funky_function( x ):
    return jnp.logical_not(x % 9)

def one_step( context, static ):   
    # read one message, and increment the iterator
    msg = jax.tree_map( lambda x : x[ context['in_iter'].begin ], context['in_array'] )
    context['in_iter'] = InIter( context['in_iter'].begin+1, context['in_iter'].end )
    
    keep_this_msg = a_funky_function( msg['value'] )         

    if not static.CONDITION_1:
        context['out'] = push( context['out'], msg, static )
    else:
        context['out'] = cond( keep_this_msg,
                      lambda q, val: push(q, val, static),
                      lambda q, val: q,
                      context['out'], msg)

    return context

def flush_return_none( buf, static ):            
    if not static.CONDITION_2:
        return None
    # hcb.id_tap also evidences it
    hcb.id_print( buf['value'][0] )
    hcb.id_print( buf['time'][0] )         

def push( out, val, static ):
    '''
    Push val into the out-buffer defined by out
    '''    
    it, buf = out
    new_it = OutIter( (it.end+1) % it.N, it.N)    
    new_buf = jax.tree_multimap( lambda q,m: q.at[it.end].set(m), buf, val)
    
    # if we exhausted and looped back to start after N, 'tap out' the ring buffer to host 
    cond( new_it.end == 0,
         lambda buf: flush_return_none( buf, static ),
         lambda buf: None,
         new_buf
        )
    
    return new_it, new_buf

def next_exists( context ):    
    in_iter = context['in_iter']
    in_array = context['in_array']
    
    next_event_time = cond( in_iter.begin != in_iter.end,
         lambda i_it, i_ar : jax.tree_map( lambda x : x[i_it.begin], i_ar )['time'],
         lambda i_it, i_ar : MAX_EVENT,
         in_iter, in_array)
    
    return next_event_time < MAX_EVENT


def main_loop_body( context, static ):                
    in_iter = context['in_iter']
    in_array = context['in_array']
    next_event = jax.tree_map( lambda x : x[in_iter.begin], in_array )['time']
            
    do_next_step = next_event < MAX_EVENT
    
    context = cond( do_next_step,
                    partial(one_step, static=static),
                    lambda x:x,
                    context)    
    return context

def main( context, static ):
    context = while_loop( lambda context: next_exists(context),
                         partial(main_loop_body, static=static),
                         context)
    return context

The operation is basically to iterate over an input array, and copy it to the output array, but conditional on the value of function, which is simplest version of repro i could make.
( The out-buffer is big enough in this case, that evaluation never even hits the branch flush_return_none - but the presence of the branch in the code is enough to make a big slowdown )

InIter = namedtuple('InIter',['begin', 'end'])
OutIter = namedtuple('OutIter',['end', 'N'])
# in-data
in_sz = 2**15-1
in_iter = InIter( jnp.array(0), jnp.array(in_sz) )
times = jnp.arange(in_sz)
values = jnp.arange(in_sz)
in_array = {'time': times, 'value': values}

# out-data
out_sz = 2**16
out_iter = OutIter( jnp.array(0), out_sz )
out_buffer = {'time':jnp.zeros( out_sz ),
              'value':jnp.zeros( out_sz )
             }
MAX_EVENT= 2**17

context = {'in_iter': in_iter, 'in_array': in_array, 'out': ( out_iter, out_buffer ) }

Now run it with different static-settings

Conditions = namedtuple('Conditions',['CONDITION_1', 'CONDITION_2'])
for static in [ Conditions(False,False), Conditions(False, True), Conditions(True, False), Conditions(True,True) ] :    
    jmain = jit( main, static_argnames=('static',) )
    jmain._clear_cache()
    jmain(context, static=static);
    tic = dt.datetime.now()
    jmain(context, static=static)
    toc = dt.datetime.now()
    print( (toc-tic) / dt.timedelta(milliseconds=1), '(ms)', static)

Prints the following - showing significant slowdown if both conditions are true.

1.062 (ms) Conditions(CONDITION_1=False, CONDITION_2=False)
1.236 (ms) Conditions(CONDITION_1=False, CONDITION_2=True)
1.416 (ms) Conditions(CONDITION_1=True, CONDITION_2=False)
1065.826 (ms) Conditions(CONDITION_1=True, CONDITION_2=True)

(I'm running on cpu only)