1// RUN: mlir-opt %s --sparse-compiler | \
2// RUN: mlir-cpu-runner \
3// RUN:  -e entry -entry-point-result=void  \
4// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
5// RUN: FileCheck %s
6
7#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
8#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
9
10//
11// Traits for tensor operations.
12//
13#trait_vec_scale = {
14  indexing_maps = [
15    affine_map<(i) -> (i)>,  // a (in)
16    affine_map<(i) -> (i)>   // x (out)
17  ],
18  iterator_types = ["parallel"]
19}
20#trait_mat_scale = {
21  indexing_maps = [
22    affine_map<(i,j) -> (i,j)>,  // A (in)
23    affine_map<(i,j) -> (i,j)>   // X (out)
24  ],
25  iterator_types = ["parallel", "parallel"]
26}
27
28module {
29  // Invert the structure of a sparse vector. Present values become missing.
30  // Missing values are filled with 1 (i32).
31  func.func @vector_complement(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector> {
32    %c = arith.constant 0 : index
33    %ci1 = arith.constant 1 : i32
34    %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
35    %xv = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector>
36    %0 = linalg.generic #trait_vec_scale
37       ins(%arga: tensor<?xf64, #SparseVector>)
38        outs(%xv: tensor<?xi32, #SparseVector>) {
39        ^bb(%a: f64, %x: i32):
40          %1 = sparse_tensor.unary %a : f64 to i32
41            present={}
42            absent={
43              sparse_tensor.yield %ci1 : i32
44            }
45          linalg.yield %1 : i32
46    } -> tensor<?xi32, #SparseVector>
47    return %0 : tensor<?xi32, #SparseVector>
48  }
49
50  // Negate existing values. Fill missing ones with +1.
51  func.func @vector_negation(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
52    %c = arith.constant 0 : index
53    %cf1 = arith.constant 1.0 : f64
54    %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
55    %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
56    %0 = linalg.generic #trait_vec_scale
57       ins(%arga: tensor<?xf64, #SparseVector>)
58        outs(%xv: tensor<?xf64, #SparseVector>) {
59        ^bb(%a: f64, %x: f64):
60          %1 = sparse_tensor.unary %a : f64 to f64
61            present={
62              ^bb0(%x0: f64):
63                %ret = arith.negf %x0 : f64
64                sparse_tensor.yield %ret : f64
65            }
66            absent={
67              sparse_tensor.yield %cf1 : f64
68            }
69          linalg.yield %1 : f64
70    } -> tensor<?xf64, #SparseVector>
71    return %0 : tensor<?xf64, #SparseVector>
72  }
73
74  // Clips values to the range [3, 7].
75  func.func @matrix_clip(%argx: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
76    %c0 = arith.constant 0 : index
77    %c1 = arith.constant 1 : index
78    %cfmin = arith.constant 3.0 : f64
79    %cfmax = arith.constant 7.0 : f64
80    %d0 = tensor.dim %argx, %c0 : tensor<?x?xf64, #DCSR>
81    %d1 = tensor.dim %argx, %c1 : tensor<?x?xf64, #DCSR>
82    %xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR>
83    %0 = linalg.generic #trait_mat_scale
84       ins(%argx: tensor<?x?xf64, #DCSR>)
85        outs(%xv: tensor<?x?xf64, #DCSR>) {
86        ^bb(%a: f64, %x: f64):
87          %1 = sparse_tensor.unary %a: f64 to f64
88            present={
89              ^bb0(%x0: f64):
90                %mincmp = arith.cmpf "ogt", %x0, %cfmin : f64
91                %x1 = arith.select %mincmp, %x0, %cfmin : f64
92                %maxcmp = arith.cmpf "olt", %x1, %cfmax : f64
93                %x2 = arith.select %maxcmp, %x1, %cfmax : f64
94                sparse_tensor.yield %x2 : f64
95            }
96            absent={}
97          linalg.yield %1 : f64
98    } -> tensor<?x?xf64, #DCSR>
99    return %0 : tensor<?x?xf64, #DCSR>
100  }
101
102  // Dumps a sparse vector of type f64.
103  func.func @dump_vec_f64(%arg0: tensor<?xf64, #SparseVector>) {
104    // Dump the values array to verify only sparse contents are stored.
105    %c0 = arith.constant 0 : index
106    %d0 = arith.constant -1.0 : f64
107    %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf64>
108    %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<32xf64>
109    vector.print %1 : vector<32xf64>
110    // Dump the dense vector to verify structure is correct.
111    %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64>
112    %2 = bufferization.to_memref %dv : memref<?xf64>
113    %3 = vector.transfer_read %2[%c0], %d0: memref<?xf64>, vector<32xf64>
114    vector.print %3 : vector<32xf64>
115    memref.dealloc %2 : memref<?xf64>
116    return
117  }
118
119  // Dumps a sparse vector of type i32.
120  func.func @dump_vec_i32(%arg0: tensor<?xi32, #SparseVector>) {
121    // Dump the values array to verify only sparse contents are stored.
122    %c0 = arith.constant 0 : index
123    %d0 = arith.constant -1 : i32
124    %0 = sparse_tensor.values %arg0 : tensor<?xi32, #SparseVector> to memref<?xi32>
125    %1 = vector.transfer_read %0[%c0], %d0: memref<?xi32>, vector<24xi32>
126    vector.print %1 : vector<24xi32>
127    // Dump the dense vector to verify structure is correct.
128    %dv = sparse_tensor.convert %arg0 : tensor<?xi32, #SparseVector> to tensor<?xi32>
129    %2 = bufferization.to_memref %dv : memref<?xi32>
130    %3 = vector.transfer_read %2[%c0], %d0: memref<?xi32>, vector<32xi32>
131    vector.print %3 : vector<32xi32>
132    memref.dealloc %2 : memref<?xi32>
133    return
134  }
135
136  // Dump a sparse matrix.
137  func.func @dump_mat(%arg0: tensor<?x?xf64, #DCSR>) {
138    %c0 = arith.constant 0 : index
139    %d0 = arith.constant -1.0 : f64
140    %0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64>
141    %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<16xf64>
142    vector.print %1 : vector<16xf64>
143    %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #DCSR> to tensor<?x?xf64>
144    %2 = bufferization.to_memref %dm : memref<?x?xf64>
145    %3 = vector.transfer_read %2[%c0, %c0], %d0: memref<?x?xf64>, vector<4x8xf64>
146    vector.print %3 : vector<4x8xf64>
147    memref.dealloc %2 : memref<?x?xf64>
148    return
149  }
150
151  // Driver method to call and verify vector kernels.
152  func.func @entry() {
153    %c0 = arith.constant 0 : index
154
155    // Setup sparse vectors.
156    %v1 = arith.constant sparse<
157       [ [0], [3], [11], [17], [20], [21], [28], [29], [31] ],
158         [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ]
159    > : tensor<32xf64>
160    %sv1 = sparse_tensor.convert %v1 : tensor<32xf64> to tensor<?xf64, #SparseVector>
161
162    // Setup sparse matrices.
163    %m1 = arith.constant sparse<
164       [ [0,0], [0,1], [1,7], [2,2], [2,4], [2,7], [3,0], [3,2], [3,3] ],
165         [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ]
166    > : tensor<4x8xf64>
167    %sm1 = sparse_tensor.convert %m1 : tensor<4x8xf64> to tensor<?x?xf64, #DCSR>
168
169    // Call sparse vector kernels.
170    %0 = call @vector_complement(%sv1)
171       : (tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector>
172    %1 = call @vector_negation(%sv1)
173       : (tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector>
174
175
176    // Call sparse matrix kernels.
177    %2 = call @matrix_clip(%sm1)
178      : (tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR>
179
180    //
181    // Verify the results.
182    //
183    // 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 )
184    // 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 )
185    // 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 )
186    // 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 )
187    // 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 )
188    // 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 )
189    // CHECK-NEXT: ( 3, 3, 3, 4, 5, 6, 7, 7, 7, -1, -1, -1, -1, -1, -1, -1 )
190    // 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 ) )
191    //
192    call @dump_vec_f64(%sv1) : (tensor<?xf64, #SparseVector>) -> ()
193    call @dump_vec_i32(%0) : (tensor<?xi32, #SparseVector>) -> ()
194    call @dump_vec_f64(%1) : (tensor<?xf64, #SparseVector>) -> ()
195    call @dump_mat(%2) : (tensor<?x?xf64, #DCSR>) -> ()
196
197    // Release the resources.
198    sparse_tensor.release %sv1 : tensor<?xf64, #SparseVector>
199    sparse_tensor.release %sm1 : tensor<?x?xf64, #DCSR>
200    sparse_tensor.release %0 : tensor<?xi32, #SparseVector>
201    sparse_tensor.release %1 : tensor<?xf64, #SparseVector>
202    sparse_tensor.release %2 : tensor<?x?xf64, #DCSR>
203    return
204  }
205}
206