`hlo_call` and `@code_hlo` use different array shapes

[dubosipsl]

Generating HLO code then compiling fails due to array shape mismatch.
I have a similar issue for HLO code generated by JAX.
I think this was working in April or so but cannot find the precise version of Reactant I was using.

version: julia 1.10.9, Reactant 0.2.143 (Google Colab)

using Reactant: Ops, to_rarray, @compile
using Random

a = to_rarray(randn(Float32, 2,3))
b = copy(a)

@compile(
          Ops.hlo_call(
              """
              module {
                func.func @main(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>) -> tensor<2x3xf32> {
                  %0 = stablehlo.add %arg0, %arg1 : tensor<2x3xf32>
                  return %0 : tensor<2x3xf32>
                }
              }
              """, a, b)
          ) # works

f(a,b) = a.*b
@compile f(a,b)  # works

code = string(@code_hlo f(a,b))
println(code)

f2(a,b) = Ops.hlo_call(code, a, b)
@compile f2(a,b)  # fails

[mofeing]

IMO this is expected behavior as of the way is currently implemented.
the reason why is that Julia and MLIR use different array orderings.

Julia is columnn-major (like Fortran) so the 2x3 matrix is stored in memory like this

[a_{11}, a_{21}, a_{12}, a_{22}, a_{13}, a_{23}]

MLIR and Python (well, the numpy convention) are row-major (like C) so the 2x3 matrix is stored in memory like this

[a_{11}, a_{12}, a_{13}, a_{21}, a_{22}, a_{23}]

so in order to be compatible with the MLIR ordering, we introduce some stablehlo.transpose ops for each input argument at the very beginning and for each result at the very end. you can see them if you disable optimizations with @code_hlo optimize=false. (%0 and %1 in @main)

julia> @code_hlo optimize=false f(a,b)
module @reactant_f attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func private @"*_broadcast_scalar"(%arg0: tensor<f32>, %arg1: tensor<f32>) -> (tensor<f32>, tensor<f32>, tensor<f32>) {
    %0 = stablehlo.multiply %arg0, %arg1 : tensor<f32>
    return %0, %arg0, %arg1 : tensor<f32>, tensor<f32>, tensor<f32>
  }
  func.func @main(%arg0: tensor<3x2xf32> {tf.aliasing_output = 1 : i32}, %arg1: tensor<3x2xf32> {tf.aliasing_output = 2 : i32}) -> (tensor<3x2xf32>, tensor<3x2xf32>, tensor<3x2xf32>) {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<3x2xf32>) -> tensor<2x3xf32>
    %1 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<3x2xf32>) -> tensor<2x3xf32>
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<2x3xf32>
    %2 = stablehlo.broadcast_in_dim %0, dims = [0, 1] : (tensor<2x3xf32>) -> tensor<2x3xf32>
    %3 = stablehlo.broadcast_in_dim %1, dims = [0, 1] : (tensor<2x3xf32>) -> tensor<2x3xf32>
    %4:3 = enzyme.batch @"*_broadcast_scalar"(%2, %3) {batch_shape = array<i64: 2, 3>} : (tensor<2x3xf32>, tensor<2x3xf32>) -> (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>)
    %5 = stablehlo.convert %4#0 : tensor<2x3xf32>
    %6 = stablehlo.transpose %5, dims = [1, 0] : (tensor<2x3xf32>) -> tensor<3x2xf32>
    %7 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2x3xf32>) -> tensor<3x2xf32>
    %8 = stablehlo.transpose %1, dims = [1, 0] : (tensor<2x3xf32>) -> tensor<3x2xf32>
    return %6, %7, %8 : tensor<3x2xf32>, tensor<3x2xf32>, tensor<3x2xf32>
  }
}

all these intermediate ops are optimized away but that the argument dimension remain invariant

julia> println(code)
module @reactant_f attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<3x2xf32>, %arg1: tensor<3x2xf32>) -> tensor<3x2xf32> {
    %0 = stablehlo.multiply %arg0, %arg1 : tensor<3x2xf32>
    return %0 : tensor<3x2xf32>
  }
}

check out that this code is in the MLIR expected order (i.e. 3x2), but hlo_call is expected to reflect the shape of the arrays at the Julia level (i.e. 2x3). if you add a couple of transposes to compensate for it, shapes will match and the resulting code will be equal to code:

julia> f2(a,b) = Ops.transpose(Ops.hlo_call(code, Ops.transpose(a, [2, 1]), Ops.transpose(b, [2, 1]))[1], [2, 1])
f2 (generic function with 1 method)

julia> @code_hlo f2(a,b)
module @reactant_f2 attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<3x2xf32>, %arg1: tensor<3x2xf32>) -> tensor<3x2xf32> {
    %0 = stablehlo.multiply %arg0, %arg1 : tensor<3x2xf32>
    return %0 : tensor<3x2xf32>
  }
}

i know this can be annoying but i'm not sure how we could improve over this

[wsmoses]

I'd be in favor of also adding the transpose to the hlo_call for consistency, or at least add a kwarg

[dubosipsl]

Thanks for these clarifications. Indeed

using Reactant: Ops, to_rarray, @compile, @code_hlo
using Random

a = randn(Float32, 2,3)
at = to_rarray(Array(transpose(a)))
a = to_rarray(a)
b = copy(a)
bt = copy(at)

f(a,b) = a.*b
@compile f(a,b)  # works

code = string(@code_hlo f(a,b))
codet = string(@code_hlo f(at,bt))
println(code)
println(codet)

ff(a,b) = Ops.hlo_call(code, a, b)
ft(a,b) = Ops.hlo_call(codet, a, b)

@compile ff(at,bt)  # works
@compile ft(a,b)  # works

Maybe worth mentioning in the doc for @code_hlo . A flag for hlo_call could be usedul too.

[mofeing]

Is there a reason why you can't just @compile the function directly and you must take the string from @code_hlo and pass it to Ops.hlo_call?

We can add the feature to transpose but I want to better understand your use case as there might be an easier way.

[dubosipsl]

My real use case is to have JAX generate hlo code, compile it with Reactant and use it from within Julia (and possibly differentiate with Enzyme).

This is here (warning, WIP):
https://github.com/ClimFlows/JaxCall.jl/blob/main/src/JaxCall.jl

I use Reactant.@code_hlo to test my Julia code without resorting to JAX. This is where I ran into these issue.

I guess the most important part is to be aware of the convention. Having call_hlo do the transpose might be useful but can also be done by the user.

[mofeing]

aha, the thing is that we already kinda support Jax! Some examples are:

• https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/12
• https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/15

The integration is done in https://github.com/EnzymeAD/Reactant.jl/blob/main/ext/ReactantPythonCallExt/pycall.jl

The only restriction as of right now, is that the arguments and results must be arrays only.

[dubosipsl]

Thanks, I actually started from these examples. What I have seen (I am far from expert) is that the weights of JAX ML models are stored in nested dictionaries, so I need to pass somehow this kind of argument to the compiled function.

Right now I convert the Python dict once to a nested NamedTuple, and I flatten it to a plain Tuple of arrays at each call.

Also, I want to call the function from "normal" Julia code using Arrays, so I convert to/from RArrays before/after calling.

Is this convoluted ? What could be made simpler ?