[mlir] Simplify `affine.min/max (a)` -> `a`

[qcolombet]

Currently we don't fold/canonicalize affine.min/maxs that only return their one value, e.g., %min = affine.min affine_map<()[s0] -> (s0)> ()[%val].

This is not a problem after the lowering of affine because we properly propagate that one value, but this may prevent other optimizations earlier in the pipeline to follow the def-uses chain of this one value.

To reproduce:

func.func @min.oneval(%arg0: index) -> index {
  %min = affine.min affine_map<()[s0] -> (s0)> ()[%arg0]
  return %min: index
}

mlir-opt -test-constant-fold <input>.mlir

Result:

func.func @min.oneval(%arg0: index) -> index {
  %min = affine.min affine_map<()[s0] -> (s0)> ()[%arg0]
  return %min: index
}

Expected result:

func.func @min.oneval(%arg0: index) -> index {
  return %arg0 : index
}

Similarly we don't optimize silly patterns like:

func.func @min.twovals(%arg0: index) -> index {
  %min = affine.min affine_map<()[s0] -> (s0, s0)> ()[%arg0]
  return %min: index
}

Which generates silly code after lowering:

func.func @min.twovals(%arg0: index) -> index {
  %0 = arith.cmpi slt, %arg0, %arg0 : index
  %1 = arith.select %0, %arg0, %arg0 : index
  return %1 : index
}

Hopefully that gets optimize away later, but if we can avoid it that would be great.

Note: The simplification could probably be generalized to something like "if all expressions of the affine_map are a constant function of the same symbol(s), find the min/max of these expressions and propagate this one" (e.g., by materializing an affine.apply of the relevant expression). Finding the min/max expression may be complicated though.

[llvmbot]

@llvm/issue-subscribers-mlir

[llvmbot]

@llvm/issue-subscribers-mlir-affine

[Lewuathe]

I've prepared a patch here. https://reviews.llvm.org/D141354

[qcolombet]

Thanks for the fix!