feat: generalize loop raising passes to work with affine indexing

[avik-pal]

fixes #1675

• loop info
  • affine index propagation
  • remove the loop offset map
• loop to batch
  • general affine indexing support
  • step != 1
  • negative scaling
  • iota indexing with scaled index
  • constant iteration argument
• while is copy simplify
  • support affine indexing
  • step != 1
  • negative scaling
• Miscellaneous issues
  • fix regression in nested_loop_conv test
  • regression in other testing
  • infinite compile time
• Multiple indices depend on ind var

[github-actions]

EnzymeJAX Benchmarks

Benchmark suite	Current: 858c436	Previous: 8df9e47	Ratio
scatter_sum / JaX / cpu / Primal	0.00000435670900005789 s	0.000004262526800084742 s	1.02
scatter_sum / JaXPipe / cpu / Primal	0.000004330562999894028 s	0.000004347323699948902 s	1.00
scatter_sum / JaX / tpu / Primal	0.0001417065681001 s	0.0001549653354999 s	0.91
scatter_sum / JaXPipe / tpu / Primal	0.0001430293372002 s	0.0001542208233999 s	0.93

This comment was automatically generated by workflow using github-action-benchmark.

[avik-pal]

using Reactant

Reactant.Compiler.WHILE_UNROLL_THRESHOLD[] = 0

function negative_mul_indexing(x)
    y = similar(x)
    fill!(y, 0)
    @trace for i in 1:2:6
        @allowscalar y[2:6, 20 - 2i, 3:4] = cos.(x[1:5, 20 - 3i, 1:2])
    end
    return y
end

x = Reactant.to_rarray(rand(Float32, 6, 20, 5));

@code_hlo negative_mul_indexing(x)

module {
  func.func @main(%arg0: tensor<5x20x6xf32> {enzymexla.memory_effects = []}) -> tensor<5x20x6xf32> attributes {enzymexla.memory_effects = []} {
    %c = stablehlo.constant dense<[[2, 17, 1], [2, 13, 1], [2, 9, 1]]> : tensor<3x3xi32>
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<5x20x6xf32>
    %0 = stablehlo.slice %arg0 [0:2, 4:17:6, 0:5] : (tensor<5x20x6xf32>) -> tensor<2x3x5xf32>
    %1 = stablehlo.reverse %0, dims = [1] : tensor<2x3x5xf32>
    %2 = stablehlo.cosine %1 : tensor<2x3x5xf32>
    %3 = "stablehlo.scatter"(%cst, %c, %2) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0, 2], inserted_window_dims = [1], scatter_dims_to_operand_dims = [0, 1, 2], index_vector_dim = 1>, unique_indices = true}> ({
    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
      stablehlo.return %arg2 : tensor<f32>
    }) : (tensor<5x20x6xf32>, tensor<3x3xi32>, tensor<2x3x5xf32>) -> tensor<5x20x6xf32>
    return %3 : tensor<5x20x6xf32>
  }
}

[wsmoses]

the jax integration tests complain

[avik-pal]

looking into them