[mlir][sparse]Add more integration tests for sparse_tensor.unary

[Peiming]

Summary

Previously, the sparse_tensor.unary integration test does not contain cases with the use of linalg.index (previoulsy unsupported), this commit adds test cases that use linalg.index operators.

[aartbik]

Nice direction. One last nit.

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

126–127

This is a bit misleading kernel. You say that we only "keep" a certain part, but what you really do is output zeros for that part. In the sparse format, you can see that there are a lot of explicitly stored zeros.

To make this a more realistic example, perhaps say that we keep values in the corner, and set all others to some value (nonzero), e.g. 99 or so. Just to avoid the misleading "only keep"

[Peiming]

A less misleading kernel

[aartbik]

Nice work, Peiming! Ship it!

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

%0 = linalg.generic #trait_vec_scale
             absent={
               sparse_tensor.yield %cf1 : f64
             }
           linalg.yield %1 : f64
     } -> tensor<?xf64, #SparseVector>
     return %0 : tensor<?xf64, #SparseVector>
   }
 
+  // Performs B[i] = i * A[i].
+  func.func @vector_magnify(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
+    %c = arith.constant 0 : index
+    %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
+    %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
+    %0 = linalg.generic #trait_vec_scale
+       ins(%arga: tensor<?xf64, #SparseVector>)
+        outs(%xv: tensor<?xf64, #SparseVector>) {
+        ^bb(%a: f64, %x: f64):
+          %idx = linalg.index 0 : index
+          %1 = sparse_tensor.unary %a : f64 to f64
+            present={
+              ^bb0(%x0: f64):
+                %tmp = arith.index_cast %idx : index to i64
+                %idxf = arith.sitofp %tmp : i64 to f64
+                %ret = arith.mulf %x0, %idxf : f64
+                sparse_tensor.yield %ret : f64
+            }
+            absent={} 
+          linalg.yield %1 : f64
+    } -> tensor<?xf64, #SparseVector>
+    return %0 : tensor<?xf64, #SparseVector>
+  }
+
   // Clips values to the range [3, 7].
   func.func @matrix_clip(%argx: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
     %cfmin = arith.constant 3.0 : f64
     %cfmax = arith.constant 7.0 : f64
     %d0 = tensor.dim %argx, %c0 : tensor<?x?xf64, #DCSR>
     %d1 = tensor.dim %argx, %c1 : tensor<?x?xf64, #DCSR>
%0 = linalg.generic #trait_mat_scale
                 sparse_tensor.yield %x2 : f64
             }
             absent={}
           linalg.yield %1 : f64
     } -> tensor<?x?xf64, #DCSR>
     return %0 : tensor<?x?xf64, #DCSR>
   }
 
