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.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.func @dumpAndRelease_234(%arg0: tensor<2x3x4xf64>) {
53    call @dump(%arg0) : (tensor<2x3x4xf64>) -> ()
54    return
55  }
56  func.func @dumpAndRelease_p34(%arg0: tensor<?x3x4xf64>) {
57    %0 = tensor.cast %arg0 : tensor<?x3x4xf64> to tensor<2x3x4xf64>
58    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
59    return
60  }
61  func.func @dumpAndRelease_2p4(%arg0: tensor<2x?x4xf64>) {
62    %0 = tensor.cast %arg0 : tensor<2x?x4xf64> to tensor<2x3x4xf64>
63    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
64    return
65  }
66  func.func @dumpAndRelease_23p(%arg0: tensor<2x3x?xf64>) {
67    %0 = tensor.cast %arg0 : tensor<2x3x?xf64> to tensor<2x3x4xf64>
68    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
69    return
70  }
71  func.func @dumpAndRelease_2pp(%arg0: tensor<2x?x?xf64>) {
72    %0 = tensor.cast %arg0 : tensor<2x?x?xf64> to tensor<2x3x4xf64>
73    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
74    return
75  }
76  func.func @dumpAndRelease_p3p(%arg0: tensor<?x3x?xf64>) {
77    %0 = tensor.cast %arg0 : tensor<?x3x?xf64> to tensor<2x3x4xf64>
78    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
79    return
80  }
81  func.func @dumpAndRelease_pp4(%arg0: tensor<?x?x4xf64>) {
82    %0 = tensor.cast %arg0 : tensor<?x?x4xf64> to tensor<2x3x4xf64>
83    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
84    return
85  }
86  func.func @dumpAndRelease_ppp(%arg0: tensor<?x?x?xf64>) {
87    %0 = tensor.cast %arg0 : tensor<?x?x?xf64> to tensor<2x3x4xf64>
88    call @dump(%0) : (tensor<2x3x4xf64>) -> ()
89    return
90  }
91
92  //
93  // Main driver.
94  //
95  func.func @entry() {
96    //
97    // Initialize a 3-dim dense tensor.
98    //
99    %src = arith.constant dense<[
100       [  [  1.0,  2.0,  3.0,  4.0 ],
101          [  5.0,  6.0,  7.0,  8.0 ],
102          [  9.0, 10.0, 11.0, 12.0 ] ],
103       [  [ 13.0, 14.0, 15.0, 16.0 ],
104          [ 17.0, 18.0, 19.0, 20.0 ],
105          [ 21.0, 22.0, 23.0, 24.0 ] ]
106    ]> : tensor<2x3x4xf64>
107
108    //
109    // Convert dense tensor directly to various sparse tensors.
110    //
111    %s2341 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor1>
112    %s2342 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor2>
113    %s2343 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor3>
114    %s2344 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor4>
115    %s2345 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor5>
116    %s2346 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor6>
117
118    %sp344 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<?x3x4xf64, #Tensor4>
119    %sp345 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<?x3x4xf64, #Tensor5>
120    %sp346 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<?x3x4xf64, #Tensor6>
121    %s2p44 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x4xf64, #Tensor4>
122    %s2p45 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x4xf64, #Tensor5>
123    %s2p46 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x4xf64, #Tensor6>
124    %s23p4 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x?xf64, #Tensor4>
125    %s23p5 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x?xf64, #Tensor5>
126    %s23p6 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x?xf64, #Tensor6>
127    %s2pp4 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x?xf64, #Tensor4>
128    %s2pp5 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x?xf64, #Tensor5>
129    %s2pp6 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x?x?xf64, #Tensor6>
130
131    //
132    // Convert sparse tensor back to dense.
133    //
134    %d2341 = sparse_tensor.convert %s2341 : tensor<2x3x4xf64, #Tensor1> to tensor<2x3x4xf64>
135    %d2342 = sparse_tensor.convert %s2342 : tensor<2x3x4xf64, #Tensor2> to tensor<2x3x4xf64>
136    %d2343 = sparse_tensor.convert %s2343 : tensor<2x3x4xf64, #Tensor3> to tensor<2x3x4xf64>
137    %d2344 = sparse_tensor.convert %s2344 : tensor<2x3x4xf64, #Tensor4> to tensor<2x3x4xf64>
138    %d2345 = sparse_tensor.convert %s2345 : tensor<2x3x4xf64, #Tensor5> to tensor<2x3x4xf64>
139    %d2346 = sparse_tensor.convert %s2346 : tensor<2x3x4xf64, #Tensor6> to tensor<2x3x4xf64>
140
141    %dp344 = sparse_tensor.convert %sp344 : tensor<?x3x4xf64, #Tensor4> to tensor<?x3x4xf64>
142    %dp345 = sparse_tensor.convert %sp345 : tensor<?x3x4xf64, #Tensor5> to tensor<?x3x4xf64>
143    %dp346 = sparse_tensor.convert %sp346 : tensor<?x3x4xf64, #Tensor6> to tensor<?x3x4xf64>
144    %d2p44 = sparse_tensor.convert %s2p44 : tensor<2x?x4xf64, #Tensor4> to tensor<2x?x4xf64>
145    %d2p45 = sparse_tensor.convert %s2p45 : tensor<2x?x4xf64, #Tensor5> to tensor<2x?x4xf64>
146    %d2p46 = sparse_tensor.convert %s2p46 : tensor<2x?x4xf64, #Tensor6> to tensor<2x?x4xf64>
147    %d23p4 = sparse_tensor.convert %s23p4 : tensor<2x3x?xf64, #Tensor4> to tensor<2x3x?xf64>
148    %d23p5 = sparse_tensor.convert %s23p5 : tensor<2x3x?xf64, #Tensor5> to tensor<2x3x?xf64>
149    %d23p6 = sparse_tensor.convert %s23p6 : tensor<2x3x?xf64, #Tensor6> to tensor<2x3x?xf64>
150    %d2pp4 = sparse_tensor.convert %s2pp4 : tensor<2x?x?xf64, #Tensor4> to tensor<2x?x?xf64>
151    %d2pp5 = sparse_tensor.convert %s2pp5 : tensor<2x?x?xf64, #Tensor5> to tensor<2x?x?xf64>
152    %d2pp6 = sparse_tensor.convert %s2pp6 : tensor<2x?x?xf64, #Tensor6> to tensor<2x?x?xf64>
153
154    %dp3p4 = sparse_tensor.convert %sp344 : tensor<?x3x4xf64, #Tensor4> to tensor<?x3x?xf64>
155    %dp3p5 = sparse_tensor.convert %sp345 : tensor<?x3x4xf64, #Tensor5> to tensor<?x3x?xf64>
156    %dp3p6 = sparse_tensor.convert %sp346 : tensor<?x3x4xf64, #Tensor6> to tensor<?x3x?xf64>
157    %dpp44 = sparse_tensor.convert %s2p44 : tensor<2x?x4xf64, #Tensor4> to tensor<?x?x4xf64>
158    %dpp45 = sparse_tensor.convert %s2p45 : tensor<2x?x4xf64, #Tensor5> to tensor<?x?x4xf64>
159    %dpp46 = sparse_tensor.convert %s2p46 : tensor<2x?x4xf64, #Tensor6> to tensor<?x?x4xf64>
160    %dppp4 = sparse_tensor.convert %s2pp4 : tensor<2x?x?xf64, #Tensor4> to tensor<?x?x?xf64>
161    %dppp5 = sparse_tensor.convert %s2pp5 : tensor<2x?x?xf64, #Tensor5> to tensor<?x?x?xf64>
162    %dppp6 = sparse_tensor.convert %s2pp6 : tensor<2x?x?xf64, #Tensor6> to tensor<?x?x?xf64>
163
164    //
165    // Check round-trip equality.  And release dense tensors.
166    //
167    // 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 ) ) )
168    call @dump(%src) : (tensor<2x3x4xf64>) -> ()
169    call @dumpAndRelease_234(%d2341) : (tensor<2x3x4xf64>) -> ()
170    call @dumpAndRelease_234(%d2342) : (tensor<2x3x4xf64>) -> ()
171    call @dumpAndRelease_234(%d2343) : (tensor<2x3x4xf64>) -> ()
172    call @dumpAndRelease_234(%d2344) : (tensor<2x3x4xf64>) -> ()
173    call @dumpAndRelease_234(%d2345) : (tensor<2x3x4xf64>) -> ()
174    call @dumpAndRelease_234(%d2346) : (tensor<2x3x4xf64>) -> ()
175    call @dumpAndRelease_p34(%dp344) : (tensor<?x3x4xf64>) -> ()
176    call @dumpAndRelease_p34(%dp345) : (tensor<?x3x4xf64>) -> ()
177    call @dumpAndRelease_p34(%dp346) : (tensor<?x3x4xf64>) -> ()
178    call @dumpAndRelease_2p4(%d2p44) : (tensor<2x?x4xf64>) -> ()
179    call @dumpAndRelease_2p4(%d2p45) : (tensor<2x?x4xf64>) -> ()
180    call @dumpAndRelease_2p4(%d2p46) : (tensor<2x?x4xf64>) -> ()
181    call @dumpAndRelease_23p(%d23p4) : (tensor<2x3x?xf64>) -> ()
182    call @dumpAndRelease_23p(%d23p5) : (tensor<2x3x?xf64>) -> ()
183    call @dumpAndRelease_23p(%d23p6) : (tensor<2x3x?xf64>) -> ()
184    call @dumpAndRelease_2pp(%d2pp4) : (tensor<2x?x?xf64>) -> ()
185    call @dumpAndRelease_2pp(%d2pp5) : (tensor<2x?x?xf64>) -> ()
186    call @dumpAndRelease_2pp(%d2pp6) : (tensor<2x?x?xf64>) -> ()
187    call @dumpAndRelease_p3p(%dp3p4) : (tensor<?x3x?xf64>) -> ()
188    call @dumpAndRelease_p3p(%dp3p5) : (tensor<?x3x?xf64>) -> ()
189    call @dumpAndRelease_p3p(%dp3p6) : (tensor<?x3x?xf64>) -> ()
190    call @dumpAndRelease_pp4(%dpp44) : (tensor<?x?x4xf64>) -> ()
191    call @dumpAndRelease_pp4(%dpp45) : (tensor<?x?x4xf64>) -> ()
192    call @dumpAndRelease_pp4(%dpp46) : (tensor<?x?x4xf64>) -> ()
193    call @dumpAndRelease_ppp(%dppp4) : (tensor<?x?x?xf64>) -> ()
194    call @dumpAndRelease_ppp(%dppp5) : (tensor<?x?x?xf64>) -> ()
195    call @dumpAndRelease_ppp(%dppp6) : (tensor<?x?x?xf64>) -> ()
196
197    //
198    // Release sparse tensors.
199    //
200    bufferization.dealloc_tensor %s2341 : tensor<2x3x4xf64, #Tensor1>
201    bufferization.dealloc_tensor %s2342 : tensor<2x3x4xf64, #Tensor2>
202    bufferization.dealloc_tensor %s2343 : tensor<2x3x4xf64, #Tensor3>
203    bufferization.dealloc_tensor %s2344 : tensor<2x3x4xf64, #Tensor4>
204    bufferization.dealloc_tensor %s2345 : tensor<2x3x4xf64, #Tensor5>
205    bufferization.dealloc_tensor %s2346 : tensor<2x3x4xf64, #Tensor6>
206    bufferization.dealloc_tensor %sp344 : tensor<?x3x4xf64, #Tensor4>
207    bufferization.dealloc_tensor %sp345 : tensor<?x3x4xf64, #Tensor5>
208    bufferization.dealloc_tensor %sp346 : tensor<?x3x4xf64, #Tensor6>
209    bufferization.dealloc_tensor %s2p44 : tensor<2x?x4xf64, #Tensor4>
210    bufferization.dealloc_tensor %s2p45 : tensor<2x?x4xf64, #Tensor5>
211    bufferization.dealloc_tensor %s2p46 : tensor<2x?x4xf64, #Tensor6>
212    bufferization.dealloc_tensor %s23p4 : tensor<2x3x?xf64, #Tensor4>
213    bufferization.dealloc_tensor %s23p5 : tensor<2x3x?xf64, #Tensor5>
214    bufferization.dealloc_tensor %s23p6 : tensor<2x3x?xf64, #Tensor6>
215    bufferization.dealloc_tensor %s2pp4 : tensor<2x?x?xf64, #Tensor4>
216    bufferization.dealloc_tensor %s2pp5 : tensor<2x?x?xf64, #Tensor5>
217    bufferization.dealloc_tensor %s2pp6 : tensor<2x?x?xf64, #Tensor6>
218
219    return
220  }
221}
222