[Mlir-commits] [mlir] 08b4cf3 - [mlir][math] Added basic support for IPowI operation.

[Slava Zakharin]

Author: Slava Zakharin
Date: 2022-08-10T10:11:53-07:00
New Revision: 08b4cf362059bba7a38a22b571395c9801137dad

URL: https://github.com/llvm/llvm-project/commit/08b4cf362059bba7a38a22b571395c9801137dad
DIFF: https://github.com/llvm/llvm-project/commit/08b4cf362059bba7a38a22b571395c9801137dad.diff

LOG: [mlir][math] Added basic support for IPowI operation.

The operation computes pow(b, p), where 'b' and 'p' are signed integers
of the same width. The result's type matches the operands' type.

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

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 966a9a1241b1..bc9be20035c8 100644
--- a/mlir/include/mlir/Dialect/Math/IR/MathOps.td
+++ b/mlir/include/mlir/Dialect/Math/IR/MathOps.td
@@ -43,6 +43,17 @@ class Math_FloatUnaryOp<string mnemonic, list<Trait> traits = []> :
   let assemblyFormat = "$operand attr-dict `:` type($result)";
 }
 
+// Base class for binary math operations on integer types. Require two
+// operands and one result of the same type. This type can be an integer
+// type, vector or tensor thereof.
+class Math_IntegerBinaryOp<string mnemonic, list<Trait> traits = []> :
+    Math_Op<mnemonic, traits # [SameOperandsAndResultType]> {
+  let arguments = (ins SignlessIntegerLike:$lhs, SignlessIntegerLike:$rhs);
+  let results = (outs SignlessIntegerLike:$result);
+
+  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($result)";
+}
+
 // Base class for binary math operations on floating point types. Require two
 // operands and one result of the same type. This type can be a floating point
 // type, vector or tensor thereof.
@@ -503,6 +514,32 @@ def Math_FmaOp : Math_FloatTernaryOp<"fma"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// IPowIOp
+//===----------------------------------------------------------------------===//
+
+def Math_IPowIOp : Math_IntegerBinaryOp<"ipowi"> {
+  let summary = "signed integer raised to the power of operation";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `math.ipowi` ssa-use `,` ssa-use `:` type
+    ```
+
+    The `ipowi` operation takes two operands of integer type (i.e., scalar,
+    tensor or vector) and returns one result of the same type. Operands
+    must have the same type.
+
+    Example:
+
+    ```mlir
+    // Scalar signed integer exponentiation.
+    %a = math.ipowi %b, %c : i32
+    ```
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // LogOp
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/Math/ops.mlir b/mlir/test/Dialect/Math/ops.mlir
index bcf085e54102..a2b959a783f2 100644
--- a/mlir/test/Dialect/Math/ops.mlir
+++ b/mlir/test/Dialect/Math/ops.mlir
@@ -218,3 +218,15 @@ func.func @round(%f: f32, %v: vector<4xf32>, %t: tensor<4x4x?xf32>) {
   %2 = math.round %t : tensor<4x4x?xf32>
   return
 }
+
+// CHECK-LABEL: func @ipowi(
+// CHECK-SAME:             %[[I:.*]]: i32, %[[V:.*]]: vector<4xi32>, %[[T:.*]]: tensor<4x4x?xi32>)
+func.func @ipowi(%i: i32, %v: vector<4xi32>, %t: tensor<4x4x?xi32>) {
+  // CHECK: %{{.*}} = math.ipowi %[[I]], %[[I]] : i32
+  %0 = math.ipowi %i, %i : i32
+  // CHECK: %{{.*}} = math.ipowi %[[V]], %[[V]] : vector<4xi32>
+  %1 = math.ipowi %v, %v : vector<4xi32>
+  // CHECK: %{{.*}} = math.ipowi %[[T]], %[[T]] : tensor<4x4x?xi32>
+  %2 = math.ipowi %t, %t : tensor<4x4x?xi32>
+  return
+}