HLO Broadcast with decreasing broadcast_dimensions issues

[apivovarov]

I faced the following issues with hlo Broadcast Operation in case when broadcast_dimensions are in decreasing order (swapped)

Issue 1 - Jax does not support it - broadcast_dimensions must be strictly increasing:

import jax
import jax.numpy as jnp
x=jnp.arange(9*12).reshape((9,12))
y=jax.lax.broadcast_in_dim(x, (1,2,12,9), (3,2))

TypeError: broadcast_in_dim broadcast_dimensions must be strictly increasing; got broadcast_dimensions (3, 2)

Issue 2 - run_hlo_module validation failed. CPU and the Interpreter outputs are different

• Interpreter output is transposed
• CPU output is just reshaped

ENTRY %main.3 {
  %Arg_0.1 = s32[9,12]{1,0} parameter(0)
  ROOT %broadcast.2 = s32[1,2,12,9]{3,2,1,0} broadcast(%Arg_0.1), dimensions={3,2}
}

./bazel-bin/xla/tools/run_hlo_module \
--input_format=hlo \
--platform=CPU a.hlotxt \
--print_literals=true

1/1 runs failed.

Questions:

1. Is it a valid hlo when broadcast_dimensions are in decreasing order?
2. Should we add additional validation to xla/service/hlo_verifier.cc to check that broadcast->dimensions() are in increasing order?
3. If decreasing order is actually allowed what should be in the output - transposed or just reshaped input?

[apivovarov]

Got reply from Benjamin Chetioui @bchetioui

It shouldn't be! As you can see on the page describing the semantics of broadcasting, XLA expects that "The order of broadcast dimensions must be strictly increasing.", matching the semantics of JAX.

Kevin, Sandeep,
I noticed that in stablehlo broadcast_in_dim doc uses decreasing broadcast dimensions. Semantically the example is doing transpose 1x3 to 3x1 and then broadcast to 2x3x2.
Is it true that stablehlo spec allows to use decreasing broadcast dimensions? It is not allowed in hlo and jax (and in pytorch/TF broadcast_to ops too).
@sdasgup3 @GleasonK

[sdasgup3]

Thanks @apivovarov!
Let me start with the answers which are handy to me.

Issue 2 - run_hlo_module validation failed. CPU and the Interpreter outputs are different

I am curious to know the reason for the failure.
What I am getting on my end is a verifier error prohibiting run_hlo_module to proceed.

For example,
the following hlo should run fine:

HloModule module_main.3, entry_computation_layout={(s32[9,12]{1,0})->s32[1,2,12,9]{2,3,1,0}}
ENTRY %main.3 {
  %Arg_0.1 = s32[9,12]{1,0} parameter(0)
  ROOT %broadcast.2 = s32[1,2,12,9] broadcast(%Arg_0.1), dimensions={3,2}
}

with interpreter and CPU results close enough.

The reason the interpreter/CPU results are similar is because the implementation of this op is agnostic of the order of the broadcast dimensions. Being a data movement op, the implementation only needs to determine which result index a particular source index would map to. (https://github.com/tensorflow/tensorflow/blob/49959d2f7d5e497edf90d6cdebc9afa3247234b7/tensorflow/compiler/xla/service/elemental_ir_emitter.cc#L2562)

This is similar to to semantics proposed in StableHLO where it does not matter, from the POV of the execution semantics, which order (strictly increasing or not) the broadcast dimensions are specified.

formally, result[result_index] = operand[operand_index] where for all d in axes(operand):

operand_index[d] = 0 if dim(operand, d) = 1.
operand_index[d] = result_index[broadcast_dimensions[d]] otherwise.

3. If decreasing order is actually allowed what should be in the output - transposed or just reshaped input?

Currently, In stableHLO the broadcast dimensions can be specified in any order. The above program is equivalent to

HloModule module_main.4, entry_computation_layout={(s32[9,12]{1,0})->s32[1,2,12,9]{3,2,1,0}}

ENTRY %main.4 (Arg_0.1: s32[9,12]) -> s32[1,2,12,9] {
  %Arg_0.1 = s32[9,12]{1,0} parameter(0)
  %transpose.2 = s32[12,9] transpose(s32[9,12] %Arg_0.1), dimensions={1,0}, metadata={source_file="<stdin>" source_line=3}
  ROOT %broadcast.3 = s32[1,2,12,9] broadcast(s32[12,9] %transpose.2), dimensions={2,3}, metadata={source_file="<stdin>" source_line=3}
}

[apivovarov]

hlotxt

ENTRY %main.3 {
  %Arg_0.1 = s32[9,12]{1,0} parameter(0)
  ROOT %broadcast.2 = s32[1,2,12,9]{3,2,1,0} broadcast(%Arg_0.1), dimensions={3,2}
}

if run_hlo_module is built with --config dbg or -c fastbuild then it shows validation error:

INTERNAL: during context [hlo verifier (debug)]: Shape of the root instruction of entry computation (s32[1,2,12,9]{2,3,1,0}) should be compatible to one specified in module's entry computation layout (s32[1,2,12,9]{3,2,1,0})

if run_hlo_module is built for release -DNDEBUG (default build) (bazel build //xla/tools:run_hlo_module)
then run_hlo_module does not show the above validation message.

The results for CPU and interpreter are different:

• Interpreter output is transposed input
• CPU output is just reshaped input

Small Example with outputs:
run_hlo_module - release build

$ cat c.hlotxt                                                                      
ENTRY %main.3 {
  %Arg_0.1 = s32[2,3] parameter(0)
  ROOT %broadcast.2 = s32[1,3,2] broadcast(%Arg_0.1), dimensions={2,1}
}


$ ./run_hlo_module --input_format=hlo --platform=CPU  c.hlotxt --print_literals=true

** Argument 0 **
s32[2,3] {
  { -1865008400, 470211269, -689632771 },
  { 1115438162, -1009961146, 143542609 }
}
Running HLO module with runner Host...
** Result with test runner Host **
s32[1,3,2] {
{
  { -1865008400, 470211269 },
  { -689632771, 1115438162 },
  { -1009961146, 143542609 }
}
}
Running HLO module with runner Interpreter...
** Result with reference runner Interpreter **
s32[1,3,2] {
{
  { -1865008400, 1115438162 },
  { 470211269, -1009961146 },
  { -689632771, 143542609 }
}
}

1/1 runs failed.