feat: general concat reshape to batching pass

[avik-pal]

Summarizing the transform pass:

1. Pattern: concat of insert dims (at the same dim) with equivalent ops.
2. For all the parent operands, we insert a dim at dim=0 and concat those.
3. Insert a batch op and immediately resolve the op into a callop using batching utilities
4. Move the first batched dim to the concat dim.

Later we can use this same pass for ops like cholesky / tringular_solve / lu where we need to check that the concat dim is not among the last 2 dims.

For lu it is a bit more complicated since it has 4 returns and we need to check whether all the other returns are either (1) unused or (2) also concatenated.

[wsmoses]

rather than call batch, why not just call the batchopinterface if defined? or other utility function from the implementation of batch lowering

[wsmoses]

e.g. we can make https://github.com/EnzymeAD/Enzyme/blob/db0181320d6e425ee963bd496ed0d8dbb615be18/enzyme/Enzyme/MLIR/Passes/EnzymeBatchPass.cpp#L130 into a batchOperation utility or something, which we call directly here instead of outlining into a new function

[avik-pal]

yeah that sounds like a better thing to do

[avik-pal]

somehow I made it segfault only on mac...

[wsmoses]

That's probably an issue of f(a(), b()) running a before b on Linux and opposite on Mac

[avik-pal]

@wsmoses any idea what https://github.com/EnzymeAD/Enzyme-JAX/actions/runs/17003076741/job/48208359246?pr=983#step:19:133 is?

[avik-pal]

dumping f before crash

func.func private @"Z%\BF\E3=[\00\00\9ClB\1A\E1T\FD\0Bhed_concatReshapeOpToBatch_34"(%arg0: tensor<288xf32>, %arg1: tensor<768xf32>) -> tensor<288x768xf32> {
  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [] x [], precision = [DEFAULT, DEFAULT] : (tensor<288xf32>, tensor<768xf32>) -> tensor<288x768xf32>
  return %0 : tensor<288x768xf32>
}

[avik-pal]

am I doing something wrong here? why is the name not batched_.... https://github.com/EnzymeAD/Enzyme/blob/30735308bbcf50124107ca8204b08ef6dc28fbc6/enzyme/Enzyme/MLIR/Passes/EnzymeBatchPass.cpp#L246C39-L246C47?

[wsmoses]

er.....wat re the linked failure

[wsmoses]

that seems like some sort of use after free undefined memory error

[wsmoses]

we can always just change that to
std::String new_name = ("batched_" + fn.getName()).str();