1// RUN: mlir-opt %s \
2// RUN:   --sparsification --sparse-tensor-conversion \
3// RUN:   --convert-vector-to-scf --convert-scf-to-std \
4// RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
5// RUN:   --std-bufferize --finalizing-bufferize  \
6// RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-std-to-llvm --reconcile-unrealized-casts | \
7// RUN: mlir-cpu-runner \
8// RUN:  -e entry -entry-point-result=void  \
9// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
10// RUN: FileCheck %s
11
12//
13// Several common sparse storage schemes.
14//
15
16#Dense  = #sparse_tensor.encoding<{
17  dimLevelType = [ "dense", "dense" ]
18}>
19
20#CSR  = #sparse_tensor.encoding<{
21  dimLevelType = [ "dense", "compressed" ]
22}>
23
24#DCSR = #sparse_tensor.encoding<{
25  dimLevelType = [ "compressed", "compressed" ]
26}>
27
28#CSC = #sparse_tensor.encoding<{
29  dimLevelType = [ "dense", "compressed" ],
30  dimOrdering = affine_map<(i,j) -> (j,i)>
31}>
32
33#DCSC = #sparse_tensor.encoding<{
34  dimLevelType = [ "compressed", "compressed" ],
35  dimOrdering = affine_map<(i,j) -> (j,i)>
36}>
37
38#BlockRow = #sparse_tensor.encoding<{
39  dimLevelType = [ "compressed", "dense" ]
40}>
41
42#BlockCol = #sparse_tensor.encoding<{
43  dimLevelType = [ "compressed", "dense" ],
44  dimOrdering = affine_map<(i,j) -> (j,i)>
45}>
46
47//
48// Integration test that looks "under the hood" of sparse storage schemes.
49//
50module {
51  //
52  // Main driver that initializes a sparse tensor and inspects the sparse
53  // storage schemes in detail. Note that users of the MLIR sparse compiler
54  // are typically not concerned with such details, but the test ensures
55  // everything is working "under the hood".
56  //
57  func @entry() {
58    %c0 = arith.constant 0 : index
59    %c1 = arith.constant 1 : index
60    %d0 = arith.constant 0.0 : f64
61
62    //
63    // Initialize a dense tensor.
64    //
65    %t = arith.constant dense<[
66       [ 1.0,  0.0,  2.0,  0.0,  0.0,  0.0,  0.0,  3.0],
67       [ 0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0],
68       [ 0.0,  0.0,  4.0,  0.0,  0.0,  0.0,  0.0,  0.0],
69       [ 0.0,  0.0,  0.0,  5.0,  0.0,  0.0,  0.0,  0.0],
70       [ 0.0,  0.0,  0.0,  0.0,  6.0,  0.0,  0.0,  0.0],
71       [ 0.0,  7.0,  8.0,  0.0,  0.0,  0.0,  0.0,  9.0],
72       [ 0.0,  0.0, 10.0,  0.0,  0.0,  0.0, 11.0, 12.0],
73       [ 0.0, 13.0, 14.0,  0.0,  0.0,  0.0, 15.0, 16.0],
74       [ 0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0],
75       [ 0.0,  0.0,  0.0,  0.0,  0.0,  0.0, 17.0,  0.0]
76    ]> : tensor<10x8xf64>
77
78    //
79    // Convert dense tensor to various sparse tensors.
80    //
81    %0 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #Dense>
82    %1 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #CSR>
83    %2 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #DCSR>
84    %3 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #CSC>
85    %4 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #DCSC>
86    %x = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #BlockRow>
87    %y = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #BlockCol>
88
89    //
90    // Inspect storage scheme of Dense.
91    //
92    // CHECK:    ( 1, 0, 2, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0,
93    // CHECK-SAME: 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0,
94    // CHECK-SAME: 0, 0, 0, 0, 6, 0, 0, 0, 0, 7, 8, 0, 0, 0, 0, 9,
95    // CHECK-SAME: 0, 0, 10, 0, 0, 0, 11, 12, 0, 13, 14, 0, 0, 0, 15, 16,
96    // CHECK-SAME: 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 17, 0 )
97    //
98    %5 = sparse_tensor.values %0 : tensor<10x8xf64, #Dense> to memref<?xf64>
99    %6 = vector.transfer_read %5[%c0], %d0: memref<?xf64>, vector<80xf64>
100    vector.print %6 : vector<80xf64>
101
102    //
103    // Inspect storage scheme of CSR.
104    //
105    // pointers(1)
106    // indices(1)
107    // values
108    //
109    // CHECK: ( 0, 3, 3, 4, 5, 6, 9, 12, 16, 16, 17 )
110    // CHECK: ( 0, 2, 7, 2, 3, 4, 1, 2, 7, 2, 6, 7, 1, 2, 6, 7, 6 )
111    // CHECK: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17 )
112    //
113    %7 = sparse_tensor.pointers %1, %c1 : tensor<10x8xf64, #CSR> to memref<?xindex>
114    %8 = vector.transfer_read %7[%c0], %c0: memref<?xindex>, vector<11xindex>
115    vector.print %8 : vector<11xindex>
116    %9 = sparse_tensor.indices %1, %c1 : tensor<10x8xf64, #CSR> to memref<?xindex>
117    %10 = vector.transfer_read %9[%c0], %c0: memref<?xindex>, vector<17xindex>
118    vector.print %10 : vector<17xindex>
119    %11 = sparse_tensor.values %1 : tensor<10x8xf64, #CSR> to memref<?xf64>
120    %12 = vector.transfer_read %11[%c0], %d0: memref<?xf64>, vector<17xf64>
121    vector.print %12 : vector<17xf64>
122
123    //
124    // Inspect storage scheme of DCSR.
125    //
126    // pointers(0)
127    // indices(0)
128    // pointers(1)
129    // indices(1)
130    // values
131    //
132    // CHECK: ( 0, 8 )
133    // CHECK: ( 0, 2, 3, 4, 5, 6, 7, 9 )
134    // CHECK: ( 0, 3, 4, 5, 6, 9, 12, 16, 17 )
135    // CHECK: ( 0, 2, 7, 2, 3, 4, 1, 2, 7, 2, 6, 7, 1, 2, 6, 7, 6 )
136    // CHECK: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17 )
137    //
138    %13 = sparse_tensor.pointers %2, %c0 : tensor<10x8xf64, #DCSR> to memref<?xindex>
139    %14 = vector.transfer_read %13[%c0], %c0: memref<?xindex>, vector<2xindex>
140    vector.print %14 : vector<2xindex>
141    %15 = sparse_tensor.indices %2, %c0 : tensor<10x8xf64, #DCSR> to memref<?xindex>
142    %16 = vector.transfer_read %15[%c0], %c0: memref<?xindex>, vector<8xindex>
143    vector.print %16 : vector<8xindex>
144    %17 = sparse_tensor.pointers %2, %c1 : tensor<10x8xf64, #DCSR> to memref<?xindex>
145    %18 = vector.transfer_read %17[%c0], %c0: memref<?xindex>, vector<9xindex>
146    vector.print %18 : vector<9xindex>
147    %19 = sparse_tensor.indices %2, %c1 : tensor<10x8xf64, #DCSR> to memref<?xindex>
148    %20 = vector.transfer_read %19[%c0], %c0: memref<?xindex>, vector<17xindex>
149    vector.print %20 : vector<17xindex>
150    %21 = sparse_tensor.values %2 : tensor<10x8xf64, #DCSR> to memref<?xf64>
151    %22 = vector.transfer_read %21[%c0], %d0: memref<?xf64>, vector<17xf64>
152    vector.print %22 : vector<17xf64>
153
154    //
155    // Inspect storage scheme of CSC.
156    //
157    // pointers(1)
158    // indices(1)
159    // values
160    //
161    // CHECK: ( 0, 1, 3, 8, 9, 10, 10, 13, 17 )
162    // CHECK: ( 0, 5, 7, 0, 2, 5, 6, 7, 3, 4, 6, 7, 9, 0, 5, 6, 7 )
163    // CHECK: ( 1, 7, 13, 2, 4, 8, 10, 14, 5, 6, 11, 15, 17, 3, 9, 12, 16 )
164    //
165    %23 = sparse_tensor.pointers %3, %c1 : tensor<10x8xf64, #CSC> to memref<?xindex>
166    %24 = vector.transfer_read %23[%c0], %c0: memref<?xindex>, vector<9xindex>
167    vector.print %24 : vector<9xindex>
168    %25 = sparse_tensor.indices %3, %c1 : tensor<10x8xf64, #CSC> to memref<?xindex>
169    %26 = vector.transfer_read %25[%c0], %c0: memref<?xindex>, vector<17xindex>
170    vector.print %26 : vector<17xindex>
171    %27 = sparse_tensor.values %3 : tensor<10x8xf64, #CSC> to memref<?xf64>
172    %28 = vector.transfer_read %27[%c0], %d0: memref<?xf64>, vector<17xf64>
173    vector.print %28 : vector<17xf64>
174
175    //
176    // Inspect storage scheme of DCSC.
177    //
178    // pointers(0)
179    // indices(0)
180    // pointers(1)
181    // indices(1)
182    // values
183    //
184    // CHECK: ( 0, 7 )
185    // CHECK: ( 0, 1, 2, 3, 4, 6, 7 )
186    // CHECK: ( 0, 1, 3, 8, 9, 10, 13, 17 )
187    // CHECK: ( 0, 5, 7, 0, 2, 5, 6, 7, 3, 4, 6, 7, 9, 0, 5, 6, 7 )
188    // CHECK: ( 1, 7, 13, 2, 4, 8, 10, 14, 5, 6, 11, 15, 17, 3, 9, 12, 16 )
189    //
190    %29 = sparse_tensor.pointers %4, %c0 : tensor<10x8xf64, #DCSC> to memref<?xindex>
191    %30 = vector.transfer_read %29[%c0], %c0: memref<?xindex>, vector<2xindex>
192    vector.print %30 : vector<2xindex>
193    %31 = sparse_tensor.indices %4, %c0 : tensor<10x8xf64, #DCSC> to memref<?xindex>
194    %32 = vector.transfer_read %31[%c0], %c0: memref<?xindex>, vector<7xindex>
195    vector.print %32 : vector<7xindex>
196    %33 = sparse_tensor.pointers %4, %c1 : tensor<10x8xf64, #DCSC> to memref<?xindex>
197    %34 = vector.transfer_read %33[%c0], %c0: memref<?xindex>, vector<8xindex>
198    vector.print %34 : vector<8xindex>
199    %35 = sparse_tensor.indices %4, %c1 : tensor<10x8xf64, #DCSC> to memref<?xindex>
200    %36 = vector.transfer_read %35[%c0], %c0: memref<?xindex>, vector<17xindex>
201    vector.print %36 : vector<17xindex>
202    %37 = sparse_tensor.values %4 : tensor<10x8xf64, #DCSC> to memref<?xf64>
203    %38 = vector.transfer_read %37[%c0], %d0: memref<?xf64>, vector<17xf64>
204    vector.print %38 : vector<17xf64>
205
206    //
207    // Inspect storage scheme of BlockRow.
208    //
209    // pointers(0)
210    // indices(0)
211    // values
212    //
213    // CHECK: ( 0, 8 )
214    // CHECK: ( 0, 2, 3, 4, 5, 6, 7, 9 )
215    // CHECK: ( 1, 0, 2, 0, 0, 0, 0, 3, 0, 0, 4, 0, 0, 0, 0, 0,
216    // CHECK-SAME: 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0,
217    // CHECK-SAME: 0, 7, 8, 0, 0, 0, 0, 9, 0, 0, 10, 0, 0, 0, 11, 12,
218    // CHECK-SAME: 0, 13, 14, 0, 0, 0, 15, 16, 0, 0, 0, 0, 0, 0, 17, 0 )
219    //
220    %39 = sparse_tensor.pointers %x, %c0 : tensor<10x8xf64, #BlockRow> to memref<?xindex>
221    %40 = vector.transfer_read %39[%c0], %c0: memref<?xindex>, vector<2xindex>
222    vector.print %40 : vector<2xindex>
223    %41 = sparse_tensor.indices %x, %c0 : tensor<10x8xf64, #BlockRow> to memref<?xindex>
224    %42 = vector.transfer_read %41[%c0], %c0: memref<?xindex>, vector<8xindex>
225    vector.print %42 : vector<8xindex>
226    %43 = sparse_tensor.values %x : tensor<10x8xf64, #BlockRow> to memref<?xf64>
227    %44 = vector.transfer_read %43[%c0], %d0: memref<?xf64>, vector<64xf64>
228    vector.print %44 : vector<64xf64>
229
230    //
231    // Inspect storage scheme of BlockCol.
232    //
233    // pointers(0)
234    // indices(0)
235    // values
236    //
237    // CHECK: ( 0, 7 )
238    // CHECK: ( 0, 1, 2, 3, 4, 6, 7 )
239    // CHECK: ( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 0, 13, 0, 0, 2, 0, 4, 0,
240    // CHECK-SAME: 0, 8, 10, 14, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0,
241    // CHECK-SAME: 0, 0, 0, 0, 0, 0, 0, 0, 0, 11, 15, 0, 17, 3, 0, 0, 0, 0, 9, 12, 16, 0, 0 )
242    //
243    %45 = sparse_tensor.pointers %y, %c0 : tensor<10x8xf64, #BlockCol> to memref<?xindex>
244    %46 = vector.transfer_read %45[%c0], %c0: memref<?xindex>, vector<2xindex>
245    vector.print %46 : vector<2xindex>
246    %47 = sparse_tensor.indices %y, %c0 : tensor<10x8xf64, #BlockCol> to memref<?xindex>
247    %48 = vector.transfer_read %47[%c0], %c0: memref<?xindex>, vector<7xindex>
248    vector.print %48 : vector<7xindex>
249    %49 = sparse_tensor.values %y : tensor<10x8xf64, #BlockCol> to memref<?xf64>
250    %50 = vector.transfer_read %49[%c0], %d0: memref<?xf64>, vector<70xf64>
251    vector.print %50 : vector<70xf64>
252
253    // Release the resources.
254    sparse_tensor.release %0 : tensor<10x8xf64, #Dense>
255    sparse_tensor.release %1 : tensor<10x8xf64, #CSR>
256    sparse_tensor.release %2 : tensor<10x8xf64, #DCSR>
257    sparse_tensor.release %3 : tensor<10x8xf64, #CSC>
258    sparse_tensor.release %4 : tensor<10x8xf64, #DCSC>
259    sparse_tensor.release %x : tensor<10x8xf64, #BlockRow>
260    sparse_tensor.release %y : tensor<10x8xf64, #BlockCol>
261
262    return
263  }
264}
265