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:  | \
6// RUN: FileCheck %s
7
8#Tensor1  = #sparse_tensor.encoding<{
9  dimLevelType = [ "compressed", "compressed", "compressed" ],
10  dimOrdering = affine_map<(i,j,k) -> (i,j,k)>
11}>
12
13#Tensor2  = #sparse_tensor.encoding<{
14  dimLevelType = [ "compressed", "compressed", "compressed" ],
15  dimOrdering = affine_map<(i,j,k) -> (j,k,i)>
16}>
17
18#Tensor3  = #sparse_tensor.encoding<{
19  dimLevelType = [ "compressed", "compressed", "compressed" ],
20  dimOrdering = affine_map<(i,j,k) -> (k,i,j)>
21}>
22
23#Tensor4  = #sparse_tensor.encoding<{
24  dimLevelType = [ "dense", "compressed", "compressed" ],
25  dimOrdering = affine_map<(i,j,k) -> (i,j,k)>
26}>
27
28#Tensor5  = #sparse_tensor.encoding<{
29  dimLevelType = [ "dense", "compressed", "compressed" ],
30  dimOrdering = affine_map<(i,j,k) -> (j,k,i)>
31}>
32
33#Tensor6  = #sparse_tensor.encoding<{
34  dimLevelType = [ "dense", "compressed", "compressed" ],
35  dimOrdering = affine_map<(i,j,k) -> (k,i,j)>
36}>
37
38//
39// Integration test that tests conversions from sparse to dense tensors.
40//
41module {
42  //
43  // Utilities for output and releasing memory.
44  //
45  func @dump(%arg0: tensor<2x3x4xf64>) {
46    %c0 = arith.constant 0 : index
47    %d0 = arith.constant -1.0 : f64
48    %0 = vector.transfer_read %arg0[%c0, %c0, %c0], %d0: tensor<2x3x4xf64>, vector<2x3x4xf64>
49    vector.print %0 : vector<2x3x4xf64>
50    return
51  }
52  func @dumpAndRelease_234(%arg0: tensor<2x3x4xf64>) {
53    call @dump(%arg0) : (tensor<2x3x4xf64>) -> ()
54    %1 = bufferization.to_memref %arg0 : memref<2x3x4xf64>
55    memref.dealloc %1 : memref<2x3x4xf64>
56    return
57  }
58  func @dumpAndRelease_p34(%arg0: tensor<?x3x4xf64>) {
59    %0 = tensor.cast %arg0 : tensor<?x3x4xf64> to tensor<2x3x4xf64>
60    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
61    %1 = bufferization.to_memref %arg0 : memref<?x3x4xf64>
62    memref.dealloc %1 : memref<?x3x4xf64>
63    return
64  }
65  func @dumpAndRelease_2p4(%arg0: tensor<2x?x4xf64>) {
66    %0 = tensor.cast %arg0 : tensor<2x?x4xf64> to tensor<2x3x4xf64>
67    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
68    %1 = bufferization.to_memref %arg0 : memref<2x?x4xf64>
69    memref.dealloc %1 : memref<2x?x4xf64>
70    return
71  }
72  func @dumpAndRelease_23p(%arg0: tensor<2x3x?xf64>) {
73    %0 = tensor.cast %arg0 : tensor<2x3x?xf64> to tensor<2x3x4xf64>
74    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
75    %1 = bufferization.to_memref %arg0 : memref<2x3x?xf64>
76    memref.dealloc %1 : memref<2x3x?xf64>
77    return
78  }
79  func @dumpAndRelease_2pp(%arg0: tensor<2x?x?xf64>) {
80    %0 = tensor.cast %arg0 : tensor<2x?x?xf64> to tensor<2x3x4xf64>
81    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
82    %1 = bufferization.to_memref %arg0 : memref<2x?x?xf64>
83    memref.dealloc %1 : memref<2x?x?xf64>
84    return
85  }
86  func @dumpAndRelease_p3p(%arg0: tensor<?x3x?xf64>) {
87    %0 = tensor.cast %arg0 : tensor<?x3x?xf64> to tensor<2x3x4xf64>
88    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
89    %1 = bufferization.to_memref %arg0 : memref<?x3x?xf64>
90    memref.dealloc %1 : memref<?x3x?xf64>
91    return
92  }
93  func @dumpAndRelease_pp4(%arg0: tensor<?x?x4xf64>) {
94    %0 = tensor.cast %arg0 : tensor<?x?x4xf64> to tensor<2x3x4xf64>
95    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
96    %1 = bufferization.to_memref %arg0 : memref<?x?x4xf64>
97    memref.dealloc %1 : memref<?x?x4xf64>
98    return
99  }
100  func @dumpAndRelease_ppp(%arg0: tensor<?x?x?xf64>) {
101    %0 = tensor.cast %arg0 : tensor<?x?x?xf64> to tensor<2x3x4xf64>
102    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
103    %1 = bufferization.to_memref %arg0 : memref<?x?x?xf64>
104    memref.dealloc %1 : memref<?x?x?xf64>
105    return
106  }
107
108  //
109  // Main driver.
110  //
111  func @entry() {
112    //
113    // Initialize a 3-dim dense tensor.
114    //
115    %src = arith.constant dense<[
116       [  [  1.0,  2.0,  3.0,  4.0 ],
117          [  5.0,  6.0,  7.0,  8.0 ],
118          [  9.0, 10.0, 11.0, 12.0 ] ],
119       [  [ 13.0, 14.0, 15.0, 16.0 ],
120          [ 17.0, 18.0, 19.0, 20.0 ],
121          [ 21.0, 22.0, 23.0, 24.0 ] ]
122    ]> : tensor<2x3x4xf64>
123
124    //
125    // Convert dense tensor directly to various sparse tensors.
126    //
127    %s2341 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor1>
128    %s2342 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor2>
129    %s2343 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor3>
130    %s2344 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor4>
131    %s2345 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor5>
132    %s2346 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor6>
133
134    %sp344 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<?x3x4xf64, #Tensor4>
135    %sp345 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<?x3x4xf64, #Tensor5>
136    %sp346 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<?x3x4xf64, #Tensor6>
137    %s2p44 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x4xf64, #Tensor4>
138    %s2p45 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x4xf64, #Tensor5>
139    %s2p46 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x4xf64, #Tensor6>
140    %s23p4 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x?xf64, #Tensor4>
141    %s23p5 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x?xf64, #Tensor5>
142    %s23p6 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x?xf64, #Tensor6>
143    %s2pp4 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x?xf64, #Tensor4>
144    %s2pp5 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x?xf64, #Tensor5>
145    %s2pp6 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x?xf64, #Tensor6>
146
147    //
148    // Convert sparse tensor back to dense.
149    //
150    %d2341 = sparse_tensor.convert %s2341 : tensor<2x3x4xf64, #Tensor1> to tensor<2x3x4xf64>
151    %d2342 = sparse_tensor.convert %s2342 : tensor<2x3x4xf64, #Tensor2> to tensor<2x3x4xf64>
152    %d2343 = sparse_tensor.convert %s2343 : tensor<2x3x4xf64, #Tensor3> to tensor<2x3x4xf64>
153    %d2344 = sparse_tensor.convert %s2344 : tensor<2x3x4xf64, #Tensor4> to tensor<2x3x4xf64>
154    %d2345 = sparse_tensor.convert %s2345 : tensor<2x3x4xf64, #Tensor5> to tensor<2x3x4xf64>
155    %d2346 = sparse_tensor.convert %s2346 : tensor<2x3x4xf64, #Tensor6> to tensor<2x3x4xf64>
156
157    %dp344 = sparse_tensor.convert %sp344 : tensor<?x3x4xf64, #Tensor4> to tensor<?x3x4xf64>
158    %dp345 = sparse_tensor.convert %sp345 : tensor<?x3x4xf64, #Tensor5> to tensor<?x3x4xf64>
159    %dp346 = sparse_tensor.convert %sp346 : tensor<?x3x4xf64, #Tensor6> to tensor<?x3x4xf64>
160    %d2p44 = sparse_tensor.convert %s2p44 : tensor<2x?x4xf64, #Tensor4> to tensor<2x?x4xf64>
161    %d2p45 = sparse_tensor.convert %s2p45 : tensor<2x?x4xf64, #Tensor5> to tensor<2x?x4xf64>
162    %d2p46 = sparse_tensor.convert %s2p46 : tensor<2x?x4xf64, #Tensor6> to tensor<2x?x4xf64>
163    %d23p4 = sparse_tensor.convert %s23p4 : tensor<2x3x?xf64, #Tensor4> to tensor<2x3x?xf64>
164    %d23p5 = sparse_tensor.convert %s23p5 : tensor<2x3x?xf64, #Tensor5> to tensor<2x3x?xf64>
165    %d23p6 = sparse_tensor.convert %s23p6 : tensor<2x3x?xf64, #Tensor6> to tensor<2x3x?xf64>
166    %d2pp4 = sparse_tensor.convert %s2pp4 : tensor<2x?x?xf64, #Tensor4> to tensor<2x?x?xf64>
167    %d2pp5 = sparse_tensor.convert %s2pp5 : tensor<2x?x?xf64, #Tensor5> to tensor<2x?x?xf64>
168    %d2pp6 = sparse_tensor.convert %s2pp6 : tensor<2x?x?xf64, #Tensor6> to tensor<2x?x?xf64>
169
170    %dp3p4 = sparse_tensor.convert %sp344 : tensor<?x3x4xf64, #Tensor4> to tensor<?x3x?xf64>
171    %dp3p5 = sparse_tensor.convert %sp345 : tensor<?x3x4xf64, #Tensor5> to tensor<?x3x?xf64>
172    %dp3p6 = sparse_tensor.convert %sp346 : tensor<?x3x4xf64, #Tensor6> to tensor<?x3x?xf64>
173    %dpp44 = sparse_tensor.convert %s2p44 : tensor<2x?x4xf64, #Tensor4> to tensor<?x?x4xf64>
174    %dpp45 = sparse_tensor.convert %s2p45 : tensor<2x?x4xf64, #Tensor5> to tensor<?x?x4xf64>
175    %dpp46 = sparse_tensor.convert %s2p46 : tensor<2x?x4xf64, #Tensor6> to tensor<?x?x4xf64>
176    %dppp4 = sparse_tensor.convert %s2pp4 : tensor<2x?x?xf64, #Tensor4> to tensor<?x?x?xf64>
177    %dppp5 = sparse_tensor.convert %s2pp5 : tensor<2x?x?xf64, #Tensor5> to tensor<?x?x?xf64>
178    %dppp6 = sparse_tensor.convert %s2pp6 : tensor<2x?x?xf64, #Tensor6> to tensor<?x?x?xf64>
179
180    //
181    // Check round-trip equality.  And release dense tensors.
182    //
183    // CHECK-COUNT-28: ( ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) ), ( ( 13, 14, 15, 16 ), ( 17, 18, 19, 20 ), ( 21, 22, 23, 24 ) ) )
184    call @dump(%src) : (tensor<2x3x4xf64>) -> ()
185    call @dumpAndRelease_234(%d2341) : (tensor<2x3x4xf64>) -> ()
186    call @dumpAndRelease_234(%d2342) : (tensor<2x3x4xf64>) -> ()
187    call @dumpAndRelease_234(%d2343) : (tensor<2x3x4xf64>) -> ()
188    call @dumpAndRelease_234(%d2344) : (tensor<2x3x4xf64>) -> ()
189    call @dumpAndRelease_234(%d2345) : (tensor<2x3x4xf64>) -> ()
190    call @dumpAndRelease_234(%d2346) : (tensor<2x3x4xf64>) -> ()
191    call @dumpAndRelease_p34(%dp344) : (tensor<?x3x4xf64>) -> ()
192    call @dumpAndRelease_p34(%dp345) : (tensor<?x3x4xf64>) -> ()
193    call @dumpAndRelease_p34(%dp346) : (tensor<?x3x4xf64>) -> ()
194    call @dumpAndRelease_2p4(%d2p44) : (tensor<2x?x4xf64>) -> ()
195    call @dumpAndRelease_2p4(%d2p45) : (tensor<2x?x4xf64>) -> ()
196    call @dumpAndRelease_2p4(%d2p46) : (tensor<2x?x4xf64>) -> ()
197    call @dumpAndRelease_23p(%d23p4) : (tensor<2x3x?xf64>) -> ()
198    call @dumpAndRelease_23p(%d23p5) : (tensor<2x3x?xf64>) -> ()
199    call @dumpAndRelease_23p(%d23p6) : (tensor<2x3x?xf64>) -> ()
200    call @dumpAndRelease_2pp(%d2pp4) : (tensor<2x?x?xf64>) -> ()
201    call @dumpAndRelease_2pp(%d2pp5) : (tensor<2x?x?xf64>) -> ()
202    call @dumpAndRelease_2pp(%d2pp6) : (tensor<2x?x?xf64>) -> ()
203    call @dumpAndRelease_p3p(%dp3p4) : (tensor<?x3x?xf64>) -> ()
204    call @dumpAndRelease_p3p(%dp3p5) : (tensor<?x3x?xf64>) -> ()
205    call @dumpAndRelease_p3p(%dp3p6) : (tensor<?x3x?xf64>) -> ()
206    call @dumpAndRelease_pp4(%dpp44) : (tensor<?x?x4xf64>) -> ()
207    call @dumpAndRelease_pp4(%dpp45) : (tensor<?x?x4xf64>) -> ()
208    call @dumpAndRelease_pp4(%dpp46) : (tensor<?x?x4xf64>) -> ()
209    call @dumpAndRelease_ppp(%dppp4) : (tensor<?x?x?xf64>) -> ()
210    call @dumpAndRelease_ppp(%dppp5) : (tensor<?x?x?xf64>) -> ()
211    call @dumpAndRelease_ppp(%dppp6) : (tensor<?x?x?xf64>) -> ()
212
213    //
214    // Release sparse tensors.
215    //
216    sparse_tensor.release %s2341 : tensor<2x3x4xf64, #Tensor1>
217    sparse_tensor.release %s2342 : tensor<2x3x4xf64, #Tensor2>
218    sparse_tensor.release %s2343 : tensor<2x3x4xf64, #Tensor3>
219    sparse_tensor.release %s2344 : tensor<2x3x4xf64, #Tensor4>
220    sparse_tensor.release %s2345 : tensor<2x3x4xf64, #Tensor5>
221    sparse_tensor.release %s2346 : tensor<2x3x4xf64, #Tensor6>
222    sparse_tensor.release %sp344 : tensor<?x3x4xf64, #Tensor4>
223    sparse_tensor.release %sp345 : tensor<?x3x4xf64, #Tensor5>
224    sparse_tensor.release %sp346 : tensor<?x3x4xf64, #Tensor6>
225    sparse_tensor.release %s2p44 : tensor<2x?x4xf64, #Tensor4>
226    sparse_tensor.release %s2p45 : tensor<2x?x4xf64, #Tensor5>
227    sparse_tensor.release %s2p46 : tensor<2x?x4xf64, #Tensor6>
228    sparse_tensor.release %s23p4 : tensor<2x3x?xf64, #Tensor4>
229    sparse_tensor.release %s23p5 : tensor<2x3x?xf64, #Tensor5>
230    sparse_tensor.release %s23p6 : tensor<2x3x?xf64, #Tensor6>
231    sparse_tensor.release %s2pp4 : tensor<2x?x?xf64, #Tensor4>
232    sparse_tensor.release %s2pp5 : tensor<2x?x?xf64, #Tensor5>
233    sparse_tensor.release %s2pp6 : tensor<2x?x?xf64, #Tensor6>
234
235    return
236  }
237}
238