How to obtain partitioned HLO from pjit'ed function?

[leiteg]

Hey, I have been experimenting with the pjit transformation and was wondering if it is possible to obtain the HLO module for this function after the partitioning pass in XLA has ran. From what I gather from the GSPMD paper, the communication primitives between devices should be automatically inserted by the compiler.

Consider the following scenario: a sharded matrix A is multiplied by a replicated vector x and the output is another vector that should be also replicated across devices. This is implemented in the short example below.

#!/usr/bin/env python3

import os
import jax
import jax.numpy as jnp
import numpy as np

from jax.experimental import PartitionSpec
from jax.experimental.maps import Mesh
from jax.experimental.pjit import pjit

# Expose 4 devices to JAX/XLA
# ------------------------------------------------
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'
jax.config.update('jax_platform_name', 'cpu')

# Create 2x2 device mesh
# ------------------------------------------------
devices = jax.devices()
device_array = np.asarray(devices).reshape(2, 2)
device_mesh = Mesh(device_array, ("x", "y"))

# Create matrix A and vector x
# ------------------------------------------------
A = jnp.eye(16)
x = jnp.arange(16).reshape((-1, 1))

# Transform `dot` operator using `pjit`
# ------------------------------------------------
spec = PartitionSpec("x", "y")

f = pjit(jnp.dot,
         in_axis_resources=(spec, None),
         out_axis_resources=None)

# Obtain HLO representation
# ------------------------------------------------
with device_mesh:
    ir = f.lower(A, x).compiler_ir("hlo")

print(ir.as_hlo_text())

Executing this code produces the following output:

HloModule pjit_dot.7

ENTRY main.7 {
  Arg_0.1 = f32[16,16]{1,0} parameter(0), sharding={devices=[2,2]0,1,2,3}
  Arg_1.2 = s32[16,1]{1,0} parameter(1), sharding={replicated}
  convert.3 = f32[16,1]{1,0} convert(Arg_1.2)
  dot.4 = f32[16,1]{1,0} dot(Arg_0.1, convert.3), lhs_contracting_dims={1}, rhs_contracting_dims={0}
  tuple.5 = (f32[16,1]{1,0}) tuple(dot.4)
  ROOT get-tuple-element.6 = f32[16,1]{1,0} get-tuple-element(tuple.5), index=0, sharding={replicated}
}

The HLO module above includes the correct sharding annotations for all the tensors but there are no communication primitives whatsover. I suspect that the GSPMD Partitioner has not run at this point yet. Is there a way to obtain the HLO module after this compiler pass has executed? Thanks in advance!

[leiteg]

After investigating JAX source code a little bit further I found how to obtain the compiled (post-partitioning) HLO.

with device_mesh:
    modules = f.lower(A, x).compile().compiler_ir()
    for hlo in modules:
        print(hlo.to_string())

The resulting HLO looks like the listing below. Notice the all-reduce and fusion operations in the entry function.

HloModule pjit_dot.7

%add (x: f32[], y: f32[]) -> f32[] {
  %x = f32[] parameter(0)
  %y = f32[] parameter(1)
  ROOT %add = f32[] add(f32[] %x, f32[] %y)
}

%add.1 (x.1: f32[], y.1: f32[]) -> f32[] {
  %x.1 = f32[] parameter(0)
  %y.1 = f32[] parameter(1)
  ROOT %add.2 = f32[] add(f32[] %x.1, f32[] %y.1)
}

