How to Verify In-Place Updates by the Compiler in JAX

[bheijden]

Hello,

I am trying to understand how to verify if in-place updates are performed by the JAX compiler.

Some background: I've noticed that my program slows down significantly (by magnitudes) when I double the size of certain arrays used as buffers. This leads me to believe that in-place updates might not be happening when writing to these buffers. If the updates were truly in-place, the buffer size shouldn't impact the performance as much.

I am considering rewriting my code to ensure in-place updates, but I need guidance on how to check if an array is being updated in-place or if it's being copied. Are there specific commands or indicators in the HLO/jaxpr representations that I should look for?

Any pointers or advice would be greatly appreciated.

Thank you!

[jakevdp]

There's no direct way to check for this, but you can get an idea of what the compiler is doing with your code by printing the optimized HLO. For example:

import jax

def f(x):
  x = x.at[0].set(1)
  return x.sum()

x = jax.numpy.arange(10)

print(jax.jit(f).lower(x).compile().as_text())

HloModule jit_f, entry_computation_layout={(s32[10]{0})->s32[]}, allow_spmd_sharding_propagation_to_parameters={true}, allow_spmd_sharding_propagation_to_output={true}

%region_1.9 (Arg_0.10: s32[], Arg_1.11: s32[]) -> s32[] {
  %Arg_0.10 = s32[] parameter(0)
  %Arg_1.11 = s32[] parameter(1)
  ROOT %add.12 = s32[] add(s32[] %Arg_0.10, s32[] %Arg_1.11), metadata={op_name="jit(f)/jit(main)/reduce_sum[axes=(0,)]" source_file="<ipython-input-3-178b9b2715a2>" source_line=5}
}

%fused_computation (param_0.2: s32[10]) -> s32[] {
  %param_0.2 = s32[10]{0} parameter(0)
  %constant.28..sunk.3 = s32[1]{0} constant({1})
  %constant.34 = s32[] constant(0)
  %dynamic-update-slice.4 = s32[10]{0} dynamic-update-slice(s32[10]{0} %param_0.2, s32[1]{0} %constant.28..sunk.3, s32[] %constant.34), metadata={op_name="jit(f)/jit(main)/scatter[update_jaxpr=None update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(), inserted_window_dims=(0,), scatter_dims_to_operand_dims=(0,)) indices_are_sorted=True unique_indices=True mode=GatherScatterMode.FILL_OR_DROP]" source_file="<ipython-input-3-178b9b2715a2>" source_line=4}
  ROOT %reduce.1 = s32[] reduce(s32[10]{0} %dynamic-update-slice.4, s32[] %constant.34), dimensions={0}, to_apply=%region_1.9, metadata={op_name="jit(f)/jit(main)/reduce_sum[axes=(0,)]" source_file="<ipython-input-3-178b9b2715a2>" source_line=5}
}

ENTRY %main.14 (Arg_0.1: s32[10]) -> s32[] {
  %Arg_0.1 = s32[10]{0} parameter(0)
  ROOT %fusion = s32[] fusion(s32[10]{0} %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(f)/jit(main)/reduce_sum[axes=(0,)]" source_file="<ipython-input-3-178b9b2715a2>" source_line=5}
}

Here the dynamic_update_slice and the reduce happen within a single fused operation, which means that the compiler has created a single kernel for that sequence of ops. This is a good indication that the array update will be done in-place.

Hope that helps!

[bheijden]

Thank you for the detailed explanation! I have a few follow-up questions to make sure I understand this correctly:

1. Determination of In-Place Updates: If I understand correctly, it's only in the final stage, when the HLO is optimized, that it’s determined whether in-place updates or copying will occur. Is that accurate?

2. Fused Operations: Is everything within an outer pair of {} in the HLO a single fused operation, or are there exceptions to this?

3. Locating Relevant Operations: The function I am JIT compiling is quite large, so I imagine that locating the relevant operations in the optimized HLO will be challenging. Do you have any tips on how to efficiently find the parts of the HLO that correspond to specific operations in my code?

4. Using Metadata for Code Correspondence: I noticed that metadata keywords like op_name, source_file, and source_line seem to map back to the original code. Is there a way to "wrap" my JAX code with a NOOP or some other identifier to make it easier to sift through the HLO? Or is it difficult to achieve a 1-to-1 correspondence between the source code and the HLO due to optimizations?

5. Reading HLO Text: You seem to be quite skilled at reading HLO text. Could you point me to any documentation or resources that would help me learn how to read the optimized HLO? Additionally, are there any tips on how to connect specific lines in the HLO text back to the actual JAX code?

Thanks again for your help!

[jakevdp]

If I understand correctly, it's only in the final stage, when the HLO is optimized, that it’s determined whether in-place updates or copying will occur. Is that accurate?

Yes, that's my understanding. The compiler makes the final decision of how the program will be executed, including whether any buffers need to be copied. That makes it difficult to surface the information anywhere but in the compiled HLO

Is everything within an outer pair of {} in the HLO a single fused operation, or are there exceptions to this?

Everything within a block labeled %fused_computation will be fused.

Locating Relevant Operations:

It's difficult to do this for large programs, unfortunately.

Using Metadata for Code Correspondence:

I've never explored this, but it seems like it could be helpful

Could you point me to any documentation or resources that would help me learn how to read the optimized HLO?

I don't know of any such documentation, but once you look at a few HLO dumps you get used to it