[Mlir-commits] [mlir] 8dd07e3 - [mlir][sparse] enable integral abs recognition
Aart Bik
llvmlistbot at llvm.org
Fri Aug 12 11:36:49 PDT 2022
Author: Aart Bik
Date: 2022-08-12T11:36:40-07:00
New Revision: 8dd07e36ca0f56bad6ca910b3a0ffe5ed38bc34a
URL: https://github.com/llvm/llvm-project/commit/8dd07e36ca0f56bad6ca910b3a0ffe5ed38bc34a
DIFF: https://github.com/llvm/llvm-project/commit/8dd07e36ca0f56bad6ca910b3a0ffe5ed38bc34a.diff
LOG: [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
Added:
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_abs.mlir
Modified:
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
mlir/unittests/Dialect/SparseTensor/MergerTest.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h b/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
index 86d36b1bca218..0e3f5d0c03080 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
@@ -32,6 +32,7 @@ enum Kind {
// Unary operations.
kAbsF,
kAbsC,
+ kAbsI,
kCeilF,
kFloorF,
kSqrtF,
diff --git a/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp b/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
index b55388bb37698..e691cb44e2656 100644
--- a/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
+++ b/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:
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_abs.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_abs.mlir
new file mode 100644
index 0000000000000..6927b91d9a7f8
--- /dev/null
+++ b/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
+ }
+}
diff --git a/mlir/unittests/Dialect/SparseTensor/MergerTest.cpp b/mlir/unittests/Dialect/SparseTensor/MergerTest.cpp
index a44f1f67054c2..91101c4d77db5 100644
--- a/mlir/unittests/Dialect/SparseTensor/MergerTest.cpp
+++ b/mlir/unittests/Dialect/SparseTensor/MergerTest.cpp
@@ -230,6 +230,7 @@ class MergerTestBase : public ::testing::Test {
// Unary operations.
case kAbsF:
case kAbsC:
+ case kAbsI:
case kCeilF:
case kFloorF:
case kSqrtF:
More information about the Mlir-commits
mailing list