LBM kernel optimizations

This is LBM after megakernelization, arg refinement and running enzyme-hlo-opt.

```mlir
// Run with:
// enzymexlamlir-opt --enzyme-hlo-opt lbm_gather_scatter_mwes.mlir
//
// These are reduced versions of the four LBM indexing cases discussed in
// lbm.refined.out.mlir. Small dimensions are used here:
//
// physical row width = 8
// logical tile = 2 x 3
//
// The corresponding LBM dimensions are a physical row width of 128 and a
// logical tile of 120 x 18000.

module {
// ------------------------------------------------------------------------
// Case 1: float gather whose second affine iota is hidden under
// broadcast_in_dim(multiply(iota, stride)).
//
// indices[i, j] = 8 + i + 8*j
//
// This is a regular 2-D tile. It can be implemented as:
// reshape 64 -> 8x8
// slice [1:4, 0:2] -> 3x2
// transpose -> 2x3
//
// Currently the gather remains because detectIotaLikeTensor does not look
// through the broadcast_in_dim.
// ------------------------------------------------------------------------
func.func @gather_broadcast_scaled_iota(
%input: tensor<64xf32>) -> tensor<2x3xf32> {
%stride = stablehlo.constant dense<8> : tensor<3xi64>
%offset = stablehlo.constant dense<8> : tensor<2x3x1xi64>

%i = stablehlo.iota dim = 0 : tensor<2x3x1xi64>
%j = stablehlo.iota dim = 0 : tensor<3xi64>
%scaled_j = stablehlo.multiply %j, %stride : tensor<3xi64>
%broadcast_j = stablehlo.broadcast_in_dim %scaled_j, dims = [1]
: (tensor<3xi64>) -> tensor<2x3x1xi64>
%grid = stablehlo.add %i, %broadcast_j : tensor<2x3x1xi64>
%indices = stablehlo.add %grid, %offset : tensor<2x3x1xi64>

%result = "stablehlo.gather"(%input, %indices) <{
dimension_numbers = #stablehlo.gather<
collapsed_slice_dims = [0],
start_index_map = [0],
index_vector_dim = 2>,
indices_are_sorted = false,
slice_sizes = array
}> : (tensor<64xf32>, tensor<2x3x1xi64>) -> tensor<2x3xf32>
return %result : tensor<2x3xf32>
}

// ------------------------------------------------------------------------
// Case 2: byte/flag gather whose first affine iota is hidden under a
// reshape(multiply(iota, byte_stride)).
//
// indices[i, j] = 16 + 4*i + 32*j
//
// This models the LBM flag load after the f32 buffer is bitcast to bytes.
// Currently the gather remains because detectIotaLikeTensor does not look
// through the reshape.
// ------------------------------------------------------------------------
func.func @gather_reshaped_scaled_iota(
%bytes: tensor<256xi8>) -> tensor<2x3xi8> {
%inner_stride = stablehlo.constant dense<4> : tensor<2x3xi64>
%row_stride = stablehlo.constant dense<32> : tensor<2x3x1xi64>
%offset = stablehlo.constant dense<16> : tensor<2x3x1xi64>

%i = stablehlo.iota dim = 0 : tensor<2x3xi64>
%scaled_i = stablehlo.multiply %i, %inner_stride : tensor<2x3xi64>
%reshaped_i = stablehlo.reshape %scaled_i
: (tensor<2x3xi64>) -> tensor<2x3x1xi64>

%j = stablehlo.iota dim = 1 : tensor<2x3x1xi64>
%scaled_j = stablehlo.multiply %j, %row_stride : tensor<2x3x1xi64>
%grid = stablehlo.add %reshaped_i, %scaled_j : tensor<2x3x1xi64>
%indices = stablehlo.add %grid, %offset : tensor<2x3x1xi64>

%result = "stablehlo.gather"(%bytes, %indices) <{
dimension_numbers = #stablehlo.gather<
collapsed_slice_dims = [0],
start_index_map = [0],
index_vector_dim = 2>,
indices_are_sorted = false,
slice_sizes = array
}> : (tensor<256xi8>, tensor<2x3x1xi64>) -> tensor<2x3xi8>
return %result : tensor<2x3xi8>
}

// ------------------------------------------------------------------------
// Case 3: non-wrapping rectangular set-index scatter with two affine iotas.
//
// indices[i, j] = 9 + i + 8*j
//
// In an 8-column view this updates rows [1, 4) and columns [1, 3).
// It can be implemented as:
// reshape destination 64 -> 8x8
// transpose updates 2x3 -> 3x2
// dynamic_update_slice at [1, 1]
// reshape 8x8 -> 64
//
// Currently the scatter remains because ScatterOpCanon only recognizes a
// single affine iota, not a sum of independently varying iotas.
// ------------------------------------------------------------------------
func.func @scatter_rectangular_multi_iota(
%destination: tensor<64xf32>,
%updates: tensor<2x3xf32>) -> tensor<64xf32> {
%stride = stablehlo.constant dense<8> : tensor<2x3x1xi64>
%offset = stablehlo.constant dense<9> : tensor<2x3x1xi64>

%i = stablehlo.iota dim = 0 : tensor<2x3x1xi64>
%j = stablehlo.iota dim = 1 : tensor<2x3x1xi64>
%scaled_j = stablehlo.multiply %j, %stride : tensor<2x3x1xi64>
%grid = stablehlo.add %i, %scaled_j : tensor<2x3x1xi64>
%indices = stablehlo.add %grid, %offset : tensor<2x3x1xi64>

%result = "stablehlo.scatter"(%destination, %indices, %updates) <{
indices_are_sorted = false,
scatter_dimension_numbers = #stablehlo.scatter<
inserted_window_dims = [0],
scatter_dims_to_operand_dims = [0],
index_vector_dim = 2>,
unique_indices = true
}> ({
^bb0(%old: tensor, %update: tensor):
stablehlo.return %update : tensor
}) : (tensor<64xf32>, tensor<2x3x1xi64>, tensor<2x3xf32>)
-> tensor<64xf32>
return %result : tensor<64xf32>
}

// ------------------------------------------------------------------------
// Case 4: set-index scatter whose tile starts in the final physical column.
//
// indices[i, j] = 7 + i + 8*j
//
// The i=0 values update column 7 of rows [0, 3), while the i=1 values wrap
// to column 0 of rows [1, 4). Consequently this is not one rectangular
// dynamic_update_slice. A correct lowering must split it into two updates:
//
// updates[0, :] -> transpose to 3x1 -> DUS at [0, 7]
// updates[1, :] -> transpose to 3x1 -> DUS at [1, 0]
//
// This is the reduced analogue of the LBM offsets whose remainder modulo
// 128 is 127.
// ------------------------------------------------------------------------
func.func @scatter_wraps_physical_row(
%destination: tensor<64xf32>,
%updates: tensor<2x3xf32>) -> tensor<64xf32> {
%stride = stablehlo.constant dense<8> : tensor<2x3x1xi64>
%offset = stablehlo.constant dense<7> : tensor<2x3x1xi64>

%i = stablehlo.iota dim = 0 : tensor<2x3x1xi64>
%j = stablehlo.iota dim = 1 : tensor<2x3x1xi64>
%scaled_j = stablehlo.multiply %j, %stride : tensor<2x3x1xi64>
%grid = stablehlo.add %i, %scaled_j : tensor<2x3x1xi64>
%indices = stablehlo.add %grid, %offset : tensor<2x3x1xi64>

%result = "stablehlo.scatter"(%destination, %indices, %updates) <{
indices_are_sorted = false,
scatter_dimension_numbers = #stablehlo.scatter<
inserted_window_dims = [0],
scatter_dims_to_operand_dims = [0],
index_vector_dim = 2>,
unique_indices = true
}> ({
^bb0(%old: tensor, %update: tensor):
stablehlo.return %update : tensor
}) : (tensor<64xf32>, tensor<2x3x1xi64>, tensor<2x3xf32>)
-> tensor<64xf32>
return %result : tensor<64xf32>
}
}

```r

Start by running enzymexlamlir-opt --enzyme-hlo-opt on lbm_gather_scatter_mwes.mlir and compare the four cases with lbm.refined.out.mlir. Read detectIotaLikeTensor and ScatterOpCanon first. Done means the broadcast/reshape iotas and rectangular scatters are optimized, while the wrapping case is split into the two described updates.

Written by the indexing model from the issue text.