+  // Slices matrix and only keeps the lower-right corner of the original matrix
+  // (i.e., A[2/d0 ..][2/d1 ..]).
+  func.func @matrix_slice(%argx: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %d0 = tensor.dim %argx, %c0 : tensor<?x?xf64, #DCSR>
+    %d1 = tensor.dim %argx, %c1 : tensor<?x?xf64, #DCSR>
+    %xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR>
+    %0 = linalg.generic #trait_mat_scale
+       ins(%argx: tensor<?x?xf64, #DCSR>)
+        outs(%xv: tensor<?x?xf64, #DCSR>) {
+        ^bb(%a: f64, %x: f64):
+          %row = linalg.index 0 : index
+          %col = linalg.index 1 : index
+          %1 = sparse_tensor.unary %a: f64 to f64
+            present={
+              ^bb0(%x0: f64):
+                %zero = arith.constant 0.0 : f64
+                %two = arith.constant 2 : index
+                %r = arith.muli %two, %row : index
+                %c = arith.muli %two, %col : index
+                %cmp1 = arith.cmpi "ult", %r, %d0 : index
+                %tmp = arith.select %cmp1, %zero, %x0 : f64
+                %cmp2 = arith.cmpi "ult", %c, %d1 : index
+                %result = arith.select %cmp2, %zero, %tmp : f64
+                sparse_tensor.yield %result : f64
+            }
+            absent={}
+          linalg.yield %1 : f64
+    } -> tensor<?x?xf64, #DCSR>
+    return %0 : tensor<?x?xf64, #DCSR>
+  }
+
   // Dumps a sparse vector of type f64.
   func.func @dump_vec_f64(%arg0: tensor<?xf64, #SparseVector>) {
     // Dump the values array to verify only sparse contents are stored.
     %c0 = arith.constant 0 : index
     %d0 = arith.constant -1.0 : f64
     %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf64>
     %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<32xf64>
     vector.print %1 : vector<32xf64>
func.func @entry() {
     > : tensor<4x8xf64>
     %sm1 = sparse_tensor.convert %m1 : tensor<4x8xf64> to tensor<?x?xf64, #DCSR>
 
     // Call sparse vector kernels.
     %0 = call @vector_complement(%sv1)
        : (tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector>
     %1 = call @vector_negation(%sv1)
        : (tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector>
+    %2 = call @vector_magnify(%sv1)
+       : (tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector>
     
 
     // Call sparse matrix kernels.
-    %2 = call @matrix_clip(%sm1)
+    %3 = call @matrix_clip(%sm1)
+      : (tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR>
+    %4 = call @matrix_slice(%sm1)
       : (tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR>
 
     //
     // Verify the results.
     //
     // CHECK:      ( 1, 2, 3, 4, 5, 6, 7, 8, 9, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 )
     // CHECK-NEXT: ( 1, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 4, 0, 0, 5, 6, 0, 0, 0, 0, 0, 0, 7, 8, 0, 9 )
     // CHECK-NEXT: ( 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, -1 )
     // CHECK-NEXT: ( 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 0 )
     // CHECK-NEXT: ( -1, 1, 1, -2, 1, 1, 1, 1, 1, 1, 1, -3, 1, 1, 1, 1, 1, -4, 1, 1, -5, -6, 1, 1, 1, 1, 1, 1, -7, -8, 1, -9 )
     // CHECK-NEXT: ( -1, 1, 1, -2, 1, 1, 1, 1, 1, 1, 1, -3, 1, 1, 1, 1, 1, -4, 1, 1, -5, -6, 1, 1, 1, 1, 1, 1, -7, -8, 1, -9 )
+    // CHECK-NEXT: ( 0, 6, 33, 68, 100, 126, 196, 232, 279, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 33, 0, 0, 0, 0, 0, 68, 0, 0, 100, 126, 0, 0, 0, 0, 0, 0, 196, 232, 0, 279 )
     // CHECK-NEXT: ( 3, 3, 3, 4, 5, 6, 7, 7, 7, -1, -1, -1, -1, -1, -1, -1 )
     // CHECK-NEXT: ( ( 3, 3, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 3 ), ( 0, 0, 4, 0, 5, 0, 0, 6 ), ( 7, 0, 7, 7, 0, 0, 0, 0 ) )
+    // CHECK-NEXT: ( 0, 0, 0, 0, 5, 6, 0, 0, 0, -1, -1, -1, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( ( 0, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 5, 0, 0, 6 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ) )
     //
     call @dump_vec_f64(%sv1) : (tensor<?xf64, #SparseVector>) -> ()
     call @dump_vec_i32(%0) : (tensor<?xi32, #SparseVector>) -> ()
     call @dump_vec_f64(%1) : (tensor<?xf64, #SparseVector>) -> ()
-    call @dump_mat(%2) : (tensor<?x?xf64, #DCSR>) -> ()
+    call @dump_vec_f64(%2) : (tensor<?xf64, #SparseVector>) -> ()
+    call @dump_mat(%3) : (tensor<?x?xf64, #DCSR>) -> ()
+    call @dump_mat(%4) : (tensor<?x?xf64, #DCSR>) -> ()
     
     // Release the resources.
     sparse_tensor.release %sv1 : tensor<?xf64, #SparseVector>
     sparse_tensor.release %sm1 : tensor<?x?xf64, #DCSR>
     sparse_tensor.release %0 : tensor<?xi32, #SparseVector>
     sparse_tensor.release %1 : tensor<?xf64, #SparseVector>
-    sparse_tensor.release %2 : tensor<?x?xf64, #DCSR>
+    sparse_tensor.release %2 : tensor<?xf64, #SparseVector>
+    sparse_tensor.release %3 : tensor<?x?xf64, #DCSR>
+    sparse_tensor.release %4 : tensor<?x?xf64, #DCSR>
     return
   }
 }