[RFC] Add min/max operations

[RFC] Add min/max operations.

I propose to add MinOp, MaxOp operations to Standard or Math dialects. There were plans to move arithmetics from Standard dialect into some new dialects, if that is still the case, then I volunteer to do the split.

Background

Legalization of higher-level dialects that operate on tensors/memrefs, e.g. HLO, TF, results in a loop nest or some operation like linalg.generic with scalar operations in its region. For most operations there is a corresponding op in Standard or Math dialects. Sum of two tensors mhlo.add can be expressed with std.addf, hyperbolic tangent is mapped to math.tanh.

Unfortunately, mhlo.maximum is rewritten as 2 ops when we deal with integers

  %0 = cmpi sgt, %lhs, %rhs : i32
  %1 = select %0, %lhs, %rhs : i32

and 5 ops when we have an op on floats that correctly propagates NANs

  %NAN = constant 0x7FC00000 : f32
  %0 = cmpf ogt, %lhs, %rhs : f32
  %1 = select  %0, %lhs, %rhs : f32
  %2 = cmpf uno, %lhs, %rhs : f32
  %3 = select  %2, %NAN, %37 : f32

There is nothing wrong with this approach if the compiler does not require to perform vectorization or conversion to atomic operations and it just lowers the IR gradually to LLVM. The only downside that I see here is that MLIR users have to reinvent the wheel and write conversion for maximum/minimum which leads to duplicate code.

It becomes unnecessarily more complex when one needs to pattern match cmp and select to convert to std.atomic_rmw or to a corresponding vector operation. Moreover, LLVM has llvm.maximum, llvm.umax and llvm.smax intrinsics and I think it would make sense to add std.max or math.max to MLIR.

3 Likes

FYI: the split is already underway - [RFC] Split the `int` and `float` dialects from `std` - the current consensus is to have a single arith dialect rather than separate int and float.

Adding max/min operations makes sense, but regardless of where it ends up, I’d recommend maintaining the umax/smax/fmax separation. These are different operations just like divide is different across the three domains.

-Chris

We need to update this section of the doc if/when we move forward with this: MLIR Rationale - MLIR

A big +1 on this proposal.

The plain operations are much more amenable to analysis than their cmp-select counterpart. Kernels that perform min-max operations can simply start with the new ops in a true MLIR sense. And if we want to recognize min-max operations in IR coming from “external” sources, we can provide a single recognizer that deals with all floating-point subtleties in a single place only once, rather than having to repeat this recognition everywhere, with the risk that some recognizers miss the subtleties.

Also,for integral min/max, saturation arithmetic is much easier to recognize from direct min-max operations (and we would even provide the clipping ranges as an attribute).

1 Like

Thanks for moving this forward, Alex! It sounds good to me. We should define the semantics properly. llvm.maximum/llvm.minimum/llvm.maxnum/llvm.minnum are intrinsics for the C/C++ Standard Library. If we want to lower the MLIR variants to those, we should be ok with their semantics, esp. around the Nan/+0.0/-0.0 handling. Other than that, it looks good!

Looking at the LLVM Language Reference, regarding vectorization, we only have the vector.reduce.fmin/fmax variants for now, which don’t align with the NaN semantics that you described. It should be easy to implement a specific lowering for those missing variants, though. Regarding atomicrmw, it seems there is no floating point version for min/max. Not sure what the best approach would be here. Maybe implementing some sort of critical region?

1 Like

Great to see consensus on this, looking forward to having this landed, thanks @pifon2a !

Good catch. Speaking for myself, my understanding of the situation has changed. MLIR is much different compared to LLVM in terms of reducing the cost of introducing operations like this. Even LLVM felt that it was the right tradeoff to eventually add these.

-Chris

1 Like

Sounds great to me too! I had comments similar to @dcaballe’s.

https://reviews.llvm.org/D110540