add @code_hlo macro to get mlir code

[Pangoraw]

Example usage:

julia> using Reactant

julia> W = Reactant.ConcreteRArray(randn(Float32, 10, 20))
       x = Reactant.ConcreteRArray(randn(Float32, 20, 5))
       Reactant.@code_hlo W * x
Module:
module attributes {transform.with_named_sequence} {
  func.func @main(%arg0: tensor<20x10xf32>, %arg1: tensor<5x20xf32>) -> tensor<5x10xf32> {
    %0 = stablehlo.dot_general %arg1, %arg0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<5x20xf32>, tensor<20x10xf32>) -> tensor<5x10xf32>
    return %0 : tensor<5x10xf32>
  }
}

cc @mofeing

[mofeing]

Awesome! Did you manage to fix the problem with sum? Also why does transform.with_named_sequence appears there?

EDIT: sum works and the transform.with_named_sequence seems to be a consequence of some pass in Enzyme-JAX.

[wsmoses]

Oh this is amazing!

One extra thing, would it be possible to have a flag that shows HLO before the pass pipeline of optimizations?

And yeah that's a leftover from the transform dialect we use for optimizations

[wsmoses]

Also @Pangoraw if you'd like, you're more than welcome to add yourself to the contributors [and also if you'd be interested in helping @mofeing and myself out getting some cool things implemented!]

[Pangoraw]

I added a run_pipeline (default=true) option such that the pass pipeline is not run on the module:

julia> W = Reactant.ConcreteRArray(randn(Float32, 10, 20))
       x = Reactant.ConcreteRArray(randn(Float32, 20, 5))
       run_pipeline = false
       Reactant.@code_hlo run_pipeline = run_pipeline W * x
Module:
module {
  func.func @main(%arg0: tensor<5x20xf32>, %arg1: tensor<20x10xf32>) -> (tensor<20x10xf32>, tensor<5x20xf32>, tensor<5x10xf32>) {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<5x20xf32>) -> tensor<20x5xf32>
    %1 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<20x10xf32>) -> tensor<10x20xf32>
    %2 = stablehlo.dot_general %1, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<10x20xf32>, tensor<20x5xf32>) -> tensor<10x5xf32>
    %3 = stablehlo.transpose %1, dims = [1, 0] : (tensor<10x20xf32>) -> tensor<20x10xf32>
    %4 = stablehlo.transpose %0, dims = [1, 0] : (tensor<20x5xf32>) -> tensor<5x20xf32>
    %5 = stablehlo.transpose %2, dims = [1, 0] : (tensor<10x5xf32>) -> tensor<5x10xf32>
    return %3, %4, %5 : tensor<20x10xf32>, tensor<5x20xf32>, tensor<5x10xf32>
  }
}

[wsmoses]

LGTM!