[Mlir-commits] [mlir] 6f3c7df - [mlir][sparse] add sparse sign integration test

[Aart Bik]

Author: Aart Bik
Date: 2022-05-12T15:56:36-07:00
New Revision: 6f3c7dfb7746b04cc98c2f29395dec5027e8d7f3

URL: https://github.com/llvm/llvm-project/commit/6f3c7dfb7746b04cc98c2f29395dec5027e8d7f3
DIFF: https://github.com/llvm/llvm-project/commit/6f3c7dfb7746b04cc98c2f29395dec5027e8d7f3.diff

LOG: [mlir][sparse] add sparse sign integration test

Implements a floating-point sign operator (using the new semi-ring ops)
that accomodates +/-Inf and +/-NaN in consistent way.

Reviewed By: bixia

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

Added: 
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sign.mlir

Modified: 
    

Removed: 
    


################################################################################
diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sign.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sign.mlir
new file mode 100644
index 0000000000000..acce488155670
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sign.mlir
@@ -0,0 +1,100 @@
+// 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 {
+  // Performs sign operation (using semi-ring unary op)
+  // with semantics that
+  // > 0 : +1.0
+  // < 0 : -1.0
+  // +Inf: +1.0
+  // -Inf: -1.0
+  // +NaN: +NaN
+  // -NaN: -NaN
+  // +0.0: +0.0
+  // -0.0: -0.0
+  func.func @sparse_sign(%arg0: tensor<?xf64, #SparseVector>)
+                             -> tensor<?xf64, #SparseVector> {
+    %c0 = arith.constant 0 : index
+    %d = tensor.dim %arg0, %c0 : tensor<?xf64, #SparseVector>
+    %xin = sparse_tensor.init [%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 = sparse_tensor.unary %a : f64 to f64
+          present={
+            ^bb1(%s: f64):
+              %z = arith.constant 0.0 : f64
+              %1 = arith.cmpf one, %s, %z : f64
+              %2 = arith.uitofp %1 : i1 to f64
+              %3 = math.copysign %2, %s : f64
+              %4 = arith.cmpf uno, %s, %s : f64
+              %5 = arith.select %4, %s, %3 : f64
+              sparse_tensor.yield %5 : f64
+          }
+          absent={}
+        linalg.yield %result : f64
+    } -> tensor<?xf64, #SparseVector>
+    return %0 : tensor<?xf64, #SparseVector>
+  }
+
+  // Driver method to call and verify sign kernel.
+  func.func @entry() {
+    %c0 = arith.constant 0 : index
+    %du = arith.constant 99.99 : f64
+
+    %pnan = arith.constant 0x7FF0000001000000 : f64
+    %nnan = arith.constant 0xFFF0000001000000 : f64
+    %pinf = arith.constant 0x7FF0000000000000 : f64
+    %ninf = arith.constant 0xFFF0000000000000 : f64
+
+    // Setup sparse vector.
+    %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>
+    %sv1 = sparse_tensor.convert %v1
+         : tensor<32xf64> to tensor<?xf64, #SparseVector>
+
+    // Call sign kernel.
+    %0 = call @sparse_sign(%sv1) : (tensor<?xf64, #SparseVector>)
+                                 -> tensor<?xf64, #SparseVector>
+
+    //
+    // Verify the results.
+    //
+    // CHECK: ( -1, 1, -1, 1, 1, -1, nan, -nan, 1, -1, -0, 0, 99.99 )
+    //
+    %1 = sparse_tensor.values %0 : tensor<?xf64, #SparseVector> to memref<?xf64>
+    %2 = vector.transfer_read %1[%c0], %du: memref<?xf64>, vector<13xf64>
+    vector.print %2 : vector<13xf64>
+
+    // Release the resources.
+    sparse_tensor.release %sv1 : tensor<?xf64, #SparseVector>
+    sparse_tensor.release %0 : tensor<?xf64, #SparseVector>
+    return
+  }
+}
+
+