%fused_computation (param_0.1: f32[8,1], param_1.1: u32[], param_2.3: u32[]) -> f32[16,1] {
  %constant.18 = f32[] constant(0)
  %broadcast.2 = f32[16,1]{1,0} broadcast(f32[] %constant.18), dimensions={}
  %param_0.1 = f32[8,1]{1,0} parameter(0)
  %constant.16 = u32[4]{0} constant({0, 0, 1, 1})
  %param_2.3 = u32[] parameter(2)
  %constant.15 = u32[] constant(4)
  %multiply.5 = u32[] multiply(u32[] %param_2.3, u32[] %constant.15)
  %param_1.1 = u32[] parameter(1)
  %add.3 = u32[] add(u32[] %multiply.5, u32[] %param_1.1)
  %dynamic-slice.5 = u32[1]{0} dynamic-slice(u32[4]{0} %constant.16, u32[] %add.3), dynamic_slice_sizes={1}
  %reshape.8 = u32[] reshape(u32[1]{0} %dynamic-slice.5)
  %constant.14 = u32[] constant(8)
  %multiply.4 = u32[] multiply(u32[] %reshape.8, u32[] %constant.14)
  %constant.17 = u32[] constant(0)
  ROOT %dynamic-update-slice.1 = f32[16,1]{1,0} dynamic-update-slice(f32[16,1]{1,0} %broadcast.2, f32[8,1]{1,0} %param_0.1, u32[] %multiply.4, u32[] %constant.17)
}

%fused_computation.1 (param_0.4: s32[16,1], param_1.6: u32[]) -> f32[8] {
  %param_0.4 = s32[16,1]{1,0} parameter(0)
  %constant.23 = u32[] constant(0)
  %constant.22 = u32[4]{0} constant({0, 1, 0, 1}), metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %param_1.6 = u32[] parameter(1)
  %dynamic-slice.7 = u32[1]{0} dynamic-slice(u32[4]{0} %constant.22, u32[] %param_1.6), dynamic_slice_sizes={1}, metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %reshape.10 = u32[] reshape(u32[1]{0} %dynamic-slice.7), metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %constant.21 = u32[] constant(1), metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %clamp.1 = u32[] clamp(u32[] %constant.23, u32[] %reshape.10, u32[] %constant.21), metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %convert.4 = s32[] convert(u32[] %clamp.1), metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %constant.20 = s32[] constant(8), metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %multiply.6 = s32[] multiply(s32[] %convert.4, s32[] %constant.20), metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %constant.19 = s32[] constant(0), metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %dynamic-slice.6 = s32[8,1]{1,0} dynamic-slice(s32[16,1]{1,0} %param_0.4, s32[] %multiply.6, s32[] %constant.19), dynamic_slice_sizes={8,1}, metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %convert.2 = f32[8,1]{1,0} convert(s32[8,1]{1,0} %dynamic-slice.6), metadata={op_name="pjit(dot)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  ROOT %reshape.9 = f32[8]{0} reshape(f32[8,1]{1,0} %convert.2)
}

ENTRY %main.7_spmd (param: f32[8,8], param.1: s32[16,1]) -> f32[16,1] {
  %param = f32[8,8]{1,0} parameter(0), sharding={devices=[2,2]0,1,2,3}
  %param.1 = s32[16,1]{1,0} parameter(1), sharding={replicated}
  %partition-id = u32[] partition-id()
  %fusion.1 = f32[8]{0} fusion(s32[16,1]{1,0} %param.1, u32[] %partition-id), kind=kLoop, calls=%fused_computation.1
  %dot.1 = f32[8]{0} dot(f32[8,8]{1,0} %param, f32[8]{0} %fusion.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
  %bitcast = f32[8,1]{1,0} bitcast(f32[8]{0} %dot.1), metadata={op_name="pjit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %all-reduce = f32[8,1]{1,0} all-reduce(f32[8,1]{1,0} %bitcast), channel_id=1, replica_groups={{0,1},{2,3}}, use_global_device_ids=true, to_apply=%add, metadata={op_name="pjit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="/scratch/playground/jax/./pmap.py" source_line=39}
  %replica-id = u32[] replica-id()
  %fusion = f32[16,1]{1,0} fusion(f32[8,1]{1,0} %all-reduce, u32[] %partition-id, u32[] %replica-id), kind=kLoop, calls=%fused_computation
  ROOT %all-reduce.1 = f32[16,1]{1,0} all-reduce(f32[16,1]{1,0} %fusion), channel_id=2, replica_groups={{0,2},{1,3}}, use_global_device_ids=true, to_apply=%add.1
}