[mlir] Add mlir::LLVM::FastmathFlags to LLVM instrinsic vector reductions

[bmyerz]

Summary

Rationale:
The LLVM dialect supports passing fastmath flags from floating point ops to LLVMIR instructions. However, not all LLVM ops have the required attribute. This change adds support for fastmath flags to llvm.intr.vector.reduce.{fmin,fmax}. One scenario where this is useful is in lowering llvm.intr.vector.reduce.{fmax,fmin} to LLVMIR with nnan (NoNans) flag so it may be lowered to a shuffle reduction.

Changes:

• Make LLVM_VecReductionF implement the FastmathFlagsInterface; change is modeled on LLVM_UnaryIntrOpF
• Add an assembly format for LLVM_VecReductionF ops. The purpose is to keep existing functionality: avoid printing the fastmath flags attribute when it has its default value (none). Change is modeled on LLVM_UnaryIntrOpBase

[bmyerz]

remove unintended .arcconfig change

[gysit]

Thanks!

LGTM modulo some minor comments.

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

389

nit: drop spaces

404

nit: drop spaces

mlir/test/Dialect/LLVMIR/roundtrip.mlir

520–523

ultra nit: it is slightly preferred if the CHECK lines are directly above the tested IR to improve readability.

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

257–260

It seems the fastmath flags lowering is already tested in llvmmir.mlir (@fastmathFlags)? I would thus drop the two tests added here.

[bmyerz]

address review comments on 1. whitespace, 2. placement of CHECK, 3. redundant test

[gysit]

Great thanks!

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

def LLVM_StackRestoreOp : LLVM_ZeroResultIntrOp<"stackrestore"> {
   let assemblyFormat = "$ptr attr-dict `:` qualified(type($ptr))";
 }
 
 //
 // Vector Reductions.
 //
 
 // LLVM vector reduction over a single vector.
-class LLVM_VecReductionBase<string mnem, Type element>
+class LLVM_VecReductionBase<string mnem, Type element, bit requiresFastmath=0>
     : LLVM_OneResultIntrOp<"vector.reduce." # mnem, [], [0],
-                           [Pure, SameOperandsAndResultElementType]>,
-      Arguments<(ins LLVM_VectorOf<element>)>;
+                           [Pure, SameOperandsAndResultElementType],
+                           requiresFastmath> {
+      dag commonArgs = (ins LLVM_VectorOf<element>:$in);
+}
 
 class LLVM_VecReductionF<string mnem>
-    : LLVM_VecReductionBase<mnem, AnyFloat>;
+    : LLVM_VecReductionBase<mnem, AnyFloat, /*requiresFastmath=*/1> {
+  dag fmfArg = (
+    ins DefaultValuedAttr<LLVM_FastmathFlagsAttr, "{}">:$fastmathFlags);
+  let arguments = !con(commonArgs, fmfArg);
+
+  let assemblyFormat = "`(` operands `)` custom<LLVMOpAttrs>(attr-dict) `:` "
+      "functional-type(operands, results)";
+}
 
 class LLVM_VecReductionI<string mnem>
