avoid excessive inlining by moving YaoBlocks.mat to func.func defs

[mofeing]

this is the MWE i'm working with

using Reactant
using YaoBlocks

θ = ConcreteRNumber(rand())

f(x) = mat(ComplexF64, Rz(x))

@code_hlo optimize = false f(θ)

on @code_hlo optimize=false seems to work well

#= /Users/mofeing/Developer/fujitsu/rz.jl:11 =# @code_hlo(optimize = false, f(θ)) = "builtin.module"() ({
  "func.func"() <{function_type = (tensor<f64>) -> tensor<2x2xcomplex<f64>>, sym_name = "rz_Float64_ComplexF64"}> ({
  ^bb0(%arg1: tensor<f64>):
    %4 = "stablehlo.constant"() <{value = dense<(0.000000e+00,0.000000e+00)> : tensor<2x2xcomplex<f64>>}> : () -> tensor<2x2xcomplex<f64>>
    %5 = "stablehlo.constant"() <{value = dense<(0.000000e+00,1.000000e+00)> : tensor<complex<f64>>}> : () -> tensor<complex<f64>>
    %6 = "stablehlo.convert"(%arg1) : (tensor<f64>) -> tensor<complex<f64>>
    %7 = "stablehlo.multiply"(%5, %6) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
    %8 = "stablehlo.constant"() <{value = dense<(2.000000e+00,0.000000e+00)> : tensor<complex<f64>>}> : () -> tensor<complex<f64>>
    %9 = "stablehlo.divide"(%7, %8) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
    %10 = "stablehlo.exponential"(%9) : (tensor<complex<f64>>) -> tensor<complex<f64>>
    %11 = "chlo.conj"(%10) : (tensor<complex<f64>>) -> tensor<complex<f64>>
    %12 = "stablehlo.broadcast_in_dim"(%11) <{broadcast_dimensions = array<i64>}> : (tensor<complex<f64>>) -> tensor<1x1xcomplex<f64>>
    %13 = "stablehlo.constant"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
    %14 = "stablehlo.constant"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
    %15 = "stablehlo.subtract"(%13, %14) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    %16 = "stablehlo.constant"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
    %17 = "stablehlo.constant"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
    %18 = "stablehlo.subtract"(%16, %17) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    %19 = "stablehlo.dynamic_update_slice"(%4, %12, %15, %18) : (tensor<2x2xcomplex<f64>>, tensor<1x1xcomplex<f64>>, tensor<i64>, tensor<i64>) -> tensor<2x2xcomplex<f64>>
    %20 = "stablehlo.broadcast_in_dim"(%10) <{broadcast_dimensions = array<i64>}> : (tensor<complex<f64>>) -> tensor<1x1xcomplex<f64>>
    %21 = "stablehlo.constant"() <{value = dense<2> : tensor<i64>}> : () -> tensor<i64>
    %22 = "stablehlo.constant"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
    %23 = "stablehlo.subtract"(%21, %22) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    %24 = "stablehlo.constant"() <{value = dense<2> : tensor<i64>}> : () -> tensor<i64>
    %25 = "stablehlo.constant"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
    %26 = "stablehlo.subtract"(%24, %25) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    %27 = "stablehlo.dynamic_update_slice"(%19, %20, %23, %26) : (tensor<2x2xcomplex<f64>>, tensor<1x1xcomplex<f64>>, tensor<i64>, tensor<i64>) -> tensor<2x2xcomplex<f64>>
    "func.return"(%27) : (tensor<2x2xcomplex<f64>>) -> ()
  }) : () -> ()
  "func.func"() <{function_type = (tensor<f64>) -> (tensor<2x2xcomplex<f64>>, tensor<f64>), sym_name = "main"}> ({
  ^bb0(%arg0: tensor<f64>):
    %0 = "stablehlo.transpose"(%arg0) <{permutation = array<i64>}> : (tensor<f64>) -> tensor<f64>
    %1 = "func.call"(%0) : (tensor<f64>) -> tensor<2x2xcomplex<f64>>
    %2 = "stablehlo.transpose"(%1) <{permutation = array<i64: 1, 0>}> : (tensor<2x2xcomplex<f64>>) -> tensor<2x2xcomplex<f64>>
    %3 = "stablehlo.transpose"(%0) <{permutation = array<i64>}> : (tensor<f64>) -> tensor<f64>
    "func.return"(%2, %3) : (tensor<2x2xcomplex<f64>>, tensor<f64>) -> ()
  }) : () -> ()
}) : () -> ()

but when i do optimize=true, it segfaults

[69071] signal 11 (2): Segmentation fault: 11
in expression starting at /Users/mofeing/Developer/fujitsu/rz.jl:12
_ZNK4mlir13SymbolRefAttr16getRootReferenceEv at /Users/mofeing/.julia/artifacts/0b9afde9c36d7e3732aedb550a7b040a67ec37d2/lib/libReactantExtra.dylib (unknown line)
_ZL18lookupSymbolInImplPN4mlir9OperationENS_13SymbolRefAttrERN4llvm15SmallVectorImplIS1_EENS3_12function_refIFS1_S1_NS_10StringAttrEEEE at /Users/mofeing/.julia/artifacts/0b9afde9c36d7e3732aedb550a7b040a67ec37d2/lib/libReactantExtra.dylib (unknown line)
Allocations: 42523755 (Pool: 42523097; Big: 658); GC: 26
fish: Job 1, 'julia --project=vqe rz.jl' terminated by signal SIGSEGV (Address boundary error)

[mofeing]

There's sth suspicious: the func.call op doesn't seem to have any attribute (it should have one: callee).

this code...

op = func.call(
    [R.theta.mlir_data];
    result_0=[IR.TensorType((2, 2), IR.Type(T))],
    callee=IR.Attribute(symname("rz", S, T)),
)

@show IR.nattrs(op) IR.Attribute(symname("rz", S, T)) IR.name(op)

IR.attr!(op, "callee", IR.Attribute(symname("rz", S, T)))

@show IR.nattrs(op)

...prints...

IR.nattrs(op) = 0
IR.Attribute(symname("rz", S, T)) = Attribute(#= "rz_Float64_ComplexF64" =#)
IR.name(op) = "func.call"
IR.nattrs(op) = 0

UPDATE

I had to use a IR.FlatSymbolRefAttribute for the callee 🤦

@jumerckx we should add better checks to the op bindings on the new generator

[mofeing]

finally works.

some notes:

• it's adding some redundant transpose and reshape after each func.call. i guess it has to do with the Julia-MLIR interface? they are later optimized away but this shouldn't apply here:

    %1004 = call @rz_Float64_ComplexF64(%271) : (tensor<f64>) -> tensor<2x2xcomplex<f64>>
    %1005 = stablehlo.transpose %1004, dims = [1, 0] : (tensor<2x2xcomplex<f64>>) -> tensor<2x2xcomplex<f64>>
    %1006 = stablehlo.reshape %1005 : (tensor<2x2xcomplex<f64>>) -> tensor<2x2xcomplex<f64>>
    %1007 = stablehlo.transpose %1006, dims = [1, 0] : (tensor<2x2xcomplex<f64>>) -> tensor<2x2xcomplex<f64>>

• before optimization, it reduces expectation computation of the VQE simulation from 16.4k loc to 11k loc
• optimization is inlining the calls, which is not what i want because it ends up generating the same code

[Pangoraw]

You should be able to use a SymbolTable for this: https://github.com/EnzymeAD/Reactant.jl/blob/main/src/mlir/IR/SymbolTable.jl

[Pangoraw]

What is causing this ? This seems suspicious