[mlir][sparse] enable integral abs recognition

[aartbik]

Summary

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>)

[aartbik]

rebased with main, enabled the test for absi

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

 enum Kind {
   // Leaf.
   kTensor = 0,
   kInvariant,
   kIndex,
   // Unary operations.
   kAbsF,
   kAbsC,
+  kAbsI,
   kCeilF,
   kFloorF,
   kSqrtF,
   kSqrtC,
   kExpm1F,
   kExpm1C,
   kLog1pF,
   kLog1pC,

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

case kInvariant:
     break;
   case kIndex:
     assert(x != -1u && y == -1u && !v && !o);
     index = x;
     break;
   // Unary operations.
   case kAbsF:
   case kAbsC:
+  case kAbsI:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
   case kSqrtC:
   case kExpm1F:
   case kExpm1C:
   case kLog1pF:
   case kLog1pC:
bool Merger::isSingleCondition(unsigned t, unsigned e) const {
   case kTensor:
     return tensorExps[e].tensor == t;
   case kInvariant:
   case kIndex:
     return false;
   // Unary operations.
   case kAbsF:
   case kAbsC:
+  case kAbsI:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
   case kSqrtC:
   case kExpm1F:
   case kExpm1C:
   case kLog1pF:
   case kLog1pC:
case kTensor:
     return "tensor";
   case kInvariant:
     return "invariant";
   case kIndex:
     return "index";
   // Unary operations.
   case kAbsF:
   case kAbsC:
+  case kAbsI:
     return "abs";
   case kCeilF:
     return "ceil";
   case kFloorF:
     return "floor";
   case kSqrtF:
   case kSqrtC:
     return "sqrt";
case kInvariant:
     llvm::dbgs() << "invariant";
     break;
   case kIndex:
     llvm::dbgs() << "index_" << tensorExps[e].index;
     break;
   // Unary operations.
   case kAbsF:
   case kAbsC:
+  case kAbsI:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
   case kSqrtC:
   case kExpm1F:
   case kExpm1C:
   case kLog1pF:
   case kLog1pC:
if (kind == kTensor) {
         t = syntheticTensor;
     }
     latSets[s].push_back(addLat(t, i, e));
     return s;
   }
   // Unary operations.
   case kAbsF:
   case kAbsC:
+  case kAbsI:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
   case kSqrtC:
   case kExpm1F:
   case kExpm1C:
   case kLog1pF:
   case kLog1pC:
Optional<unsigned> Merger::buildTensorExp(linalg::GenericOp op, Value v) {
   if (def->getNumOperands() == 1) {
     auto x = buildTensorExp(op, def->getOperand(0));
     if (x.has_value()) {
       unsigned e = x.value();
       if (isa<math::AbsFOp>(def))
         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))
         return addExp(kFloorF, e);
       if (isa<math::SqrtOp>(def))
         return addExp(kSqrtF, e);
       if (isa<complex::SqrtOp>(def))
         return addExp(kSqrtC, e);
Value Merger::buildExp(RewriterBase &rewriter, Location loc, unsigned e,
   // Unary operations.
   case kAbsF:
     return rewriter.create<math::AbsFOp>(loc, v0);
   case kAbsC: {
     auto type = v0.getType().cast<ComplexType>();
     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:
     return rewriter.create<math::FloorOp>(loc, v0);
   case kSqrtF:
     return rewriter.create<math::SqrtOp>(loc, v0);
   case kSqrtC:
     return rewriter.create<complex::SqrtOp>(loc, v0);

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

+// 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

bool compareExpression(unsigned e, const std::shared_ptr<Pattern> &pattern) {
     case kTensor:
       return tensorExp.tensor == pattern->tensorNum;
     case kInvariant:
     case kIndex:
       llvm_unreachable("invariant not handled yet");
     // Unary operations.
     case kAbsF:
     case kAbsC:
+    case kAbsI:
     case kCeilF:
     case kFloorF:
     case kSqrtF:
     case kSqrtC:
     case kExpm1F:
     case kExpm1C:
     case kLog1pF:
     case kLog1pC: