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 --lower-affine \
6// RUN:   --convert-vector-to-llvm --convert-memref-to-llvm \
7// RUN:   --convert-std-to-llvm --reconcile-unrealized-casts | \
8// RUN: mlir-cpu-runner \
9// RUN:  -e entry -entry-point-result=void  \
10// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
11// RUN: FileCheck %s
12//
13// Do the same run, but now with SIMDization as well. This should not change the outcome.
14//
15// RUN: mlir-opt %s \
16// RUN:   --sparsification="vectorization-strategy=2 vl=2 enable-simd-index32" --sparse-tensor-conversion \
17// RUN:   --convert-vector-to-scf --convert-scf-to-std \
18// RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
19// RUN:   --std-bufferize --finalizing-bufferize --lower-affine \
20// RUN:   --convert-vector-to-llvm --convert-memref-to-llvm \
21// RUN:   --convert-std-to-llvm --reconcile-unrealized-casts | \
22// RUN: mlir-cpu-runner \
23// RUN:  -e entry -entry-point-result=void  \
24// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
25// RUN: FileCheck %s
26//
27
28#SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
29
30#trait_cast = {
31  indexing_maps = [
32    affine_map<(i) -> (i)>,  // A (in)
33    affine_map<(i) -> (i)>   // X (out)
34  ],
35  iterator_types = ["parallel"],
36  doc = "X(i) = cast A(i)"
37}
38
39//
40// Integration test that lowers a kernel annotated as sparse to actual sparse
41// code, initializes a matching sparse storage scheme from a dense vector,
42// and runs the resulting code with the JIT compiler.
43//
44module {
45  //
46  // Various kernels that cast a sparse vector from one type to another.
47  // Standard supports the following casts.
48  //   sitofp
49  //   uitofp
50  //   fptosi
51  //   fptoui
52  //   fpext
53  //   fptrunc
54  //   sexti
55  //   zexti
56  //   trunci
57  //   bitcast
58  // Since all casts are "zero preserving" unary operations, lattice computation
59  // and conversion to sparse code is straightforward.
60  //
61  func @sparse_cast_s32_to_f32(%arga: tensor<10xi32, #SV>) -> tensor<10xf32> {
62    %argx = constant dense<0.0> : tensor<10xf32>
63    %0 = linalg.generic #trait_cast
64      ins(%arga: tensor<10xi32, #SV>)
65      outs(%argx: tensor<10xf32>) {
66        ^bb(%a: i32, %x : f32):
67          %cst = sitofp %a : i32 to f32
68          linalg.yield %cst : f32
69    } -> tensor<10xf32>
70    return %0 : tensor<10xf32>
71  }
72  func @sparse_cast_u32_to_f32(%arga: tensor<10xi32, #SV>) -> tensor<10xf32> {
73    %argx = constant dense<0.0> : tensor<10xf32>
74    %0 = linalg.generic #trait_cast
75      ins(%arga: tensor<10xi32, #SV>)
76      outs(%argx: tensor<10xf32>) {
77        ^bb(%a: i32, %x : f32):
78          %cst = uitofp %a : i32 to f32
79          linalg.yield %cst : f32
80    } -> tensor<10xf32>
81    return %0 : tensor<10xf32>
82  }
83  func @sparse_cast_f32_to_s32(%arga: tensor<10xf32, #SV>) -> tensor<10xi32> {
84    %argx = constant dense<0> : tensor<10xi32>
85    %0 = linalg.generic #trait_cast
86      ins(%arga: tensor<10xf32, #SV>)
87      outs(%argx: tensor<10xi32>) {
88        ^bb(%a: f32, %x : i32):
89          %cst = fptosi %a : f32 to i32
90          linalg.yield %cst : i32
91    } -> tensor<10xi32>
92    return %0 : tensor<10xi32>
93  }
94  func @sparse_cast_f64_to_u32(%arga: tensor<10xf64, #SV>) -> tensor<10xi32> {
95    %argx = constant dense<0> : tensor<10xi32>
96    %0 = linalg.generic #trait_cast
97      ins(%arga: tensor<10xf64, #SV>)
98      outs(%argx: tensor<10xi32>) {
99        ^bb(%a: f64, %x : i32):
100          %cst = fptoui %a : f64 to i32
101          linalg.yield %cst : i32
102    } -> tensor<10xi32>
103    return %0 : tensor<10xi32>
104  }
105  func @sparse_cast_f32_to_f64(%arga: tensor<10xf32, #SV>) -> tensor<10xf64> {
106    %argx = constant dense<0.0> : tensor<10xf64>
107    %0 = linalg.generic #trait_cast
108      ins(%arga: tensor<10xf32, #SV>)
109      outs(%argx: tensor<10xf64>) {
110        ^bb(%a: f32, %x : f64):
111          %cst = fpext %a : f32 to f64
112          linalg.yield %cst : f64
113    } -> tensor<10xf64>
114    return %0 : tensor<10xf64>
115  }
116  func @sparse_cast_f64_to_f32(%arga: tensor<10xf64, #SV>) -> tensor<10xf32> {
117    %argx = constant dense<0.0> : tensor<10xf32>
118    %0 = linalg.generic #trait_cast
119      ins(%arga: tensor<10xf64, #SV>)
120      outs(%argx: tensor<10xf32>) {
121        ^bb(%a: f64, %x : f32):
122          %cst = fptrunc %a : f64 to f32
123          linalg.yield %cst : f32
124    } -> tensor<10xf32>
125    return %0 : tensor<10xf32>
126  }
127  func @sparse_cast_s32_to_u64(%arga: tensor<10xi32, #SV>) -> tensor<10xi64> {
128    %argx = constant dense<0> : tensor<10xi64>
129    %0 = linalg.generic #trait_cast
130      ins(%arga: tensor<10xi32, #SV>)
131      outs(%argx: tensor<10xi64>) {
132        ^bb(%a: i32, %x : i64):
133          %cst = sexti %a : i32 to i64
134          linalg.yield %cst : i64
135    } -> tensor<10xi64>
136    return %0 : tensor<10xi64>
137  }
138  func @sparse_cast_u32_to_s64(%arga: tensor<10xi32, #SV>) -> tensor<10xi64> {
139    %argx = constant dense<0> : tensor<10xi64>
140    %0 = linalg.generic #trait_cast
141      ins(%arga: tensor<10xi32, #SV>)
142      outs(%argx: tensor<10xi64>) {
143        ^bb(%a: i32, %x : i64):
144          %cst = zexti %a : i32 to i64
145          linalg.yield %cst : i64
146    } -> tensor<10xi64>
147    return %0 : tensor<10xi64>
148  }
149  func @sparse_cast_i32_to_i8(%arga: tensor<10xi32, #SV>) -> tensor<10xi8> {
150    %argx = constant dense<0> : tensor<10xi8>
151    %0 = linalg.generic #trait_cast
152      ins(%arga: tensor<10xi32, #SV>)
153      outs(%argx: tensor<10xi8>) {
154        ^bb(%a: i32, %x : i8):
155          %cst = trunci %a : i32 to i8
156          linalg.yield %cst : i8
157    } -> tensor<10xi8>
158    return %0 : tensor<10xi8>
159  }
160  func @sparse_cast_f32_as_s32(%arga: tensor<10xf32, #SV>) -> tensor<10xi32> {
161    %argx = constant dense<0> : tensor<10xi32>
162    %0 = linalg.generic #trait_cast
163      ins(%arga: tensor<10xf32, #SV>)
164      outs(%argx: tensor<10xi32>) {
165        ^bb(%a: f32, %x : i32):
166          %cst = bitcast %a : f32 to i32
167          linalg.yield %cst : i32
168    } -> tensor<10xi32>
169    return %0 : tensor<10xi32>
170  }
171
172  //
173  // Main driver that converts a dense tensor into a sparse tensor
174  // and then calls the sparse casting kernel.
175  //
176  func @entry() {
177    %z = constant 0 : index
178    %b = constant 0 : i8
179    %i = constant 0 : i32
180    %l = constant 0 : i64
181    %f = constant 0.0 : f32
182    %d = constant 0.0 : f64
183
184    // Initialize dense tensors, convert to a sparse vectors.
185    %0 = constant dense<[ -4, -3, -2, -1, 0, 1, 2, 3, 4, 305 ]> : tensor<10xi32>
186    %1 = sparse_tensor.convert %0 : tensor<10xi32> to tensor<10xi32, #SV>
187    %2 = constant dense<[ -4.4, -3.3, -2.2, -1.1, 0.0, 1.1, 2.2, 3.3, 4.4, 305.5 ]> : tensor<10xf32>
188    %3 = sparse_tensor.convert %2 : tensor<10xf32> to tensor<10xf32, #SV>
189    %4 = constant dense<[ -4.4, -3.3, -2.2, -1.1, 0.0, 1.1, 2.2, 3.3, 4.4, 305.5 ]> : tensor<10xf64>
190    %5 = sparse_tensor.convert %4 : tensor<10xf64> to tensor<10xf64, #SV>
191    %6 = constant dense<[ 4294967295.0, 4294967294.0, 4294967293.0, 4294967292.0,
192                          0.0, 1.1, 2.2, 3.3, 4.4, 305.5 ]> : tensor<10xf64>
193    %7 = sparse_tensor.convert %6 : tensor<10xf64> to tensor<10xf64, #SV>
194
195    //
196    // CHECK: ( -4, -3, -2, -1, 0, 1, 2, 3, 4, 305 )
197    //
198    %c0 = call @sparse_cast_s32_to_f32(%1) : (tensor<10xi32, #SV>) -> tensor<10xf32>
199    %m0 = memref.buffer_cast %c0 : memref<10xf32>
200    %v0 = vector.transfer_read %m0[%z], %f: memref<10xf32>, vector<10xf32>
201    vector.print %v0 : vector<10xf32>
202
203    //
204    // CHECK: ( 4.29497e+09, 4.29497e+09, 4.29497e+09, 4.29497e+09, 0, 1, 2, 3, 4, 305 )
205    //
206    %c1 = call @sparse_cast_u32_to_f32(%1) : (tensor<10xi32, #SV>) -> tensor<10xf32>
207    %m1 = memref.buffer_cast %c1 : memref<10xf32>
208    %v1 = vector.transfer_read %m1[%z], %f: memref<10xf32>, vector<10xf32>
209    vector.print %v1 : vector<10xf32>
210
211    //
212    // CHECK: ( -4, -3, -2, -1, 0, 1, 2, 3, 4, 305 )
213    //
214    %c2 = call @sparse_cast_f32_to_s32(%3) : (tensor<10xf32, #SV>) -> tensor<10xi32>
215    %m2 = memref.buffer_cast %c2 : memref<10xi32>
216    %v2 = vector.transfer_read %m2[%z], %i: memref<10xi32>, vector<10xi32>
217    vector.print %v2 : vector<10xi32>
218
219    //
220    // CHECK: ( 4294967295, 4294967294, 4294967293, 4294967292, 0, 1, 2, 3, 4, 305 )
221    //
222    %c3 = call @sparse_cast_f64_to_u32(%7) : (tensor<10xf64, #SV>) -> tensor<10xi32>
223    %m3 = memref.buffer_cast %c3 : memref<10xi32>
224    %v3 = vector.transfer_read %m3[%z], %i: memref<10xi32>, vector<10xi32>
225    %vu = vector.bitcast %v3 : vector<10xi32> to vector<10xui32>
226    vector.print %vu : vector<10xui32>
227
228    //
229    // CHECK: ( -4.4, -3.3, -2.2, -1.1, 0, 1.1, 2.2, 3.3, 4.4, 305.5 )
230    //
231    %c4 = call @sparse_cast_f32_to_f64(%3) : (tensor<10xf32, #SV>) -> tensor<10xf64>
232    %m4 = memref.buffer_cast %c4 : memref<10xf64>
233    %v4 = vector.transfer_read %m4[%z], %d: memref<10xf64>, vector<10xf64>
234    vector.print %v4 : vector<10xf64>
235
236    //
237    // CHECK: ( -4.4, -3.3, -2.2, -1.1, 0, 1.1, 2.2, 3.3, 4.4, 305.5 )
238    //
239    %c5 = call @sparse_cast_f64_to_f32(%5) : (tensor<10xf64, #SV>) -> tensor<10xf32>
240    %m5 = memref.buffer_cast %c5 : memref<10xf32>
241    %v5 = vector.transfer_read %m5[%z], %f: memref<10xf32>, vector<10xf32>
242    vector.print %v5 : vector<10xf32>
243
244    //
245    // CHECK: ( -4, -3, -2, -1, 0, 1, 2, 3, 4, 305 )
246    //
247    %c6 = call @sparse_cast_s32_to_u64(%1) : (tensor<10xi32, #SV>) -> tensor<10xi64>
248    %m6 = memref.buffer_cast %c6 : memref<10xi64>
249    %v6 = vector.transfer_read %m6[%z], %l: memref<10xi64>, vector<10xi64>
250    vector.print %v6 : vector<10xi64>
251
252    //
253    // CHECK: ( 4294967292, 4294967293, 4294967294, 4294967295, 0, 1, 2, 3, 4, 305 )
254    //
255    %c7 = call @sparse_cast_u32_to_s64(%1) : (tensor<10xi32, #SV>) -> tensor<10xi64>
256    %m7 = memref.buffer_cast %c7 : memref<10xi64>
257    %v7 = vector.transfer_read %m7[%z], %l: memref<10xi64>, vector<10xi64>
258    vector.print %v7 : vector<10xi64>
259
260    //
261    // CHECK: ( -4, -3, -2, -1, 0, 1, 2, 3, 4, 49 )
262    //
263    %c8 = call @sparse_cast_i32_to_i8(%1) : (tensor<10xi32, #SV>) -> tensor<10xi8>
264    %m8 = memref.buffer_cast %c8 : memref<10xi8>
265    %v8 = vector.transfer_read %m8[%z], %b: memref<10xi8>, vector<10xi8>
266    vector.print %v8 : vector<10xi8>
267
268    //
269    // CHECK: ( -1064514355, -1068289229, -1072902963, -1081291571, 0, 1066192077, 1074580685, 1079194419, 1082969293, 1134084096 )
270    //
271    %c9 = call @sparse_cast_f32_as_s32(%3) : (tensor<10xf32, #SV>) -> tensor<10xi32>
272    %m9 = memref.buffer_cast %c9 : memref<10xi32>
273    %v9 = vector.transfer_read %m9[%z], %i: memref<10xi32>, vector<10xi32>
274    vector.print %v9 : vector<10xi32>
275
276    return
277  }
278}
279