[Mlir-commits] [mlir] 79afa58 - [mlir][llvm] Add fastmathFlags to SelectOp

[Jeff Niu]

Author: Jeff Niu
Date: 2023-03-16T08:57:52+01:00
New Revision: 79afa58c2779db21d1ee7d65d32d5250e58e7226

URL: https://github.com/llvm/llvm-project/commit/79afa58c2779db21d1ee7d65d32d5250e58e7226
DIFF: https://github.com/llvm/llvm-project/commit/79afa58c2779db21d1ee7d65d32d5250e58e7226.diff

LOG: [mlir][llvm] Add fastmathFlags to SelectOp

`select` has fastmath flags, which were missing from the LLVM dialect op
definition.

Reviewed By: ftynse, dcaballe, gysit

Differential Revision: https://reviews.llvm.org/D145829

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
    mlir/test/Target/LLVMIR/Import/fastmath.ll
    mlir/test/Target/LLVMIR/llvmir.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 197a226cc478c..c2bb2f34a463a 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -771,17 +771,22 @@ def LLVM_ShuffleVectorOp : LLVM_Op<"shufflevector",
 // Misc operations.
 def LLVM_SelectOp
     : LLVM_Op<"select",
-          [Pure, AllTypesMatch<["trueValue", "falseValue", "res"]>]>,
+          [Pure, AllTypesMatch<["trueValue", "falseValue", "res"]>,
+           DeclareOpInterfaceMethods<FastmathFlagsInterface>]>,
       LLVM_Builder<
           "$res = builder.CreateSelect($condition, $trueValue, $falseValue);"> {
   let arguments = (ins LLVM_ScalarOrVectorOf<I1>:$condition,
-                   LLVM_Type:$trueValue, LLVM_Type:$falseValue);
+                   LLVM_Type:$trueValue, LLVM_Type:$falseValue,
+                   DefaultValuedAttr<LLVM_FastmathFlagsAttr,
+                                     "{}">:$fastmathFlags);
   let results = (outs LLVM_Type:$res);
   let assemblyFormat = "operands attr-dict `:` type($condition) `,` type($res)";
   string llvmInstName = "Select";
   string mlirBuilder = [{
-    $res = $_builder.create<LLVM::SelectOp>(
+    auto op = $_builder.create<LLVM::SelectOp>(
       $_location, $_resultType, $condition, $trueValue, $falseValue);
+    moduleImport.setFastmathFlagsAttr(inst, op);
+    $res = op;
   }];
 }
 def LLVM_FreezeOp : LLVM_Op<"freeze", [SameOperandsAndResultType]> {

diff  --git a/mlir/test/Target/LLVMIR/Import/fastmath.ll b/mlir/test/Target/LLVMIR/Import/fastmath.ll
index 8457c9f1c894d..9b30c3218b810 100644
--- a/mlir/test/Target/LLVMIR/Import/fastmath.ll
+++ b/mlir/test/Target/LLVMIR/Import/fastmath.ll
@@ -1,7 +1,7 @@
 ; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s
 
 ; CHECK-LABEL: @fastmath_inst
-define void @fastmath_inst(float %arg1, float %arg2) {
+define void @fastmath_inst(float %arg1, float %arg2, i1 %arg3) {
   ; CHECK: llvm.fadd %{{.*}}, %{{.*}}  {fastmathFlags = #llvm.fastmath<nnan, ninf>} : f32
   %1 = fadd nnan ninf float %arg1, %arg2
   ; CHECK: llvm.fsub %{{.*}}, %{{.*}}  {fastmathFlags = #llvm.fastmath<nsz>} : f32
@@ -12,6 +12,8 @@ define void @fastmath_inst(float %arg1, float %arg2) {
   %4 = fdiv afn reassoc float %arg1, %arg2
   ; CHECK: llvm.fneg %{{.*}}  {fastmathFlags = #llvm.fastmath<fast>} : f32
   %5 = fneg fast float %arg1
+  ; CHECK: llvm.select %{{.*}}, %{{.*}}, %{{.*}} {fastmathFlags = #llvm.fastmath<contract>} : i1, f32
+  %6 = select contract i1 %arg3, float %arg1, float %arg2
   ret void
 }
 

diff  --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index c8314fa916b05..ce65ff995709c 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -1949,6 +1949,10 @@ llvm.func @fastmathFlags(%arg0: f32, %arg1 : vector<2xf32>) {
   %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
 
+
+  %23 = llvm.mlir.constant(true) : i1
+// CHECK: select contract i1
+  %24 = llvm.select %23, %arg0, %20 {fastmathFlags = #llvm.fastmath<contract>} : i1, f32
   llvm.return
 }