[Codegen] Run LoopCoalescingPass at the end of warp reduce

[IanWood1]

I observed that this improved some of the punet dispatches using warp reduction from 80us to 60us (example). However, I'm not seeing the improvement through the noise of CI. Locally, I saw ~0.7ms improvement to punet.

[pashu123]

For reviewers:

  %2 = scf.for %arg0 = %c0 to %c60 step %c4 iter_args(%arg1 = %cst_1) -> (vector<1x4xf32>) {
    %60 = scf.for %arg2 = %c0 to %c3840 step %c256 iter_args(%arg3 = %arg1) -> (vector<1x4xf32>) {
      %61 = affine.apply affine_map<(d0)[s0] -> (d0 + s0 floordiv 64)>(%arg0)[%thread_id_x]
      %62 = affine.apply affine_map<(d0)[s0] -> (d0 + s0 * 4 - (s0 floordiv 64) * 256)>(%arg2)[%thread_id_x]
      %63 = vector.transfer_read %0[%workgroup_id_y, %workgroup_id_x, %61, %62], %cst_2 {in_bounds = [true, true]} : memref<2x32x60x3840xf16, strided<[7372800, 230400, 3840, 1], offset: 42741504>, #hal.descriptor
_type<storage_buffer>>, vector<1x4xf16>
      %64 = arith.extf %63 : vector<1x4xf16> to vector<1x4xf32>
      %65 = arith.addf %64, %arg3 : vector<1x4xf32>
      scf.yield %65 : vector<1x4xf32>
    }
    scf.yield %60 : vector<1x4xf32>
  }

It removes the outer loop with affine delinearize (apart from normalizing the loops).

 %2 = scf.for %arg0 = %c0_4 to %c225 step %c1 iter_args(%arg1 = %cst_1) -> (vector<1x4xf32>) {
      %60:2 = affine.delinearize_index %arg0 into (15, 15) : index, index
      %61 = affine.apply #map(%60#1)
      %62 = affine.apply #map1(%60#0)
      %63 = affine.apply #map2(%62)[%thread_id_x]
      %64 = affine.apply #map3(%61)[%thread_id_x]
      %65 = vector.transfer_read %0[%workgroup_id_y, %workgroup_id_x, %63, %64], %cst_2 {in_bounds = [true, true]} : memref<2x32x60x3840xf16, strided<[7372800, 230400, 3840, 1], offset: 42741504>, #hal.descriptor
_type<storage_buffer>>, vector<1x4xf16>
      %66 = arith.extf %65 : vector<1x4xf16> to vector<1x4xf32>
      %67 = arith.addf %66, %arg1 : vector<1x4xf32>
      scf.yield %67 : vector<1x4xf32>
  }

[pashu123]

LGTM! Please make sure onnx-test-suite passes.

[MaheshRavishankar]

Let's land this. Maybe try onnx suite to be sure, by this seems harmless enough. If we have issues we can always revert

[IanWood1]

Okay, I'll run the onnx tests first. workflow here: https://github.com/nod-ai/SHARK-TestSuite/actions/runs/13257180035

[kuhar]

Could we have a LIT test for this?

[IanWood1]

Could we have a LIT test for this?

Any suggestions on where to add a test? should it be added to https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_rocm.mlir?

[MaheshRavishankar]

That would be the place.