[mlir][sparse] enable integral abs recognition

The end-to-end test for this new feature also exposed a bug
in LLVM IR lowering (since then, fixed), where we need to account
for the min-poison bit as extra argument.

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

Reviewed By: bixia

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

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

@@ -32,6 +32,7 @@ enum Kind {
   // Unary operations.
   kAbsF,
   kAbsC,
+  kAbsI,
   kCeilF,
   kFloorF,
   kSqrtF,

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

@@ -40,6 +40,7 @@ TensorExp::TensorExp(Kind k, unsigned x, unsigned y, Value v, Operation *o)
   // Unary operations.
   case kAbsF:
   case kAbsC:
+  case kAbsI:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
@@ -310,6 +311,7 @@ bool Merger::isSingleCondition(unsigned t, unsigned e) const {
   // Unary operations.
   case kAbsF:
   case kAbsC:
+  case kAbsI:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
@@ -398,6 +400,7 @@ static const char *kindToOpSymbol(Kind kind) {
   // Unary operations.
   case kAbsF:
   case kAbsC:
+  case kAbsI:
     return "abs";
   case kCeilF:
     return "ceil";
@@ -497,6 +500,7 @@ void Merger::dumpExp(unsigned e) const {
   // Unary operations.
   case kAbsF:
   case kAbsC:
+  case kAbsI:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
@@ -630,6 +634,7 @@ unsigned Merger::buildLattices(unsigned e, unsigned i) {
   // Unary operations.
   case kAbsF:
   case kAbsC:
+  case kAbsI:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
@@ -896,6 +901,8 @@ Optional<unsigned> Merger::buildTensorExp(linalg::GenericOp op, Value v) {
         return addExp(kAbsF, e);
       if (isa<complex::AbsOp>(def))
         return addExp(kAbsC, e);
+      if (isa<math::AbsIOp>(def))
+        return addExp(kAbsI, e);
       if (isa<math::CeilOp>(def))
         return addExp(kCeilF, e);
       if (isa<math::FloorOp>(def))
@@ -1079,6 +1086,8 @@ Value Merger::buildExp(RewriterBase &rewriter, Location loc, unsigned e,
     auto eltType = type.getElementType().cast<FloatType>();
     return rewriter.create<complex::AbsOp>(loc, eltType, v0);
   }
+  case kAbsI:
+    return rewriter.create<math::AbsIOp>(loc, v0);
   case kCeilF:
     return rewriter.create<math::CeilOp>(loc, v0);
   case kFloorF:

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_abs.mlir

@@ -0,0 +1,110 @@
+// RUN: mlir-opt %s --sparse-compiler | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+#SparseVector = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
+
+#trait_op = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>, // a
+    affine_map<(i) -> (i)>  // x (out)
+  ],
+  iterator_types = ["parallel"],
+  doc = "x(i) = OP a(i)"
+}
+
+module {
+  func.func @sparse_absf(%arg0: tensor<?xf64, #SparseVector>)
+                             -> tensor<?xf64, #SparseVector> {
+    %c0 = arith.constant 0 : index
+    %d = tensor.dim %arg0, %c0 : tensor<?xf64, #SparseVector>
+    %xin = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
+    %0 = linalg.generic #trait_op
+      ins(%arg0: tensor<?xf64, #SparseVector>)
+      outs(%xin: tensor<?xf64, #SparseVector>) {
+      ^bb0(%a: f64, %x: f64) :
+        %result = math.absf %a : f64
+        linalg.yield %result : f64
+    } -> tensor<?xf64, #SparseVector>
+    return %0 : tensor<?xf64, #SparseVector>
+  }
+
+  func.func @sparse_absi(%arg0: tensor<?xi32, #SparseVector>)
+                             -> tensor<?xi32, #SparseVector> {
+    %c0 = arith.constant 0 : index
+    %d = tensor.dim %arg0, %c0 : tensor<?xi32, #SparseVector>
+    %xin = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector>
+    %0 = linalg.generic #trait_op
+      ins(%arg0: tensor<?xi32, #SparseVector>)
+      outs(%xin: tensor<?xi32, #SparseVector>) {
+      ^bb0(%a: i32, %x: i32) :
+        %result = math.absi %a : i32
+        linalg.yield %result : i32
+    } -> tensor<?xi32, #SparseVector>
+    return %0 : tensor<?xi32, #SparseVector>
+  }
+
+  // Driver method to call and verify sign kernel.
+  func.func @entry() {
+    %c0 = arith.constant 0 : index
+    %df = arith.constant 99.99 : f64
+    %di = arith.constant 9999 : i32
+
+    %pnan = arith.constant 0x7FF0000001000000 : f64
+    %nnan = arith.constant 0xFFF0000001000000 : f64
+    %pinf = arith.constant 0x7FF0000000000000 : f64
+    %ninf = arith.constant 0xFFF0000000000000 : f64
+
+    // Setup sparse vectors.
+    %v1 = arith.constant sparse<
+       [ [0], [3], [5], [11], [13], [17], [18], [20], [21], [28], [29], [31] ],
+         [ -1.5, 1.5, -10.2, 11.3, 1.0, -1.0,
+           0x7FF0000001000000, // +NaN
+           0xFFF0000001000000, // -NaN
+           0x7FF0000000000000, // +Inf
+           0xFFF0000000000000, // -Inf
+           -0.0,               // -Zero
+           0.0                 // +Zero
+        ]
+    > : tensor<32xf64>
+    %v2 = arith.constant sparse<
+       [ [0], [3], [5], [11], [13], [17], [18], [21], [31] ],
+         [ -2147483648, -2147483647, -1000, -1, 0,
+	   1, 1000, 2147483646, 2147483647
+	 ]
+    > : tensor<32xi32>
+    %sv1 = sparse_tensor.convert %v1
+         : tensor<32xf64> to tensor<?xf64, #SparseVector>
+    %sv2 = sparse_tensor.convert %v2
+         : tensor<32xi32> to tensor<?xi32, #SparseVector>
+
+    // Call abs kernels.
+    %0 = call @sparse_absf(%sv1) : (tensor<?xf64, #SparseVector>)
+                                 -> tensor<?xf64, #SparseVector>
+
+    %1 = call @sparse_absi(%sv2) : (tensor<?xi32, #SparseVector>)
+                                 -> tensor<?xi32, #SparseVector>
+
+    //
+    // Verify the results.
+    //
+    // CHECK:      ( 1.5, 1.5, 10.2, 11.3, 1, 1, nan, nan, inf, inf, 0, 0, 99.99 )
+    // CHECK-NEXT: ( -2147483648, 2147483647, 1000, 1, 0, 1, 1000, 2147483646, 2147483647, 9999, 9999, 9999, 9999 )
+    //
+    %x = sparse_tensor.values %0 : tensor<?xf64, #SparseVector> to memref<?xf64>
+    %y = sparse_tensor.values %1 : tensor<?xi32, #SparseVector> to memref<?xi32>
+    %a = vector.transfer_read %x[%c0], %df: memref<?xf64>, vector<13xf64>
+    %b = vector.transfer_read %y[%c0], %di: memref<?xi32>, vector<13xi32>
+    vector.print %a : vector<13xf64>
+    vector.print %b : vector<13xi32>
+
+    // Release the resources.
+    bufferization.dealloc_tensor %sv1 : tensor<?xf64, #SparseVector>
+    bufferization.dealloc_tensor %sv2 : tensor<?xi32, #SparseVector>
+    bufferization.dealloc_tensor %0 : tensor<?xf64, #SparseVector>
+    bufferization.dealloc_tensor %1 : tensor<?xi32, #SparseVector>
+    return
+  }
+}

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

@@ -230,6 +230,7 @@ protected:
     // Unary operations.
     case kAbsF:
     case kAbsC:
+    case kAbsI:
     case kCeilF:
     case kFloorF:
     case kSqrtF: