xref: /llvm-project/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir (revision 13af97a70e7202507dcca89d2f732e5126d2bbcd)
1// RUN: mlir-opt %s -split-input-file | mlir-opt -split-input-file | FileCheck %s
2
3#SV  = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
4
5// CHECK: #[[$SV:.*]] = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
6// CHECK-LABEL: func private @sparse_1d_tensor(
7// CHECK-SAME: tensor<32xf64, #[[$SV]]>)
8func.func private @sparse_1d_tensor(tensor<32xf64, #SV>)
9
10// -----
11
12#CSR = #sparse_tensor.encoding<{
13  map = (d0, d1) -> (d0 : dense, d1 : compressed),
14  posWidth = 64,
15  crdWidth = 64
16}>
17
18// CHECK: #[[$CSR:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed), posWidth = 64, crdWidth = 64 }>
19// CHECK-LABEL: func private @sparse_csr(
20// CHECK-SAME: tensor<?x?xf32, #[[$CSR]]>)
21func.func private @sparse_csr(tensor<?x?xf32, #CSR>)
22
23// -----
24
25#CSR_OnlyOnes = #sparse_tensor.encoding<{
26  map = (d0, d1) -> (d0 : dense, d1 : compressed),
27  posWidth = 64,
28  crdWidth = 64,
29  explicitVal = 1.0 : f32,
30  implicitVal = 0.0 : f32
31}>
32
33// CHECK: #[[$CSR_OnlyOnes:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed), posWidth = 64, crdWidth = 64, explicitVal = 1.000000e+00 : f32, implicitVal = 0.000000e+00 : f32 }>
34// CHECK-LABEL: func private @sparse_csr(
35// CHECK-SAME: tensor<?x?xf32, #[[$CSR_OnlyOnes]]>)
36func.func private @sparse_csr(tensor<?x?xf32, #CSR_OnlyOnes>)
37
38// -----
39
40#CSR_OnlyOnes = #sparse_tensor.encoding<{
41  map = (d0, d1) -> (d0 : dense, d1 : compressed),
42  explicitVal = 1.0 : f64,
43  implicitVal = 0.0 : f64
44}>
45
46// CHECK: #[[$CSR_OnlyOnes:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed), explicitVal = 1.000000e+00 : f64, implicitVal = 0.000000e+00 : f64 }>
47// CHECK-LABEL: func private @sparse_csr(
48// CHECK-SAME: tensor<?x?xf64, #[[$CSR_OnlyOnes]]>)
49func.func private @sparse_csr(tensor<?x?xf64, #CSR_OnlyOnes>)
50
51// -----
52
53#CSR_OnlyOnes = #sparse_tensor.encoding<{
54  map = (d0, d1) -> (d0 : dense, d1 : compressed),
55  posWidth = 64,
56  crdWidth = 64,
57  explicitVal = 1 : i32,
58  implicitVal = 0 : i32
59}>
60
61// CHECK: #[[$CSR_OnlyOnes:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed), posWidth = 64, crdWidth = 64, explicitVal = 1 : i32, implicitVal = 0 : i32 }>
62// CHECK-LABEL: func private @sparse_csr(
63// CHECK-SAME: tensor<?x?xi32, #[[$CSR_OnlyOnes]]>)
64func.func private @sparse_csr(tensor<?x?xi32, #CSR_OnlyOnes>)
65
66// -----
67
68#CSR_OnlyOnes = #sparse_tensor.encoding<{
69  map = (d0, d1) -> (d0 : dense, d1 : compressed),
70  posWidth = 64,
71  crdWidth = 64,
72  explicitVal = 1 : i64,
73  implicitVal = 0 : i64
74}>
75
76// CHECK: #[[$CSR_OnlyOnes:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed), posWidth = 64, crdWidth = 64, explicitVal = 1 : i64, implicitVal = 0 : i64 }>
77// CHECK-LABEL: func private @sparse_csr(
78// CHECK-SAME: tensor<?x?xi64, #[[$CSR_OnlyOnes]]>)
79func.func private @sparse_csr(tensor<?x?xi64, #CSR_OnlyOnes>)
80
81// -----
82
83#CSR_OnlyOnes = #sparse_tensor.encoding<{
84  map = (d0, d1) -> (d0 : dense, d1 : compressed),
85  posWidth = 64,
86  crdWidth = 64,
87  explicitVal = #complex.number<:f32 1.0, 0.0>,
88  implicitVal = #complex.number<:f32 0.0, 0.0>
89}>
90
91// CHECK: #[[$CSR_OnlyOnes:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed), posWidth = 64, crdWidth = 64, explicitVal = #complex.number<:f32 1.000000e+00, 0.000000e+00> : complex<f32>, implicitVal = #complex.number<:f32 0.000000e+00, 0.000000e+00> : complex<f32> }>
92// CHECK-LABEL: func private @sparse_csr(
93// CHECK-SAME: tensor<?x?xcomplex<f32>, #[[$CSR_OnlyOnes]]>)
94func.func private @sparse_csr(tensor<?x?xcomplex<f32>, #CSR_OnlyOnes>)
95
96// -----
97
98#BCSR = #sparse_tensor.encoding<{
99  map = (d0, d1, d2) -> (d0 : batch, d1: dense, d2 : compressed),
100}>
101
102// CHECK: #[[$BCSR:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : batch, d1 : dense, d2 : compressed) }>
103// CHECK-LABEL: func private @sparse_bcsr(
104// CHECK-SAME: tensor<?x?x?xf32, #[[$BCSR]]>)
105func.func private @sparse_bcsr(tensor<?x?x?xf32, #BCSR>)
106
107// -----
108
109#CSR_explicit = #sparse_tensor.encoding<{
110  map = {l0, l1} (d0 = l0, d1 = l1) -> (l0 = d0 : dense, l1 = d1 : compressed)
111}>
112
113// CHECK: #[[$CSR_EXPLICIT:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }>
114// CHECK-LABEL: func private @CSR_explicit(
115// CHECK-SAME: tensor<?x?xf64, #[[$CSR_EXPLICIT]]>
116func.func private @CSR_explicit(%arg0: tensor<?x?xf64, #CSR_explicit>) {
117  return
118}
119
120// -----
121
122#CSC = #sparse_tensor.encoding<{
123  map = (d0, d1) -> (d1 : dense, d0 : compressed),
124  posWidth = 0,
125  crdWidth = 0
126}>
127
128// CHECK-DAG: #[[$CSC:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d1 : dense, d0 : compressed) }>
129// CHECK-LABEL: func private @sparse_csc(
130// CHECK-SAME: tensor<?x?xf32, #[[$CSC]]>)
131func.func private @sparse_csc(tensor<?x?xf32, #CSC>)
132
133// -----
134
135#DCSC = #sparse_tensor.encoding<{
136  map = (d0, d1) -> (d1 : compressed, d0 : compressed),
137  posWidth = 0,
138  crdWidth = 64
139}>
140
141// CHECK-DAG: #[[$DCSC:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d1 : compressed, d0 : compressed), crdWidth = 64 }>
142// CHECK-LABEL: func private @sparse_dcsc(
143// CHECK-SAME: tensor<?x?xf32, #[[$DCSC]]>)
144func.func private @sparse_dcsc(tensor<?x?xf32, #DCSC>)
145
146// -----
147
148#COO = #sparse_tensor.encoding<{
149  map = (d0, d1) -> (d0 : compressed(nonunique, nonordered), d1 : singleton(nonordered))
150}>
151
152// CHECK-DAG: #[[$COO:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : compressed(nonunique, nonordered), d1 : singleton(nonordered)) }>
153// CHECK-LABEL: func private @sparse_coo(
154// CHECK-SAME: tensor<?x?xf32, #[[$COO]]>)
155func.func private @sparse_coo(tensor<?x?xf32, #COO>)
156
157// -----
158
159#COO_DENSE = #sparse_tensor.encoding<{
160  map = (d0, d1, d2) -> (d0 : compressed(nonunique), d1 : singleton, d2: dense)
161}>
162
163// CHECK-DAG: #[[$COO:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : compressed(nonunique), d1 : singleton, d2 : dense) }>
164// CHECK-LABEL: func private @sparse_coo_trailing_dense(
165// CHECK-SAME: tensor<?x?x1xf32, #[[$COO]]>)
166func.func private @sparse_coo_trailing_dense(tensor<?x?x1xf32, #COO_DENSE>)
167
168// -----
169
170#BCOO = #sparse_tensor.encoding<{
171  map = (d0, d1, d2) -> (d0 : dense, d1 : loose_compressed(nonunique), d2 : singleton)
172}>
173
174// CHECK-DAG: #[[$BCOO:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : dense, d1 : loose_compressed(nonunique), d2 : singleton) }>
175// CHECK-LABEL: func private @sparse_bcoo(
176// CHECK-SAME: tensor<?x?x?xf32, #[[$BCOO]]>)
177func.func private @sparse_bcoo(tensor<?x?x?xf32, #BCOO>)
178
179// -----
180
181#SortedCOO = #sparse_tensor.encoding<{
182  map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)
183}>
184
185// CHECK-DAG: #[[$SortedCOO:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton) }>
186// CHECK-LABEL: func private @sparse_sorted_coo(
187// CHECK-SAME: tensor<10x10xf64, #[[$SortedCOO]]>)
188func.func private @sparse_sorted_coo(tensor<10x10xf64, #SortedCOO>)
189
190// -----
191
192#COO_SoA = #sparse_tensor.encoding<{
193  map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton(soa))
194}>
195
196// CHECK-DAG: #[[$COO_SoA:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton(soa)) }>
197// CHECK-LABEL: func private @sparse_coo(
198// CHECK-SAME: tensor<?x?xf32, #[[$COO_SoA]]>)
199func.func private @sparse_coo(tensor<?x?xf32, #COO_SoA>)
200
201// -----
202
203#BSR = #sparse_tensor.encoding<{
204   map = ( i, j ) ->
205      ( i floordiv 2 : dense,
206        j floordiv 3 : compressed,
207        i mod 2      : dense,
208        j mod 3      : dense
209      )
210}>
211
212// CHECK-DAG: #[[$BSR:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : dense, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>
213// CHECK-LABEL: func private @sparse_bsr(
214// CHECK-SAME: tensor<10x60xf64, #[[$BSR]]>
215func.func private @sparse_bsr(tensor<10x60xf64, #BSR>)
216
217
218// -----
219
220#ELL = #sparse_tensor.encoding<{
221  map = [s0](d0, d1) -> (d0 * (s0 * 4) : dense, d0 : dense, d1 : compressed)
222}>
223
224// CHECK-DAG: #[[$ELL:.*]] = #sparse_tensor.encoding<{ map = [s0](d0, d1) -> (d0 * (s0 * 4) : dense, d0 : dense, d1 : compressed) }>
225// CHECK-LABEL: func private @sparse_ell(
226// CHECK-SAME: tensor<?x?xf64, #[[$ELL]]>
227func.func private @sparse_ell(tensor<?x?xf64, #ELL>)
228
229// -----
230
231#CSR_SLICE = #sparse_tensor.encoding<{
232  map = (d0 : #sparse_tensor<slice(1, 4, 1)>, d1 : #sparse_tensor<slice(1, 4, 2)>) -> (d0 : dense, d1 : compressed)
233}>
234
235// CHECK-DAG: #[[$CSR_SLICE:.*]] = #sparse_tensor.encoding<{ map = (d0 : #sparse_tensor<slice(1, 4, 1)>, d1 : #sparse_tensor<slice(1, 4, 2)>) -> (d0 : dense, d1 : compressed) }>
236// CHECK-LABEL: func private @sparse_slice(
237// CHECK-SAME: tensor<?x?xf64, #[[$CSR_SLICE]]>
238func.func private @sparse_slice(tensor<?x?xf64, #CSR_SLICE>)
239
240// -----
241
242#CSR_SLICE = #sparse_tensor.encoding<{
243  map = (d0 : #sparse_tensor<slice(1, ?, 1)>, d1 : #sparse_tensor<slice(?, 4, 2)>) -> (d0 : dense, d1 : compressed)
244}>
245
246// CHECK-DAG: #[[$CSR_SLICE:.*]] = #sparse_tensor.encoding<{ map = (d0 : #sparse_tensor<slice(1, ?, 1)>, d1 : #sparse_tensor<slice(?, 4, 2)>) -> (d0 : dense, d1 : compressed) }>
247// CHECK-LABEL: func private @sparse_slice(
248// CHECK-SAME: tensor<?x?xf64, #[[$CSR_SLICE]]>
249func.func private @sparse_slice(tensor<?x?xf64, #CSR_SLICE>)
250
251// -----
252
253#BSR = #sparse_tensor.encoding<{
254  map = ( i, j ) ->
255  ( i floordiv 2 : dense,
256    j floordiv 3 : compressed,
257    i mod 2      : dense,
258    j mod 3      : dense
259  )
260}>
261
262// CHECK-DAG: #[[$BSR:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : dense, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>
263// CHECK-LABEL: func private @BSR(
264// CHECK-SAME: tensor<?x?xf64, #[[$BSR]]>
265func.func private @BSR(%arg0: tensor<?x?xf64, #BSR>) {
266  return
267}
268
269// -----
270
271#BCSR = #sparse_tensor.encoding<{
272  map = ( i, j, k ) ->
273  ( i floordiv 2 : dense,
274    j floordiv 3 : dense,
275    k floordiv 4 : compressed,
276    i mod 2      : dense,
277    j mod 3      : dense,
278    k mod 4      : dense
279  )
280}>
281
282// CHECK-DAG: #[[$BCSR:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 floordiv 2 : dense, d1 floordiv 3 : dense, d2 floordiv 4 : compressed, d0 mod 2 : dense, d1 mod 3 : dense, d2 mod 4 : dense) }>
283// CHECK-LABEL: func private @BCSR(
284// CHECK-SAME: tensor<?x?x?xf64, #[[$BCSR]]>
285func.func private @BCSR(%arg0: tensor<?x?x?xf64, #BCSR>) {
286  return
287}
288// -----
289
290#BSR_explicit = #sparse_tensor.encoding<{
291  map =
292  {il, jl, ii, jj}
293  ( i = il * 2 + ii,
294    j = jl * 3 + jj
295  ) ->
296  ( il = i floordiv 2 : dense,
297    jl = j floordiv 3 : compressed,
298    ii = i mod 2      : dense,
299    jj = j mod 3      : dense
300  )
301}>
302
303// CHECK-DAG: #[[$BSR_explicit:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : dense, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>
304// CHECK-LABEL: func private @BSR_explicit(
305// CHECK-SAME: tensor<?x?xf64, #[[$BSR_explicit]]>
306func.func private @BSR_explicit(%arg0: tensor<?x?xf64, #BSR_explicit>) {
307  return
308}
309
310// -----
311
312#NV_24 = #sparse_tensor.encoding<{
313  map = ( i, j ) ->
314  ( i            : dense,
315    j floordiv 4 : dense,
316    j mod 4      : structured[2, 4]
317  ),
318  crdWidth = 8  // we would even like just 2-bits
319}>
320
321// CHECK-DAG: #[[$NV_24:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 floordiv 4 : dense, d1 mod 4 : structured[2, 4]), crdWidth = 8 }>
322// CHECK-LABEL: func private @NV_24(
323// CHECK-SAME: tensor<?x?xf64, #[[$NV_24]]>
324func.func private @NV_24(%arg0: tensor<?x?xf64, #NV_24>) {
325  return
326}
327
328// -----
329
330#NV_24 = #sparse_tensor.encoding<{
331  map = ( i, j, k ) ->
332  ( i            : dense,
333    j            : dense,
334    k floordiv 4 : dense,
335    k mod 4      : structured[2, 4]
336  )
337}>
338
339// CHECK-DAG: #[[$NV_24:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : dense, d1 : dense, d2 floordiv 4 : dense, d2 mod 4 : structured[2, 4]) }>
340// CHECK-LABEL: func private @NV_24(
341// CHECK-SAME: tensor<?x?x?xf64, #[[$NV_24]]>
342func.func private @NV_24(%arg0: tensor<?x?x?xf64, #NV_24>) {
343  return
344}
345
346// -----
347
348#NV_24 = #sparse_tensor.encoding<{
349  map = ( i, j, k ) ->
350  ( i            : dense,
351    k floordiv 4 : dense,
352    j            : dense,
353    k mod 4      : structured[2, 4]
354  )
355}>
356
357// CHECK-DAG: #[[$NV_24:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : dense, d2 floordiv 4 : dense, d1 : dense, d2 mod 4 : structured[2, 4]) }>
358// CHECK-LABEL: func private @NV_24(
359// CHECK-SAME: tensor<?x?x?xf64, #[[$NV_24]]>
360func.func private @NV_24(%arg0: tensor<?x?x?xf64, #NV_24>) {
361  return
362}
363
364// -----
365
366#NOutOfM = #sparse_tensor.encoding<{
367  map = ( i, j, k ) ->
368  ( i            : dense,
369    k floordiv 8 : dense,
370    j            : dense,
371    k mod 8      : structured[5, 8]
372  )
373}>
374
375// CHECK-DAG: #[[$NOutOfM:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : dense, d2 floordiv 8 : dense, d1 : dense, d2 mod 8 : structured[5, 8]) }>
376// CHECK-LABEL: func private @NOutOfM(
377// CHECK-SAME: tensor<?x?x?xf64, #[[$NOutOfM]]>
378func.func private @NOutOfM(%arg0: tensor<?x?x?xf64, #NOutOfM>) {
379  return
380}
381