print optimized code

[mariogeiger]

Hi,

I would like to see what jit and xla is able to do as optimization.
For that I tried to look at jax.make_jaxpr and jax.xla_computation but it shows me clearly unoptimized code that executes dead code (code that computes unused quantities).

See an example:

def f(a, b, c):
    return a * b * c

print(jax.make_jaxpr(jax.grad(f))(1., 2., 3.))

prints

{ lambda  ; a b c.
  let d = mul a b
      _ = mul d c
      e = mul 1.0 c
      f = mul e b
  in (f,) }

_ is clearly unused, I have the impression that I'm looking at the unoptimized code and that the optimization still has to be performed. Is ti possible to print the code after the optimization has been preceded?

I guess that the optimized code should look like

{ lambda  ; a b c.
  let d = mul c b
  in (d,) }

[hawkinsp]

There are two ways to dump the optimized HLO. One is to set the following environment variable: XLA_FLAGS=--xla_dump_to=/tmp/somewhere.

Another is the following undocumented and internal API (no promises of stability):

In [34]: import jax

In [35]: def f(x): return jnp.sin(x)

In [36]: c = jax.xla_computation(f)(1.)

In [37]: backend = jax.lib.xla_bridge.get_backend()
    ...: e = backend.compile(c)
    ...: print(e.hlo_modules()[0].to_string())
HloModule xla_computation_f__1.5

ENTRY %xla_computation_f__1.5 (parameter.1: f64[]) -> (f64[]) {
  %parameter.1 = f64[] parameter(0)
  %sine.3 = f64[] sine(f64[] %parameter.1), metadata={op_type="sin" op_name="xla_computation(f)/sin" source_file="<ipython-input-35-bb7430a03aae>" source_line=1}
  ROOT %tuple.4 = (f64[]) tuple(f64[] %sine.3)
}

[soraros]

import jaxlib.xla_extension as xla_ext

# You could use the presets
option = xla_ext.HloPrintOptions.short_parsable()
print(module.to_string(option))

option = xla_ext.HloPrintOptions.canonical()
print(module.to_string(option))

option = xla_ext.HloPrintOptions.fingerprint()
print(module.to_string(option))

# Or set each option manually
option = xla_ext.HloPrintOptions()
option.print_metadata = False
option.include_layout_in_shapes = False
print(module.to_string(option))

[mariogeiger]

Perfect!

[awf]

For future visitors: the source of HloPrintOptions serves as documentation of available options:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/hlo_instruction.h

[jecampagne]

Link is broken.

[hawkinsp]

https://github.com/openxla/xla/blob/f6052ab38fb090c078e7d3522ea24a304b4be6e4/xla/hlo/ir/hlo_instruction.h#L72

[jakevdp]

In case people come across this question in searches, there is now an easier API for printing optimized computations, described at Ahead of time Lowering and compilation. For example:

def f(a, b, c):
    return a * b * c

print(jax.jit(jax.grad(f)).lower(1., 2., 3.).compile().as_text())

HloModule jit_f, entry_computation_layout={(f32[],f32[])->f32[]}, allow_spmd_sharding_propagation_to_output={true}

ENTRY %main.4 (Arg_0.1: f32[], Arg_1.2: f32[]) -> f32[] {
  %Arg_1.2 = f32[] parameter(1), sharding={replicated}
  %Arg_0.1 = f32[] parameter(0), sharding={replicated}
  ROOT %multiply.3 = f32[] multiply(f32[] %Arg_1.2, f32[] %Arg_0.1), metadata={op_name="jit(f)/jit(main)/mul" source_file="<ipython-input-1-c4f41180e36b>" source_line=4}
}

[joaospinto]

@jakevdp Is there a way of exporting portable artifacts (i.e. function -> jit -> export -> mlir_module_serialized) with similar mappings to the original user code? Alternatively, what's the easiest way of converting the HLO file produced by the method you described to StableHLO while preserving these mappings?

[jakevdp]

There's some information on exporting JAX code at https://jax.readthedocs.io/en/latest/export/index.html Does that help?

[joaospinto]

@jakevdp I've read that in the past, but the documentation about lowering and exporting is quite incomplete, so as far as I've seen it doesn't seem to cover my question above. The portable way of exporting this (which, IIUC, does not have these code mapping metadata) is by doing

jit_f = jax.jit(f)
exported = export.export(jit_f)(dummy_inputs)
mlir_bc = exported.mlir_module_serialized

On the other hand, the following does include code mapping metadata, but is not a portable export format, nor is it even StableHLO, IIUC.

jit_f = jax.jit(f)
lowered = jit_f.lower(dummy_inputs)
lowered.compile().as_text()

Ideally there would be a way of passing an argument to export.export to request that code mapping metadata be preserved.

[jakevdp]

Thanks – I'd suggest opening a new issue with this as a feature request. Comments on years-old discussions tend to get lost in the mix.

[joaospinto]

Sure - created here.