[mlir][LLVM] Introduce reduction intrinsics for minimum/maximum

[unterumarmung]

Summary

This patch adds supports for the reduction intrinsic
for floating point minimum and maximum that have
been added to LLVM by https://reviews.llvm.org/D152370.

Related to: #63969

[dcaballe]

Awesome! I just wrote similar code yesterday! :)

[unterumarmung]

In D155869#4519796, @dcaballe wrote:

Awesome! I just wrote similar code yesterday! :)

Should I proceed with the https://github.com/llvm/llvm-project/issues/63969 or do you have the solution for that? If yes, is my understanding correct that we should just nuke createMinMaxF function that adds boilerplate for the NaN handling?

[dcaballe]

In D155869#4519803, @unterumarmung wrote:

In D155869#4519796, @dcaballe wrote:

Awesome! I just wrote similar code yesterday! :)

Should I proceed with the https://github.com/llvm/llvm-project/issues/63969 or do you have the solution for that? If yes, is my understanding correct that we should just nuke createMinMaxF function that adds boilerplate for the NaN handling?

Thanks, no I just run a few experiments. Yes, we have to replace createMinMaxF with the corresponding maximum/minimum and the reduction intrinsics similarly. Thanks a lot for looking into this!

[unterumarmung]

Address review comments

[unterumarmung]

Pushed changes to wrong revision, revert

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

 def LLVM_vector_reduce_smax : LLVM_VecReductionI<"smax">;
 def LLVM_vector_reduce_smin : LLVM_VecReductionI<"smin">;
 def LLVM_vector_reduce_umax : LLVM_VecReductionI<"umax">;
 def LLVM_vector_reduce_umin : LLVM_VecReductionI<"umin">;
 def LLVM_vector_reduce_xor : LLVM_VecReductionI<"xor">;
 
 def LLVM_vector_reduce_fmax : LLVM_VecReductionF<"fmax">;
 def LLVM_vector_reduce_fmin : LLVM_VecReductionF<"fmin">;
+def LLVM_vector_reduce_fmaximum : LLVM_VecReductionF<"fmaximum">;
+def LLVM_vector_reduce_fminimum : LLVM_VecReductionF<"fminimum">;
 
 def LLVM_vector_reduce_fadd : LLVM_VecReductionAccF<"fadd">;
 def LLVM_vector_reduce_fmul : LLVM_VecReductionAccF<"fmul">;
 
 //
 // LLVM Matrix operations.
 //
 

mlir/test/Dialect/LLVMIR/roundtrip.mlir