-    : LLVM_VecReductionBase<mnem, AnySignlessInteger>;
+    : LLVM_VecReductionBase<mnem, AnySignlessInteger> {
+      let arguments = commonArgs;
+}
 
 // LLVM vector reduction over a single vector, with an initial value,
 // and with permission to reassociate the reduction operations.
 class LLVM_VecReductionAccBase<string mnem, Type element>
     : LLVM_OneResultIntrOp<"vector.reduce." # mnem, [], [0],
                            [Pure, SameOperandsAndResultElementType]>,
       Arguments<(ins element:$start_value, LLVM_VectorOf<element>:$input,
                  DefaultValuedAttr<BoolAttr, "false">:$reassoc)> {

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 // -----
 
 func.func @reduce_fmax_f32(%arg0: vector<16xf32>, %arg1: f32) -> f32 {
   %0 = vector.reduction <maxf>, %arg0, %arg1 : vector<16xf32> into f32
   return %0 : f32
 }
 // CHECK-LABEL: @reduce_fmax_f32(
 // CHECK-SAME: %[[A:.*]]: vector<16xf32>, %[[B:.*]]: f32)
-//      CHECK: %[[V:.*]] = "llvm.intr.vector.reduce.fmax"(%[[A]]) : (vector<16xf32>) -> f32
+//      CHECK: %[[V:.*]] = llvm.intr.vector.reduce.fmax(%[[A]]) : (vector<16xf32>) -> f32
 //      CHECK: %[[C0:.*]] = llvm.fcmp "ogt" %[[V]], %[[B]] : f32
 //      CHECK: %[[S0:.*]] = llvm.select %[[C0]], %[[V]], %[[B]] : i1, f32
 //      CHECK: %[[C1:.*]] = llvm.fcmp "uno" %[[V]], %[[B]] : f32
 //      CHECK: %[[NAN:.*]] = llvm.mlir.constant(0x7FC00000 : f32) : f32
 //      CHECK: %[[R:.*]] = llvm.select %[[C1]], %[[NAN]], %[[S0]] : i1, f32
 //      CHECK: return %[[R]] : f32
 
 // -----
 
 func.func @reduce_fmin_f32(%arg0: vector<16xf32>, %arg1: f32) -> f32 {
   %0 = vector.reduction <minf>, %arg0, %arg1 : vector<16xf32> into f32
   return %0 : f32
 }
 // CHECK-LABEL: @reduce_fmin_f32(
 // CHECK-SAME: %[[A:.*]]: vector<16xf32>, %[[B:.*]]: f32)
-//      CHECK: %[[V:.*]] = "llvm.intr.vector.reduce.fmin"(%[[A]]) : (vector<16xf32>) -> f32
+//      CHECK: %[[V:.*]] = llvm.intr.vector.reduce.fmin(%[[A]]) : (vector<16xf32>) -> f32
 //      CHECK: %[[C0:.*]] = llvm.fcmp "olt" %[[V]], %[[B]] : f32
 //      CHECK: %[[S0:.*]] = llvm.select %[[C0]], %[[V]], %[[B]] : i1, f32
 //      CHECK: %[[C1:.*]] = llvm.fcmp "uno" %[[V]], %[[B]] : f32
 //      CHECK: %[[NAN:.*]] = llvm.mlir.constant(0x7FC00000 : f32) : f32
 //      CHECK: %[[R:.*]] = llvm.select %[[C1]], %[[NAN]], %[[S0]] : i1, f32
 //      CHECK: return %[[R]] : f32
 
 // -----

mlir/test/Dialect/LLVMIR/roundtrip.mlir

 
 // CHECK: {{.*}} = llvm.fneg %arg0 : f32
   %10 = llvm.fneg %arg0 {fastmathFlags = #llvm.fastmath<none>} : f32
 
 // 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
   return
 }
 
 // CHECK-LABEL: llvm.func @vararg_func
 llvm.func @vararg_func(%arg0: i32, ...) {
   // CHECK: %{{.*}} = llvm.mlir.constant(1 : i32) : i32
   // CHECK: %{{.*}} = llvm.mlir.constant(1 : i32) : i32
   %0 = llvm.mlir.constant(1 : i32) : i32

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

llvm.func @entry() {
     %7 = llvm.shufflevector %6, %4 [0, 0, 0, 0] : vector<4xf32>
     %8 = llvm.mlir.constant(1 : i64) : i64
     %9 = llvm.insertelement %1, %7[%8 : i64] : vector<4xf32>
     %10 = llvm.mlir.constant(2 : i64) : i64
     %11 = llvm.insertelement %2, %9[%10 : i64] : vector<4xf32>
     %12 = llvm.mlir.constant(3 : i64) : i64
     %v = llvm.insertelement %3, %11[%12 : i64] : vector<4xf32>
 
-    %max = "llvm.intr.vector.reduce.fmax"(%v)
+    %max = llvm.intr.vector.reduce.fmax(%v)
         : (vector<4xf32>) -> f32
     llvm.call @printF32(%max) : (f32) -> ()
     llvm.call @printNewline() : () -> ()
     // CHECK: 4
 
-    %min = "llvm.intr.vector.reduce.fmin"(%v)
+    %min = llvm.intr.vector.reduce.fmin(%v)
         : (vector<4xf32>) -> f32
     llvm.call @printF32(%min) : (f32) -> ()
     llvm.call @printNewline() : () -> ()
     // CHECK: 1
 
     %add1 = "llvm.intr.vector.reduce.fadd"(%0, %v)
         : (f32, vector<4xf32>) -> f32
     llvm.call @printF32(%add1) : (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>)
 
 ; CHECK-LABEL: @fastmath_intr
-define void @fastmath_intr(float %arg1, i32 %arg2) {
+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)
+
   ret void
 }

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

 }
 
 ; CHECK-LABEL:  llvm.func @vector_reductions
 define void @vector_reductions(float %0, <8 x float> %1, <8 x i32> %2) {
   ; CHECK: "llvm.intr.vector.reduce.add"(%{{.*}}) : (vector<8xi32>) -> i32
   %4 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.and"(%{{.*}}) : (vector<8xi32>) -> i32
   %5 = call i32 @llvm.vector.reduce.and.v8i32(<8 x i32> %2)
-  ; CHECK: "llvm.intr.vector.reduce.fmax"(%{{.*}}) : (vector<8xf32>) -> f32
+  ; CHECK: llvm.intr.vector.reduce.fmax(%{{.*}}) : (vector<8xf32>) -> f32
   %6 = call float @llvm.vector.reduce.fmax.v8f32(<8 x float> %1)
-  ; CHECK: "llvm.intr.vector.reduce.fmin"(%{{.*}}) : (vector<8xf32>) -> f32
+  ; CHECK: llvm.intr.vector.reduce.fmin(%{{.*}}) : (vector<8xf32>) -> f32
   %7 = call float @llvm.vector.reduce.fmin.v8f32(<8 x float> %1)
   ; CHECK: "llvm.intr.vector.reduce.mul"(%{{.*}}) : (vector<8xi32>) -> i32
   %8 = call i32 @llvm.vector.reduce.mul.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.or"(%{{.*}}) : (vector<8xi32>) -> i32
   %9 = call i32 @llvm.vector.reduce.or.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.smax"(%{{.*}}) : (vector<8xi32>) -> i32
   %10 = call i32 @llvm.vector.reduce.smax.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.smin"(%{{.*}}) : (vector<8xi32>) -> i32

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

 
 // CHECK-LABEL: @vector_reductions
 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
+  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
+  llvm.intr.vector.reduce.fmin(%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
   // CHECK: call i32 @llvm.vector.reduce.umax.v8i32
   "llvm.intr.vector.reduce.umax"(%arg2) : (vector<8xi32>) -> i32
   // CHECK: call i32 @llvm.vector.reduce.umin.v8i32
   "llvm.intr.vector.reduce.umin"(%arg2) : (vector<8xi32>) -> i32

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

llvm.func @vec_reduce_add_intr_wrong_type(%arg0 : vector<4xi32>) -> f32 {
   %0 = "llvm.intr.vector.reduce.add"(%arg0) : (vector<4xi32>) -> f32
   llvm.return %0 : f32
 }
 
 // -----
 
 llvm.func @vec_reduce_fmax_intr_wrong_type(%arg0 : vector<4xi32>) -> i32 {
   // expected-error @below{{op operand #0 must be LLVM dialect-compatible vector of floating-point}}
-  %0 = "llvm.intr.vector.reduce.fmax"(%arg0) : (vector<4xi32>) -> i32
+  %0 = llvm.intr.vector.reduce.fmax(%arg0) : (vector<4xi32>) -> i32
   llvm.return %0 : i32
 }
 
 // -----
 
 llvm.func @matrix_load_intr_wrong_type(%ptr : !llvm.ptr, %stride : i32) -> f32 {
   // expected-error @below{{op result #0 must be LLVM dialect-compatible vector type, but got 'f32'}}
   %0 = llvm.intr.matrix.column.major.load %ptr, <stride=%stride>

mlir/test/Target/LLVMIR/llvmir.mlir

llvm.func @useInlineAsm(%arg0: i32) {
   llvm.return
 }
 
 // -----
 
 llvm.func @fastmathFlagsFunc(f32) -> f32
 
 // CHECK-LABEL: @fastmathFlags
-llvm.func @fastmathFlags(%arg0: f32) {
+llvm.func @fastmathFlags(%arg0: f32, %arg1 : vector<2xf32>) {
 // CHECK: {{.*}} = fadd nnan ninf float {{.*}}, {{.*}}
 // CHECK: {{.*}} = fsub nnan ninf float {{.*}}, {{.*}}
 // CHECK: {{.*}} = fmul nnan ninf float {{.*}}, {{.*}}
 // CHECK: {{.*}} = fdiv nnan ninf float {{.*}}, {{.*}}
 // CHECK: {{.*}} = frem nnan ninf float {{.*}}, {{.*}}
   %0 = llvm.fadd %arg0, %arg0 {fastmathFlags = #llvm.fastmath<nnan, ninf>} : f32
   %1 = llvm.fsub %arg0, %arg0 {fastmathFlags = #llvm.fastmath<nnan, ninf>} : f32
   %2 = llvm.fmul %arg0, %arg0 {fastmathFlags = #llvm.fastmath<nnan, ninf>} : f32

 // CHECK: call afn float @llvm.copysign.f32(float {{.*}}, float {{.*}})
   %18 = "llvm.intr.copysign"(%arg0, %arg0) {fastmathFlags = #llvm.fastmath<afn>} : (f32, f32) -> f32
 
 // CHECK: call fast float @llvm.powi.f32.i32(float {{.*}}, i32 {{.*}})
   %exp = llvm.mlir.constant(1 : i32) : i32
   %19 = "llvm.intr.powi"(%arg0, %exp) {fastmathFlags = #llvm.fastmath<fast>} : (f32, i32) -> f32
 // 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
+
   llvm.return
 }
 
 // -----
 
 // CHECK-LABEL: @switch_args
 llvm.func @switch_args(%arg0: i32) -> i32 {
   %0 = llvm.mlir.constant(5 : i32) : i32