[Mlir-commits] [mlir] 07a4c4d - [mlir][math] Added arith::FastMathAttr support for math::FPowI.

[Slava Zakharin]

Author: Slava Zakharin
Date: 2022-12-13T20:47:20-08:00
New Revision: 07a4c4d601dceb14c95ab50f7bb43e39810eb278

URL: https://github.com/llvm/llvm-project/commit/07a4c4d601dceb14c95ab50f7bb43e39810eb278
DIFF: https://github.com/llvm/llvm-project/commit/07a4c4d601dceb14c95ab50f7bb43e39810eb278.diff

LOG: [mlir][math] Added arith::FastMathAttr support for math::FPowI.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Math/IR/MathOps.td
    mlir/test/Dialect/Math/ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Math/IR/MathOps.td b/mlir/include/mlir/Dialect/Math/IR/MathOps.td
index a5b28bd0891c5..3f2a8d7cb4647 100644
--- a/mlir/include/mlir/Dialect/Math/IR/MathOps.td
+++ b/mlir/include/mlir/Dialect/Math/IR/MathOps.td
@@ -857,7 +857,8 @@ def Math_TruncOp : Math_FloatUnaryOp<"trunc"> {
 //===----------------------------------------------------------------------===//
 
 def Math_FPowIOp : Math_Op<"fpowi",
-    [SameOperandsAndResultShape, AllTypesMatch<["lhs", "result"]>]> {
+    [SameOperandsAndResultShape, AllTypesMatch<["lhs", "result"]>,
+     DeclareOpInterfaceMethods<ArithFastMathInterface>]> {
   let summary = "floating point raised to the signed integer power";
   let description = [{
     Syntax:
@@ -890,9 +891,12 @@ def Math_FPowIOp : Math_Op<"fpowi",
     ```
   }];
 
-  let arguments = (ins FloatLike:$lhs, SignlessIntegerLike:$rhs);
+  let arguments = (ins FloatLike:$lhs, SignlessIntegerLike:$rhs,
+      DefaultValuedAttr<Arith_FastMathAttr,
+                        "::mlir::arith::FastMathFlags::none">:$fastmath);
   let results = (outs FloatLike:$result);
-  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs)";
+  let assemblyFormat = [{ $lhs `,` $rhs (`fastmath` `` $fastmath^)?
+                          attr-dict `:` type($lhs) `,` type($rhs) }];
 
   // TODO: add a constant folder using pow[f] for cases, when
   //       the power argument is exactly representable in floating

diff  --git a/mlir/test/Dialect/Math/ops.mlir b/mlir/test/Dialect/Math/ops.mlir
index 7e121f80dd79e..0f744f52d1c96 100644
--- a/mlir/test/Dialect/Math/ops.mlir
+++ b/mlir/test/Dialect/Math/ops.mlir
@@ -271,15 +271,18 @@ func.func @trunc(%f: f32, %v: vector<4xf32>, %t: tensor<4x4x?xf32>) {
 }
 
 // CHECK-LABEL: func @fastmath(
-// CHECK-SAME:             %[[F:.*]]: f32, %[[V:.*]]: vector<4xf32>, %[[T:.*]]: tensor<4x4x?xf32>)
-func.func @fastmath(%f: f32, %v: vector<4xf32>, %t: tensor<4x4x?xf32>) {
-  // CHECK: %{{.*}} = math.trunc %[[F]] fastmath<fast> : f32
+// CHECK-SAME:      %[[F:.*]]: f32, %[[I:.*]]: i32,
+// CHECK-SAME:      %[[V:.*]]: vector<4xf32>, %[[T:.*]]: tensor<4x4x?xf32>)
+func.func @fastmath(%f: f32, %i: i32, %v: vector<4xf32>, %t: tensor<4x4x?xf32>) {
+  // CHECK: math.trunc %[[F]] fastmath<fast> : f32
   %0 = math.trunc %f fastmath<fast> : f32
-  // CHECK: %{{.*}} = math.powf %[[V]], %[[V]] fastmath<fast> : vector<4xf32>
+  // CHECK: math.powf %[[V]], %[[V]] fastmath<fast> : vector<4xf32>
   %1 = math.powf %v, %v fastmath<reassoc,nnan,ninf,nsz,arcp,contract,afn> : vector<4xf32>
-  // CHECK: %{{.*}} = math.fma %[[T]], %[[T]], %[[T]] : tensor<4x4x?xf32>
+  // CHECK: math.fma %[[T]], %[[T]], %[[T]] : tensor<4x4x?xf32>
   %2 = math.fma %t, %t, %t fastmath<none> : tensor<4x4x?xf32>
-  // CHECK: %{{.*}} = math.absf %[[F]] fastmath<ninf> : f32
+  // CHECK: math.absf %[[F]] fastmath<ninf> : f32
   %3 = math.absf %f fastmath<ninf> : f32
+  // CHECK: math.fpowi %[[F]], %[[I]] fastmath<fast> : f32, i32
+  %4 = math.fpowi %f, %i fastmath<fast> : f32, i32
   return
 }