Forcing constant folding or set max folding time

[epignatelli]

It happens sometimes that a constant is taking long time to compile and I receive an error of this kind:

2023-06-27 11:00:09.200673: E external/xla/xla/service/slow_operation_alarm.cc:65] Constant folding an instruction is taking > 1s:

  %fusion.130 = u8[10,51,32,32,3]{4,3,2,1,0} fusion(u8[32,32,3]{2,1,0} %constant.784, s32[10,10,5]{2,1,0} %constant.1967), kind=kLoop, calls=%fused_computation.130, metadata={op_name="jit(test)/jit(main)/concatenate[dimension=1]" source_file="/home/uceeepi/repos/navix/navix/graphics.py" source_line=111}

This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.

Is there any mechanism in place to force constant folding, or set a time threshold for it?

[davisyoshida]

I recently dealt with this, so I'm curious if there's a general solution as well.

In my case just modified my function to avoid there being large constants to fold. In my case the issue was checking whether all gradients were finite (for using jmp loss scaling). I was using optax.set_to_zero() to produce zero gradients for frozen parameters, but this meant that at compile time there were 7B zeros to check the finite-ness of, along with a few other ops. My solution was switching from making gradients with zeros_like() to just returning 0., since all the downstream ops supported broadcasting. Maybe there's a similar trick you can do in your case?