// CHECK: {{.*}} = llvm.intr.sin(%arg0) {fastmathFlags = #llvm.fastmath<fast>} : (f32) -> f32
   %11 = llvm.intr.sin(%arg0) {fastmathFlags = #llvm.fastmath<fast>} : (f32) -> f32
 // CHECK: {{.*}} = llvm.intr.sin(%arg0) {fastmathFlags = #llvm.fastmath<afn>} : (f32) -> f32
   %12 = llvm.intr.sin(%arg0) {fastmathFlags = #llvm.fastmath<afn>} : (f32) -> f32
 
 // CHECK: {{.*}} = llvm.intr.vector.reduce.fmin(%arg3) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
   %13 = llvm.intr.vector.reduce.fmin(%arg3) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
 // CHECK: {{.*}} = llvm.intr.vector.reduce.fmax(%arg3) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
   %14 = llvm.intr.vector.reduce.fmax(%arg3) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
+// CHECK: {{.*}} = llvm.intr.vector.reduce.fminimum(%arg3) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
+  %15 = llvm.intr.vector.reduce.fminimum(%arg3) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
+// CHECK: {{.*}} = llvm.intr.vector.reduce.fmaximum(%arg3) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
+  %16 = llvm.intr.vector.reduce.fmaximum(%arg3) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
   return
 }
 
 // CHECK-LABEL: @lifetime
 // CHECK-SAME: %[[P:.*]]: !llvm.ptr
 llvm.func @lifetime(%p: !llvm.ptr) {
   // CHECK: llvm.intr.lifetime.start 16, %[[P]]
   llvm.intr.lifetime.start 16, %p : !llvm.ptr

mlir/test/Integration/Dialect/LLVMIR/CPU/test-vector-reductions-fp.mlir

llvm.func @entry() {
     // CHECK: 4
 
     %min = llvm.intr.vector.reduce.fmin(%v)
         : (vector<4xf32>) -> f32
     llvm.call @printF32(%min) : (f32) -> ()
     llvm.call @printNewline() : () -> ()
     // CHECK: 1
 
+    %maximum = llvm.intr.vector.reduce.fmaximum(%v)
+        : (vector<4xf32>) -> f32
+    llvm.call @printF32(%maximum) : (f32) -> ()
+    llvm.call @printNewline() : () -> ()
+    // CHECK: 4
+
+    %minimum = llvm.intr.vector.reduce.fminimum(%v)
+        : (vector<4xf32>) -> f32
+    llvm.call @printF32(%minimum) : (f32) -> ()
+    llvm.call @printNewline() : () -> ()
+    // CHECK: 1
+
     %add1 = "llvm.intr.vector.reduce.fadd"(%0, %v)
         : (f32, vector<4xf32>) -> f32
     llvm.call @printF32(%add1) : (f32) -> ()
     llvm.call @printNewline() : () -> ()
     // CHECK: 11
 
     %add1r = "llvm.intr.vector.reduce.fadd"(%0, %v)
         {reassoc = true} : (f32, vector<4xf32>) -> f32

mlir/test/Target/LLVMIR/Import/fastmath.ll

 ; // -----
 
 declare float @llvm.exp.f32(float)
 declare float @llvm.powi.f32.i32(float, i32)
 declare float @llvm.pow.f32(float, float)
 declare float @llvm.fmuladd.f32(float, float, float)
 declare float @llvm.vector.reduce.fmin.v2f32(<2 x float>)
 declare float @llvm.vector.reduce.fmax.v2f32(<2 x float>)
+declare float @llvm.vector.reduce.fminimum.v2f32(<2 x float>)
+declare float @llvm.vector.reduce.fmaximum.v2f32(<2 x float>)
 
 ; CHECK-LABEL: @fastmath_intr
 define void @fastmath_intr(float %arg1, i32 %arg2, <2 x float> %arg3) {
   ; CHECK: llvm.intr.exp(%{{.*}}) {fastmathFlags = #llvm.fastmath<nnan, ninf>} : (f32) -> f32
   %1 = call nnan ninf float @llvm.exp.f32(float %arg1)
   ; CHECK: llvm.intr.powi(%{{.*}}, %{{.*}}) {fastmathFlags = #llvm.fastmath<fast>} : (f32, i32) -> f32
   %2 = call fast float @llvm.powi.f32.i32(float %arg1, i32 %arg2)
   ; CHECK:  llvm.intr.pow(%{{.*}}, %{{.*}}) {fastmathFlags = #llvm.fastmath<fast>} : (f32, f32) -> f32
   %3 = call fast float @llvm.pow.f32(float %arg1, float %arg1)
   ; CHECK: llvm.intr.fmuladd(%{{.*}}, %{{.*}}, %{{.*}}) {fastmathFlags = #llvm.fastmath<fast>} : (f32, f32, f32) -> f32
   %4 = call fast float @llvm.fmuladd.f32(float %arg1, float %arg1, float %arg1)
   ; CHECK: %{{.*}} = llvm.intr.vector.reduce.fmin({{.*}}) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
   %5 = call nnan float @llvm.vector.reduce.fmin.v2f32(<2 x float> %arg3)
   ; CHECK: %{{.*}} = llvm.intr.vector.reduce.fmax({{.*}}) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
   %6 = call nnan float @llvm.vector.reduce.fmax.v2f32(<2 x float> %arg3)
+  ; CHECK: %{{.*}} = llvm.intr.vector.reduce.fminimum({{.*}}) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
+  %7 = call nnan float @llvm.vector.reduce.fminimum.v2f32(<2 x float> %arg3)
+  ; CHECK: %{{.*}} = llvm.intr.vector.reduce.fmaximum({{.*}}) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
+  %8 = call nnan float @llvm.vector.reduce.fmaximum.v2f32(<2 x float> %arg3)
 
   ret void
 }

mlir/test/Target/LLVMIR/Import/intrinsic.ll

define void @vector_reductions(float %0, <8 x float> %1, <8 x i32> %2) {
   ; CHECK: "llvm.intr.vector.reduce.fmul"(%{{.*}}, %{{.*}}) <{reassoc = false}> : (f32, vector<8xf32>) -> f32
   %15 = call float @llvm.vector.reduce.fmul.v8f32(float %0, <8 x float> %1)
   ; CHECK: "llvm.intr.vector.reduce.fadd"(%{{.*}}, %{{.*}}) <{reassoc = true}> : (f32, vector<8xf32>) -> f32
   %16 = call reassoc float @llvm.vector.reduce.fadd.v8f32(float %0, <8 x float> %1)
   ; CHECK: "llvm.intr.vector.reduce.fmul"(%{{.*}}, %{{.*}}) <{reassoc = true}> : (f32, vector<8xf32>) -> f32
   %17 = call reassoc float @llvm.vector.reduce.fmul.v8f32(float %0, <8 x float> %1)
   ; CHECK:  "llvm.intr.vector.reduce.xor"(%{{.*}}) : (vector<8xi32>) -> i32
   %18 = call i32 @llvm.vector.reduce.xor.v8i32(<8 x i32> %2)
+  ; CHECK: llvm.intr.vector.reduce.fmaximum(%{{.*}}) : (vector<8xf32>) -> f32
+  %19 = call float @llvm.vector.reduce.fmaximum.v8f32(<8 x float> %1)
+  ; CHECK: llvm.intr.vector.reduce.fminimum(%{{.*}}) : (vector<8xf32>) -> f32
+  %20 = call float @llvm.vector.reduce.fminimum.v8f32(<8 x float> %1)
   ret void
 }
 
 ; CHECK-LABEL: @matrix_intrinsics
 ; CHECK-SAME:  %[[VEC1:[a-zA-Z0-9]+]]
 ; CHECK-SAME:  %[[VEC2:[a-zA-Z0-9]+]]
 ; CHECK-SAME:  %[[PTR:[a-zA-Z0-9]+]]
 ; CHECK-SAME:  %[[STRIDE:[a-zA-Z0-9]+]]

 declare i32 @llvm.umax.i32(i32, i32)
 declare <8 x i32> @llvm.umax.v8i32(<8 x i32>, <8 x i32>)
 declare i32 @llvm.umin.i32(i32, i32)
 declare <8 x i32> @llvm.umin.v8i32(<8 x i32>, <8 x i32>)
 declare i32 @llvm.vector.reduce.add.v8i32(<8 x i32>)
 declare i32 @llvm.vector.reduce.and.v8i32(<8 x i32>)
 declare float @llvm.vector.reduce.fmax.v8f32(<8 x float>)
 declare float @llvm.vector.reduce.fmin.v8f32(<8 x float>)
+declare float @llvm.vector.reduce.fmaximum.v8f32(<8 x float>)
+declare float @llvm.vector.reduce.fminimum.v8f32(<8 x float>)
 declare i32 @llvm.vector.reduce.mul.v8i32(<8 x i32>)
 declare i32 @llvm.vector.reduce.or.v8i32(<8 x i32>)
 declare i32 @llvm.vector.reduce.smax.v8i32(<8 x i32>)
 declare i32 @llvm.vector.reduce.smin.v8i32(<8 x i32>)
 declare i32 @llvm.vector.reduce.umax.v8i32(<8 x i32>)
 declare i32 @llvm.vector.reduce.umin.v8i32(<8 x i32>)
 declare float @llvm.vector.reduce.fadd.v8f32(float, <8 x float>)
 declare float @llvm.vector.reduce.fmul.v8f32(float, <8 x float>)

mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

llvm.func @vector_reductions(%arg0: f32, %arg1: vector<8xf32>, %arg2: vector<8xi32>) {
   // CHECK: call i32 @llvm.vector.reduce.add.v8i32
   "llvm.intr.vector.reduce.add"(%arg2) : (vector<8xi32>) -> i32
   // CHECK: call i32 @llvm.vector.reduce.and.v8i32
   "llvm.intr.vector.reduce.and"(%arg2) : (vector<8xi32>) -> i32
   // CHECK: call float @llvm.vector.reduce.fmax.v8f32
   llvm.intr.vector.reduce.fmax(%arg1) : (vector<8xf32>) -> f32
   // CHECK: call float @llvm.vector.reduce.fmin.v8f32
   llvm.intr.vector.reduce.fmin(%arg1) : (vector<8xf32>) -> f32
+  // CHECK: call float @llvm.vector.reduce.fmaximum.v8f32
+  llvm.intr.vector.reduce.fmaximum(%arg1) : (vector<8xf32>) -> f32
+  // CHECK: call float @llvm.vector.reduce.fminimum.v8f32
+  llvm.intr.vector.reduce.fminimum(%arg1) : (vector<8xf32>) -> f32
   // CHECK: call i32 @llvm.vector.reduce.mul.v8i32
   "llvm.intr.vector.reduce.mul"(%arg2) : (vector<8xi32>) -> i32
   // CHECK: call i32 @llvm.vector.reduce.or.v8i32
   "llvm.intr.vector.reduce.or"(%arg2) : (vector<8xi32>) -> i32
   // CHECK: call i32 @llvm.vector.reduce.smax.v8i32
   "llvm.intr.vector.reduce.smax"(%arg2) : (vector<8xi32>) -> i32
   // CHECK: call i32 @llvm.vector.reduce.smin.v8i32
   "llvm.intr.vector.reduce.smin"(%arg2) : (vector<8xi32>) -> i32

mlir/test/Target/LLVMIR/llvmir.mlir

 // CHECK: call afn float @llvm.powi.f32.i32(float {{.*}}, i32 {{.*}})
   %20 = "llvm.intr.powi"(%arg0, %exp) {fastmathFlags = #llvm.fastmath<afn>} : (f32, i32) -> f32
 
 // CHECK: call nnan float @llvm.vector.reduce.fmax.v2f32(<2 x float> {{.*}})
 // CHECK: call nnan float @llvm.vector.reduce.fmin.v2f32(<2 x float> {{.*}})
   %21 = llvm.intr.vector.reduce.fmax(%arg1) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
   %22 = llvm.intr.vector.reduce.fmin(%arg1) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
 
+// CHECK: call nnan float @llvm.vector.reduce.fmaximum.v2f32(<2 x float> {{.*}})
+// CHECK: call nnan float @llvm.vector.reduce.fminimum.v2f32(<2 x float> {{.*}})
+  %23 = llvm.intr.vector.reduce.fmaximum(%arg1) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
+  %24 = llvm.intr.vector.reduce.fminimum(%arg1) {fastmathFlags = #llvm.fastmath<nnan>} : (vector<2xf32>) -> f32
 
-  %23 = llvm.mlir.constant(true) : i1
+  %25 = llvm.mlir.constant(true) : i1
 // CHECK: select contract i1
-  %24 = llvm.select %23, %arg0, %20 {fastmathFlags = #llvm.fastmath<contract>} : i1, f32
+  %26 = llvm.select %25, %arg0, %20 {fastmathFlags = #llvm.fastmath<contract>} : i1, f32
   llvm.return
 }
 
 // -----
 
 // CHECK-LABEL: @switch_args
 llvm.func @switch_args(%arg0: i32) -> i32 {
   %0 = llvm.mlir.constant(5 : i32) : i32