[mlir][math] Add `math.absi` op

Adds an integer absolute value op to the math dialect.

When converting to LLVM, this op is lowered to the LLVM `abs` intrinsic.
When converting to SPIRV, this op is lowered to `spv.GL.SAbs`.

Depends on D131325

Reviewed By: ftynse

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

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index c7e947b..937a2fa 100644 (file)
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
@@ -40,6 +40,7 @@ class LLVM_CountZerosIntrinsicOp<string func, list<Trait> traits = []> :
   let arguments = (ins LLVM_Type:$in, I<1>:$zero_undefined);
 }
 
+def LLVM_AbsOp : LLVM_UnaryIntrinsicOp<"abs">;
 def LLVM_CopySignOp : LLVM_BinarySameArgsIntrinsicOp<"copysign">;
 def LLVM_CosOp : LLVM_UnaryIntrinsicOp<"cos">;
 def LLVM_ExpOp : LLVM_UnaryIntrinsicOp<"exp">;

diff --git a/mlir/include/mlir/Dialect/Math/IR/MathOps.td b/mlir/include/mlir/Dialect/Math/IR/MathOps.td
index 27ab3a3..966a9a1 100644 (file)
--- a/mlir/include/mlir/Dialect/Math/IR/MathOps.td
+++ b/mlir/include/mlir/Dialect/Math/IR/MathOps.td
@@ -87,6 +87,27 @@ def Math_AbsFOp : Math_FloatUnaryOp<"absf"> {
 }
 
 //===----------------------------------------------------------------------===//
+// AbsIOp
+//===----------------------------------------------------------------------===//
+
+def Math_AbsIOp : Math_IntegerUnaryOp<"absi"> {
+  let summary = "integer absolute-value operation";
+  let description = [{
+    The `absi` operation computes the absolute value. It takes one operand of
+    integer type (i.e., scalar, tensor or vector) and returns one result of the
+    same type.
+
+    Example:
+
+    ```mlir
+    // Scalar absolute value.
+    %a = math.absi %b : i64
+    ```
+  }];
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // AtanOp
 //===----------------------------------------------------------------------===//
 

diff --git a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
index 1ac4bbd..77fc91d 100644 (file)
--- a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
+++ b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
@@ -19,6 +19,7 @@ using namespace mlir;
 
 namespace {
 using AbsFOpLowering = VectorConvertToLLVMPattern<math::AbsFOp, LLVM::FAbsOp>;
+using AbsIOpLowering = VectorConvertToLLVMPattern<math::AbsIOp, LLVM::AbsOp>;
 using CeilOpLowering = VectorConvertToLLVMPattern<math::CeilOp, LLVM::FCeilOp>;
 using CopySignOpLowering =
     VectorConvertToLLVMPattern<math::CopySignOp, LLVM::CopySignOp>;
@@ -269,6 +270,7 @@ void mlir::populateMathToLLVMConversionPatterns(LLVMTypeConverter &converter,
   // clang-format off
   patterns.add<
     AbsFOpLowering,
+    AbsIOpLowering,
     CeilOpLowering,
     CopySignOpLowering,
     CosOpLowering,

diff --git a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
index e53bd2b..55a2429 100644 (file)
--- a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
+++ b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
@@ -288,6 +288,7 @@ void populateMathToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
       .add<CountLeadingZerosPattern, Log1pOpPattern<spirv::GLLogOp>,
            ExpM1OpPattern<spirv::GLExpOp>, PowFOpPattern, RoundOpPattern,
            spirv::ElementwiseOpPattern<math::AbsFOp, spirv::GLFAbsOp>,
+           spirv::ElementwiseOpPattern<math::AbsIOp, spirv::GLSAbsOp>,
            spirv::ElementwiseOpPattern<math::CeilOp, spirv::GLCeilOp>,
            spirv::ElementwiseOpPattern<math::CosOp, spirv::GLCosOp>,
            spirv::ElementwiseOpPattern<math::ExpOp, spirv::GLExpOp>,

diff --git a/mlir/lib/Dialect/Math/IR/MathOps.cpp b/mlir/lib/Dialect/Math/IR/MathOps.cpp
index 50d6512..e128435 100644 (file)
--- a/mlir/lib/Dialect/Math/IR/MathOps.cpp
+++ b/mlir/lib/Dialect/Math/IR/MathOps.cpp
@@ -31,6 +31,15 @@ OpFoldResult math::AbsFOp::fold(ArrayRef<Attribute> operands) {
 }
 
 //===----------------------------------------------------------------------===//
+// AbsIOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::AbsIOp::fold(ArrayRef<Attribute> operands) {
+  return constFoldUnaryOp<IntegerAttr>(operands,
+                                       [](const APInt &a) { return a.abs(); });
+}
+
+//===----------------------------------------------------------------------===//
 // AtanOp folder
 //===----------------------------------------------------------------------===//
 

diff --git a/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir b/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
index 6378ea6..fc91a55 100644 (file)
--- a/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
+++ b/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
@@ -2,14 +2,16 @@
 
 // CHECK-LABEL: @ops
 func.func @ops(%arg0: f32, %arg1: f32, %arg2: i32, %arg3: i32, %arg4: f64) {
-// CHECK: = "llvm.intr.exp"(%{{.*}}) : (f32) -> f32
-  %13 = math.exp %arg0 : f32
-// CHECK: = "llvm.intr.exp2"(%{{.*}}) : (f32) -> f32
-  %14 = math.exp2 %arg0 : f32
-// CHECK: = "llvm.intr.sqrt"(%{{.*}}) : (f32) -> f32
-  %19 = math.sqrt %arg0 : f32
-// CHECK: = "llvm.intr.sqrt"(%{{.*}}) : (f64) -> f64
-  %20 = math.sqrt %arg4 : f64
+  // CHECK: = "llvm.intr.exp"(%{{.*}}) : (f32) -> f32
+  %0 = math.exp %arg0 : f32
+  // CHECK: = "llvm.intr.exp2"(%{{.*}}) : (f32) -> f32
+  %1 = math.exp2 %arg0 : f32
+  // CHECK: = "llvm.intr.sqrt"(%{{.*}}) : (f32) -> f32
+  %2 = math.sqrt %arg0 : f32
+  // CHECK: = "llvm.intr.sqrt"(%{{.*}}) : (f64) -> f64
+  %3 = math.sqrt %arg4 : f64
+  // CHECK: = "llvm.intr.abs"(%{{.*}}) : (i32) -> i32
+  %4 = math.absi %arg2 : i32
   func.return
 }
 

diff --git a/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir b/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
index 1302b47..662478a 100644 (file)
--- a/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
+++ b/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
@@ -79,6 +79,13 @@ func.func @float32_ternary_vector(%a: vector<4xf32>, %b: vector<4xf32>,
   return
 }
 
+// CHECK-LABEL: @int_unary
+func.func @int_unary(%arg0: i32) {
+  // CHECK: spv.GL.SAbs %{{.*}}
+  %0 = math.absi %arg0 : i32
+  return
+}
+
 // CHECK-LABEL: @ctlz_scalar
 //  CHECK-SAME: (%[[VAL:.+]]: i32)
 func.func @ctlz_scalar(%val: i32) -> i32 {