scf for loop canonicalizations dont propagate

[nirvedhmeshram]

Consider these tests

func.func @scf_for_canonicalize_min(%A : memref<i64>) {
  %c0 = arith.constant 0 : index
  %c2 = arith.constant 2 : index
  %c638 = arith.constant 638 : index

  scf.for %i = %c0 to %c638 step %c2  {
    %1 = affine.min affine_map<(d0) -> (10239, d0 * 16)>(%i)
    %2 = affine.min affine_map<(d0) -> (-d0 + 10239, 32)>(%1)
    %3 = arith.index_cast %2: index to i64
    memref.store %3, %A[]: memref<i64>
  }
  return
}

// -----
func.func @scf_for_canonicalize_min(%A : memref<i64>) {
  %c0 = arith.constant 0 : index
  %c2 = arith.constant 2 : index
  %c638 = arith.constant 638 : index

  scf.for %i = %c0 to %c638 step %c2  {
    %2 = affine.min affine_map<(d0) -> (-d0 * 16 + 10239, 32)>(%i)
    %3 = arith.index_cast %2: index to i64
    memref.store %3, %A[]: memref<i64>
  }
  return
}

With

mlir-opt test.mlir -scf-for-loop-canonicalization --split-input-file --mlir-print-local-scope

We get the following outputs

  func.func @scf_for_canonicalize_min(%arg0: memref<i64>) {
    %c0 = arith.constant 0 : index
    %c2 = arith.constant 2 : index
    %c638 = arith.constant 638 : index
    scf.for %arg1 = %c0 to %c638 step %c2 {
      %0 = affine.apply affine_map<(d0) -> (d0 * 16)>(%arg1)
      %1 = affine.min affine_map<(d0) -> (-d0 + 10239, 32)>(%0)
      %2 = arith.index_cast %1 : index to i64
      memref.store %2, %arg0[] : memref<i64>
    }
    return
  }
// -----
  func.func @scf_for_canonicalize_min(%arg0: memref<i64>) {
    %c32_i64 = arith.constant 32 : i64
    %c0 = arith.constant 0 : index
    %c2 = arith.constant 2 : index
    %c638 = arith.constant 638 : index
    scf.for %arg1 = %c0 to %c638 step %c2 {
      memref.store %c32_i64, %arg0[] : memref<i64>
    }
    return
  }

We could not canonicalize the second affine.min in the first test but can do it in the second test because it directly takes the induction var as an input. What do we need to do to propagate the canonicalization?

[nirvedhmeshram]

@matthias-springer any inputs on how to go about enabling such a canonicalization? cc @Max191

[nirvedhmeshram]

Actually we can just use this pass SimplifyAffineMinMaxPass for simplifying the affine maps so we have a solution.