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