1// RUN: mlir-opt %s -split-input-file | mlir-opt | FileCheck %s
2
3#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
4
5// CHECK-LABEL: func @sparse_new(
6// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
7//       CHECK: %[[T:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr<i8> to tensor<128xf64, #{{.*}}>
8//       CHECK: return %[[T]] : tensor<128xf64, #{{.*}}>
9func.func @sparse_new(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> {
10  %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector>
11  return %0 : tensor<128xf64, #SparseVector>
12}
13
14// -----
15
16#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
17
18// CHECK-LABEL: func @sparse_dealloc(
19// CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>
20//       CHECK: bufferization.dealloc_tensor %[[A]] : tensor<128xf64, #{{.*}}>
21//       CHECK: return
22func.func @sparse_dealloc(%arg0: tensor<128xf64, #SparseVector>) {
23  bufferization.dealloc_tensor %arg0 : tensor<128xf64, #SparseVector>
24  return
25}
26
27// -----
28
29#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
30
31// CHECK-LABEL: func @sparse_convert_1d_to_sparse(
32// CHECK-SAME: %[[A:.*]]: tensor<64xf32>)
33//       CHECK: %[[T:.*]] = sparse_tensor.convert %[[A]] : tensor<64xf32> to tensor<64xf32, #{{.*}}>
34//       CHECK: return %[[T]] : tensor<64xf32, #{{.*}}>
35func.func @sparse_convert_1d_to_sparse(%arg0: tensor<64xf32>) -> tensor<64xf32, #SparseVector> {
36  %0 = sparse_tensor.convert %arg0 : tensor<64xf32> to tensor<64xf32, #SparseVector>
37  return %0 : tensor<64xf32, #SparseVector>
38}
39
40// -----
41
42#SparseTensor = #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>
43
44// CHECK-LABEL: func @sparse_convert_3d_from_sparse(
45// CHECK-SAME: %[[A:.*]]: tensor<8x8x8xf64, #{{.*}}>)
46//       CHECK: %[[T:.*]] = sparse_tensor.convert %[[A]] : tensor<8x8x8xf64, #{{.*}}> to tensor<8x8x8xf64>
47//       CHECK: return %[[T]] : tensor<8x8x8xf64>
48func.func @sparse_convert_3d_from_sparse(%arg0: tensor<8x8x8xf64, #SparseTensor>) -> tensor<8x8x8xf64> {
49  %0 = sparse_tensor.convert %arg0 : tensor<8x8x8xf64, #SparseTensor> to tensor<8x8x8xf64>
50  return %0 : tensor<8x8x8xf64>
51}
52
53// -----
54
55#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
56
57// CHECK-LABEL: func @sparse_pointers(
58//  CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>)
59//       CHECK: %[[C:.*]] = arith.constant 0 : index
60//       CHECK: %[[T:.*]] = sparse_tensor.pointers %[[A]], %[[C]] : tensor<128xf64, #{{.*}}> to memref<?xindex>
61//       CHECK: return %[[T]] : memref<?xindex>
62func.func @sparse_pointers(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> {
63  %c = arith.constant 0 : index
64  %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex>
65  return %0 : memref<?xindex>
66}
67
68// -----
69
70#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
71
72// CHECK-LABEL: func @sparse_indices(
73//  CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>)
74//       CHECK: %[[C:.*]] = arith.constant 0 : index
75//       CHECK: %[[T:.*]] = sparse_tensor.indices %[[A]], %[[C]] : tensor<128xf64, #{{.*}}> to memref<?xindex>
76//       CHECK: return %[[T]] : memref<?xindex>
77func.func @sparse_indices(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> {
78  %c = arith.constant 0 : index
79  %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex>
80  return %0 : memref<?xindex>
81}
82
83// -----
84
85#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
86
87// CHECK-LABEL: func @sparse_values(
88//  CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>)
89//       CHECK: %[[T:.*]] = sparse_tensor.values %[[A]] : tensor<128xf64, #{{.*}}> to memref<?xf64>
90//       CHECK: return %[[T]] : memref<?xf64>
91func.func @sparse_values(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xf64> {
92  %0 = sparse_tensor.values %arg0 : tensor<128xf64, #SparseVector> to memref<?xf64>
93  return %0 : memref<?xf64>
94}
95
96// -----
97
98#DenseMatrix = #sparse_tensor.encoding<{dimLevelType = ["dense","dense"]}>
99
100// CHECK-LABEL: func @sparse_load(
101//  CHECK-SAME: %[[A:.*]]: tensor<16x32xf64, #{{.*}}>)
102//       CHECK: %[[T:.*]] = sparse_tensor.load %[[A]] : tensor<16x32xf64, #{{.*}}>
103//       CHECK: return %[[T]] : tensor<16x32xf64, #{{.*}}>
104func.func @sparse_load(%arg0: tensor<16x32xf64, #DenseMatrix>) -> tensor<16x32xf64, #DenseMatrix> {
105  %0 = sparse_tensor.load %arg0 : tensor<16x32xf64, #DenseMatrix>
106  return %0 : tensor<16x32xf64, #DenseMatrix>
107}
108
109// -----
110
111#DenseMatrix = #sparse_tensor.encoding<{dimLevelType = ["dense","dense"]}>
112
113// CHECK-LABEL: func @sparse_load_ins(
114//  CHECK-SAME: %[[A:.*]]: tensor<16x32xf64, #{{.*}}>)
115//       CHECK: %[[T:.*]] = sparse_tensor.load %[[A]] hasInserts : tensor<16x32xf64, #{{.*}}>
116//       CHECK: return %[[T]] : tensor<16x32xf64, #{{.*}}>
117func.func @sparse_load_ins(%arg0: tensor<16x32xf64, #DenseMatrix>) -> tensor<16x32xf64, #DenseMatrix> {
118  %0 = sparse_tensor.load %arg0 hasInserts : tensor<16x32xf64, #DenseMatrix>
119  return %0 : tensor<16x32xf64, #DenseMatrix>
120}
121
122// -----
123
124#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
125
126// CHECK-LABEL: func @sparse_insert(
127//  CHECK-SAME: %[[A:.*]]: tensor<128xf64, #sparse_tensor.encoding<{{.*}}>>,
128//  CHECK-SAME: %[[B:.*]]: memref<?xindex>,
129//  CHECK-SAME: %[[C:.*]]: f64) {
130//       CHECK: sparse_tensor.lex_insert %[[A]], %[[B]], %[[C]] : tensor<128xf64, #{{.*}}>, memref<?xindex>, f64
131//       CHECK: return
132func.func @sparse_insert(%arg0: tensor<128xf64, #SparseVector>, %arg1: memref<?xindex>, %arg2: f64) {
133  sparse_tensor.lex_insert %arg0, %arg1, %arg2 : tensor<128xf64, #SparseVector>, memref<?xindex>, f64
134  return
135}
136
137// -----
138
139#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
140
141// CHECK-LABEL: func @sparse_expansion(
142//  CHECK-SAME: %[[A:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>)
143//       CHECK: sparse_tensor.expand %[[A]]
144//       CHECK: return
145func.func @sparse_expansion(%arg0: tensor<8x8xf64, #SparseMatrix>) {
146  %values, %filled, %added, %count = sparse_tensor.expand %arg0
147    : tensor<8x8xf64, #SparseMatrix> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
148  return
149}
150
151// -----
152
153#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
154
155// CHECK-LABEL: func @sparse_compression(
156//  CHECK-SAME: %[[A:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>,
157//       CHECK: sparse_tensor.compress %[[A]]
158//       CHECK: return
159func.func @sparse_compression(%arg0: tensor<8x8xf64, #SparseMatrix>,
160                         %arg1: memref<?xindex>, %arg2: memref<?xf64>, %arg3: memref<?xi1>,
161                         %arg4: memref<?xindex>, %arg5: index) {
162  sparse_tensor.compress %arg0, %arg1, %arg2, %arg3, %arg4, %arg5
163    : tensor<8x8xf64, #SparseMatrix>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
164  return
165}
166
167// -----
168
169#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
170
171// CHECK-LABEL: func @sparse_out(
172//  CHECK-SAME: %[[A:.*]]: tensor<?x?xf64, #sparse_tensor.encoding<{{.*}}>>,
173//  CHECK-SAME: %[[B:.*]]: !llvm.ptr<i8>)
174//       CHECK: sparse_tensor.out %[[A]], %[[B]] : tensor<?x?xf64, #sparse_tensor.encoding<{{.*}}>>, !llvm.ptr<i8>
175//       CHECK: return
176func.func @sparse_out(%arg0: tensor<?x?xf64, #SparseMatrix>, %arg1: !llvm.ptr<i8>) {
177  sparse_tensor.out %arg0, %arg1 : tensor<?x?xf64, #SparseMatrix>, !llvm.ptr<i8>
178  return
179}
180
181// -----
182
183#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
184
185// CHECK-LABEL: func @sparse_binary(
186//  CHECK-SAME:   %[[A:.*]]: f64, %[[B:.*]]: i64) -> f64 {
187//       CHECK:   %[[Z:.*]] = arith.constant 0.000000e+00 : f64
188//       CHECK:   %[[C1:.*]] = sparse_tensor.binary %[[A]], %[[B]] : f64, i64 to f64
189//       CHECK:     overlap = {
190//       CHECK:       ^bb0(%[[A1:.*]]: f64, %[[B1:.*]]: i64):
191//       CHECK:         sparse_tensor.yield %[[A1]] : f64
192//       CHECK:     }
193//       CHECK:     left = identity
194//       CHECK:     right = {
195//       CHECK:       ^bb0(%[[A2:.*]]: i64):
196//       CHECK:         sparse_tensor.yield %[[Z]] : f64
197//       CHECK:     }
198//       CHECK:   return %[[C1]] : f64
199//       CHECK: }
200func.func @sparse_binary(%arg0: f64, %arg1: i64) -> f64 {
201  %cf0 = arith.constant 0.0 : f64
202  %r = sparse_tensor.binary %arg0, %arg1 : f64, i64 to f64
203    overlap={
204      ^bb0(%x: f64, %y: i64):
205        sparse_tensor.yield %x : f64
206    }
207    left=identity
208    right={
209      ^bb0(%y: i64):
210        sparse_tensor.yield %cf0 : f64
211    }
212  return %r : f64
213}
214
215// -----
216
217#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
218
219// CHECK-LABEL: func @sparse_unary(
220//  CHECK-SAME:   %[[A:.*]]: f64) -> f64 {
221//       CHECK:   %[[C1:.*]] = sparse_tensor.unary %[[A]] : f64 to f64
222//       CHECK:     present = {
223//       CHECK:       ^bb0(%[[A1:.*]]: f64):
224//       CHECK:         sparse_tensor.yield %[[A1]] : f64
225//       CHECK:     }
226//       CHECK:     absent = {
227//       CHECK:       %[[R:.*]] = arith.constant -1.000000e+00 : f64
228//       CHECK:       sparse_tensor.yield %[[R]] : f64
229//       CHECK:     }
230//       CHECK:   return %[[C1]] : f64
231//       CHECK: }
232func.func @sparse_unary(%arg0: f64) -> f64 {
233  %r = sparse_tensor.unary %arg0 : f64 to f64
234    present={
235      ^bb0(%x: f64):
236        sparse_tensor.yield %x : f64
237    } absent={
238      ^bb0:
239        %cf1 = arith.constant -1.0 : f64
240        sparse_tensor.yield %cf1 : f64
241    }
242  return %r : f64
243}
244
245// -----
246
247#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
248
249// CHECK-LABEL: func @sparse_unary(
250//  CHECK-SAME:   %[[A:.*]]: f64) -> i64 {
251//       CHECK:   %[[C1:.*]] = sparse_tensor.unary %[[A]] : f64 to i64
252//       CHECK:     present = {
253//       CHECK:       ^bb0(%[[A1:.*]]: f64):
254//       CHECK:         %[[R:.*]] = arith.fptosi %[[A1]] : f64 to i64
255//       CHECK:         sparse_tensor.yield %[[R]] : i64
256//       CHECK:     }
257//       CHECK:     absent = {
258//       CHECK:     }
259//       CHECK:   return %[[C1]] : i64
260//       CHECK: }
261func.func @sparse_unary(%arg0: f64) -> i64 {
262  %r = sparse_tensor.unary %arg0 : f64 to i64
263    present={
264      ^bb0(%x: f64):
265        %ret = arith.fptosi %x : f64 to i64
266        sparse_tensor.yield %ret : i64
267    }
268    absent={}
269  return %r : i64
270}
271
272// -----
273
274#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
275
276// CHECK-LABEL: func @sparse_reduce_2d_to_1d(
277//  CHECK-SAME:   %[[A:.*]]: f64, %[[B:.*]]: f64) -> f64 {
278//       CHECK:   %[[Z:.*]] = arith.constant 0.000000e+00 : f64
279//       CHECK:   %[[C1:.*]] = sparse_tensor.reduce %[[A]], %[[B]], %[[Z]] : f64 {
280//       CHECK:       ^bb0(%[[A1:.*]]: f64, %[[B1:.*]]: f64):
281//       CHECK:         sparse_tensor.yield %[[A1]] : f64
282//       CHECK:     }
283//       CHECK:   return %[[C1]] : f64
284//       CHECK: }
285func.func @sparse_reduce_2d_to_1d(%arg0: f64, %arg1: f64) -> f64 {
286  %cf0 = arith.constant 0.0 : f64
287  %r = sparse_tensor.reduce %arg0, %arg1, %cf0 : f64 {
288      ^bb0(%x: f64, %y: f64):
289        sparse_tensor.yield %x : f64
290    }
291  return %r : f64
292}