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

[Mogball]

Summary

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

[aartbik]

mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir

13

Did you test this end-to-end? The actual intrinsic takes an extra parameter

declare i32 @llvm.abs.i32(i32 <src>, i1 <is_int_min_poison>)

so I am not sure this will actually execute properly when using the execution engine on the resulting LLVM IR?

[aartbik]

mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir

13

This end-to-end test cannot run yet due to this issue:

https://reviews.llvm.org/D131712

[Mogball]

mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir

13

Good catch! Let me submit a fix.

[aartbik]

mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir

13

Thanks!

[Mogball]

mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir

13

https://reviews.llvm.org/D131785

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

 }
 
 class LLVM_CountZerosIntrinsicOp<string func, list<Trait> traits = []> :
     LLVM_OneResultIntrOp<func, [], [0],
            !listconcat([NoSideEffect], 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">;
 def LLVM_Exp2Op : LLVM_UnaryIntrinsicOp<"exp2">;
 def LLVM_FAbsOp : LLVM_UnaryIntrinsicOp<"fabs">;
 def LLVM_FCeilOp : LLVM_UnaryIntrinsicOp<"ceil">;
 def LLVM_FFloorOp : LLVM_UnaryIntrinsicOp<"floor">;
 def LLVM_FMAOp : LLVM_TernarySameArgsIntrinsicOp<"fma">;

mlir/include/mlir/Dialect/Math/IR/MathOps.td

let description = [{
     // Scalar absolute value.
     %a = math.absf %b : f64
     ```
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// 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
 //===----------------------------------------------------------------------===//
 
 def Math_AtanOp : Math_FloatUnaryOp<"atan">{
   let summary = "arcus tangent of the given value";
   let description = [{
     Syntax:
 

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/IR/TypeUtilities.h"
 
 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>;
 using CosOpLowering = VectorConvertToLLVMPattern<math::CosOp, LLVM::CosOp>;
 using CtPopFOpLowering =
     VectorConvertToLLVMPattern<math::CtPopOp, LLVM::CtPopOp>;
 using Exp2OpLowering = VectorConvertToLLVMPattern<math::Exp2Op, LLVM::Exp2Op>;
 using ExpOpLowering = VectorConvertToLLVMPattern<math::ExpOp, LLVM::ExpOp>;

 };
 } // namespace
 
 void mlir::populateMathToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                 RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
     AbsFOpLowering,
+    AbsIOpLowering,
     CeilOpLowering,
     CopySignOpLowering,
     CosOpLowering,
     CountLeadingZerosOpLowering,
     CountTrailingZerosOpLowering,
     CtPopFOpLowering,
     Exp2OpLowering,
     ExpM1OpLowering,

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

void populateMathToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
   // Core patterns
   patterns.add<CopySignPattern>(typeConverter, patterns.getContext());
 
   // GLSL patterns
   patterns
       .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>,
            spirv::ElementwiseOpPattern<math::FloorOp, spirv::GLFloorOp>,
            spirv::ElementwiseOpPattern<math::FmaOp, spirv::GLFmaOp>,
            spirv::ElementwiseOpPattern<math::LogOp, spirv::GLLogOp>,
            spirv::ElementwiseOpPattern<math::RsqrtOp, spirv::GLInverseSqrtOp>,
            spirv::ElementwiseOpPattern<math::SinOp, spirv::GLSinOp>,

mlir/lib/Dialect/Math/IR/MathOps.cpp

 //===----------------------------------------------------------------------===//
 
 OpFoldResult math::AbsFOp::fold(ArrayRef<Attribute> operands) {
   return constFoldUnaryOp<FloatAttr>(operands,
                                      [](const APFloat &a) { return abs(a); });
 }
 
 //===----------------------------------------------------------------------===//
+// AbsIOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::AbsIOp::fold(ArrayRef<Attribute> operands) {
+  return constFoldUnaryOp<IntegerAttr>(operands,
+                                       [](const APInt &a) { return a.abs(); });
+}
+
+//===----------------------------------------------------------------------===//
 // AtanOp folder
 //===----------------------------------------------------------------------===//
 
 OpFoldResult math::AtanOp::fold(ArrayRef<Attribute> operands) {
   return constFoldUnaryOpConditional<FloatAttr>(
       operands, [](const APFloat &a) -> Optional<APFloat> {
         switch (a.getSizeInBits(a.getSemantics())) {
         case 64:

mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir

 // RUN: mlir-opt %s -split-input-file -pass-pipeline="func.func(convert-math-to-llvm)" | FileCheck %s
 
 // 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
+  %0 = math.exp %arg0 : f32
   // CHECK: = "llvm.intr.exp2"(%{{.*}}) : (f32) -> f32
-  %14 = math.exp2 %arg0 : f32
+  %1 = math.exp2 %arg0 : f32
   // CHECK: = "llvm.intr.sqrt"(%{{.*}}) : (f32) -> f32
-  %19 = math.sqrt %arg0 : f32
+  %2 = math.sqrt %arg0 : f32
   // CHECK: = "llvm.intr.sqrt"(%{{.*}}) : (f64) -> f64
-  %20 = math.sqrt %arg4 : f64
+  %3 = math.sqrt %arg4 : f64
+  // CHECK: = "llvm.intr.abs"(%{{.*}}) : (i32) -> i32
+  %4 = math.absi %arg2 : i32
   func.return
 }
 
 // -----
 
 // CHECK-LABEL: func @log1p(
 // CHECK-SAME: f32
 func.func @log1p(%arg0 : f32) {

mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir

 // CHECK-LABEL: @float32_ternary_vector
 func.func @float32_ternary_vector(%a: vector<4xf32>, %b: vector<4xf32>,
                             %c: vector<4xf32>) {
   // CHECK: spv.GL.Fma %{{.*}}: vector<4xf32>
   %0 = math.fma %a, %b, %c : 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 {
   // CHECK-DAG: %[[V1:.+]] = spv.Constant 1 : i32
   // CHECK-DAG: %[[V31:.+]] = spv.Constant 31 : i32
   // CHECK-DAG: %[[V32:.+]] = spv.Constant 32 : i32
   // CHECK: %[[MSB:.+]] = spv.GL.FindUMsb %[[VAL]] : i32
   // CHECK: %[[SUB1:.+]] = spv.ISub %[[V31]], %[[MSB]] : i32