Skip to content

Commit 662d821

Browse files
aartbikaartbik
authored
[mlir][sparse] migrate datastructure tests to sparse_tensor.print (#83956)
Continuing the efforts started in #83357
1 parent d016712 commit 662d821

File tree

2 files changed

+72
-166
lines changed

2 files changed

+72
-166
lines changed

‎mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_1d.mlir

Lines changed: 23 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
1111
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
1212
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
13-
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
13+
// DEFINE: %{run_opts} = -e main -entry-point-result=void
1414
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
1515
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
1616
//
@@ -30,8 +30,6 @@
3030
// Do the same run, but now with direct IR generation and VLA vectorization.
3131
// RUN: %if mlir_arm_sve_tests %{ %{compile_sve} | %{run_sve} | FileCheck %s %}
3232

33-
// Insertion example using pure codegen (no sparse runtime support lib).
34-
3533
#SparseVector = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
3634

3735
#trait_mul_s = {
@@ -43,27 +41,7 @@
4341
}
4442

4543
module {
46-
47-
// Dumps positions, indices, values for verification.
48-
func.func @dump(%argx: tensor<1024xf32, #SparseVector>) {
49-
%c0 = arith.constant 0 : index
50-
%f0 = arith.constant 0.0 : f32
51-
%p = sparse_tensor.positions %argx { level = 0 : index }
52-
: tensor<1024xf32, #SparseVector> to memref<?xindex>
53-
%i = sparse_tensor.coordinates %argx { level = 0 : index }
54-
: tensor<1024xf32, #SparseVector> to memref<?xindex>
55-
%v = sparse_tensor.values %argx
56-
: tensor<1024xf32, #SparseVector> to memref<?xf32>
57-
%vp = vector.transfer_read %p[%c0], %c0: memref<?xindex>, vector<2xindex>
58-
%vi = vector.transfer_read %i[%c0], %c0: memref<?xindex>, vector<8xindex>
59-
%vv = vector.transfer_read %v[%c0], %f0: memref<?xf32>, vector<8xf32>
60-
vector.print %vp : vector<2xindex>
61-
vector.print %vi : vector<8xindex>
62-
vector.print %vv : vector<8xf32>
63-
return
64-
}
65-
66-
func.func @entry() {
44+
func.func @main() {
6745
%f1 = arith.constant 1.0 : f32
6846
%f2 = arith.constant 2.0 : f32
6947
%f3 = arith.constant 3.0 : f32
@@ -82,10 +60,17 @@ module {
8260
%4 = sparse_tensor.insert %f4 into %3[%c1023] : tensor<1024xf32, #SparseVector>
8361
%5 = sparse_tensor.load %4 hasInserts : tensor<1024xf32, #SparseVector>
8462

85-
// CHECK: ( 0, 4 )
86-
// CHECK-NEXT: ( 0, 1, 3, 1023
87-
// CHECK-NEXT: ( 1, 2, 3, 4
88-
call @dump(%5) : (tensor<1024xf32, #SparseVector>) -> ()
63+
//
64+
// CHECK: ---- Sparse Tensor ----
65+
// CHECK-NEXT: nse = 4
66+
// CHECK-NEXT: dim = ( 1024 )
67+
// CHECK-NEXT: lvl = ( 1024 )
68+
// CHECK-NEXT: pos[0] : ( 0, 4,
69+
// CHECK-NEXT: crd[0] : ( 0, 1, 3, 1023,
70+
// CHECK-NEXT: values : ( 1, 2, 3, 4,
71+
// CHECK-NEXT: ----
72+
//
73+
sparse_tensor.print %5 : tensor<1024xf32, #SparseVector>
8974

9075
// Build another sparse vector in a loop.
9176
%6 = tensor.empty() : tensor<1024xf32, #SparseVector>
@@ -96,18 +81,17 @@ module {
9681
}
9782
%8 = sparse_tensor.load %7 hasInserts : tensor<1024xf32, #SparseVector>
9883

99-
// CHECK-NEXT: ( 0, 8 )
100-
// CHECK-NEXT: ( 0, 3, 6, 9, 12, 15, 18, 21 )
101-
// CHECK-NEXT: ( 1, 1, 1, 1, 1, 1, 1, 1 )
10284
//
103-
call @dump(%8) : (tensor<1024xf32, #SparseVector>) -> ()
104-
105-
// CHECK-NEXT: 4
106-
// CHECK-NEXT: 8
107-
%noe1 = sparse_tensor.number_of_entries %5 : tensor<1024xf32, #SparseVector>
108-
%noe2 = sparse_tensor.number_of_entries %8 : tensor<1024xf32, #SparseVector>
109-
vector.print %noe1 : index
110-
vector.print %noe2 : index
85+
// CHECK-NEXT: ---- Sparse Tensor ----
86+
// CHECK-NEXT: nse = 8
87+
// CHECK-NEXT: dim = ( 1024 )
88+
// CHECK-NEXT: lvl = ( 1024 )
89+
// CHECK-NEXT: pos[0] : ( 0, 8,
90+
// CHECK-NEXT: crd[0] : ( 0, 3, 6, 9, 12, 15, 18, 21,
91+
// CHECK-NEXT: values : ( 1, 1, 1, 1, 1, 1, 1, 1,
92+
// CHECK-NEXT: ----
93+
//
94+
sparse_tensor.print %8 : tensor<1024xf32, #SparseVector>
11195

11296
// Free resources.
11397
bufferization.dealloc_tensor %5 : tensor<1024xf32, #SparseVector>

‎mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_2d.mlir

Lines changed: 49 additions & 127 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
1111
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
1212
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
13-
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
13+
// DEFINE: %{run_opts} = -e main -entry-point-result=void
1414
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
1515
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
1616
//
@@ -32,7 +32,7 @@
3232
}>
3333

3434
#SortedCOO = #sparse_tensor.encoding<{
35-
map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)
35+
map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton(soa))
3636
}>
3737

3838
#CSR = #sparse_tensor.encoding<{
@@ -48,94 +48,11 @@
4848
}>
4949

5050
module {
51-
52-
func.func @dump_dense(%arg0: tensor<4x3xf64, #Dense>) {
53-
%c0 = arith.constant 0 : index
54-
%fu = arith.constant 99.0 : f64
55-
%v = sparse_tensor.values %arg0 : tensor<4x3xf64, #Dense> to memref<?xf64>
56-
%vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<12xf64>
57-
vector.print %vv : vector<12xf64>
58-
return
59-
}
60-
61-
func.func @dump_coo(%arg0: tensor<4x3xf64, #SortedCOO>) {
62-
%c0 = arith.constant 0 : index
63-
%cu = arith.constant -1 : index
64-
%fu = arith.constant 99.0 : f64
65-
%p0 = sparse_tensor.positions %arg0 { level = 0 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex>
66-
%i0 = sparse_tensor.coordinates %arg0 { level = 0 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
67-
%i1 = sparse_tensor.coordinates %arg0 { level = 1 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
68-
%v = sparse_tensor.values %arg0 : tensor<4x3xf64, #SortedCOO> to memref<?xf64>
69-
%vp0 = vector.transfer_read %p0[%c0], %cu: memref<?xindex>, vector<2xindex>
70-
vector.print %vp0 : vector<2xindex>
71-
%vi0 = vector.transfer_read %i0[%c0], %cu: memref<?xindex, strided<[?], offset: ?>>, vector<4xindex>
72-
vector.print %vi0 : vector<4xindex>
73-
%vi1 = vector.transfer_read %i1[%c0], %cu: memref<?xindex, strided<[?], offset: ?>>, vector<4xindex>
74-
vector.print %vi1 : vector<4xindex>
75-
%vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<4xf64>
76-
vector.print %vv : vector<4xf64>
77-
return
78-
}
79-
80-
func.func @dump_csr(%arg0: tensor<4x3xf64, #CSR>) {
81-
%c0 = arith.constant 0 : index
82-
%cu = arith.constant -1 : index
83-
%fu = arith.constant 99.0 : f64
84-
%p1 = sparse_tensor.positions %arg0 { level = 1 : index } : tensor<4x3xf64, #CSR> to memref<?xindex>
85-
%i1 = sparse_tensor.coordinates %arg0 { level = 1 : index } : tensor<4x3xf64, #CSR> to memref<?xindex>
86-
%v = sparse_tensor.values %arg0 : tensor<4x3xf64, #CSR> to memref<?xf64>
87-
%vp1 = vector.transfer_read %p1[%c0], %cu: memref<?xindex>, vector<5xindex>
88-
vector.print %vp1 : vector<5xindex>
89-
%vi1 = vector.transfer_read %i1[%c0], %cu: memref<?xindex>, vector<4xindex>
90-
vector.print %vi1 : vector<4xindex>
91-
%vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<4xf64>
92-
vector.print %vv : vector<4xf64>
93-
return
94-
}
95-
96-
func.func @dump_dcsr(%arg0: tensor<4x3xf64, #DCSR>) {
97-
%c0 = arith.constant 0 : index
98-
%cu = arith.constant -1 : index
99-
%fu = arith.constant 99.0 : f64
100-
%p0 = sparse_tensor.positions %arg0 { level = 0 : index } : tensor<4x3xf64, #DCSR> to memref<?xindex>
101-
%i0 = sparse_tensor.coordinates %arg0 { level = 0 : index } : tensor<4x3xf64, #DCSR> to memref<?xindex>
102-
%p1 = sparse_tensor.positions %arg0 { level = 1 : index } : tensor<4x3xf64, #DCSR> to memref<?xindex>
103-
%i1 = sparse_tensor.coordinates %arg0 { level = 1 : index } : tensor<4x3xf64, #DCSR> to memref<?xindex>
104-
%v = sparse_tensor.values %arg0 : tensor<4x3xf64, #DCSR> to memref<?xf64>
105-
%vp0 = vector.transfer_read %p0[%c0], %cu: memref<?xindex>, vector<2xindex>
106-
vector.print %vp0 : vector<2xindex>
107-
%vi0 = vector.transfer_read %i0[%c0], %cu: memref<?xindex>, vector<3xindex>
108-
vector.print %vi0 : vector<3xindex>
109-
%vp1 = vector.transfer_read %p1[%c0], %cu: memref<?xindex>, vector<4xindex>
110-
vector.print %vp1 : vector<4xindex>
111-
%vi1 = vector.transfer_read %i1[%c0], %cu: memref<?xindex>, vector<4xindex>
112-
vector.print %vi1 : vector<4xindex>
113-
%vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<4xf64>
114-
vector.print %vv : vector<4xf64>
115-
return
116-
}
117-
118-
func.func @dump_row(%arg0: tensor<4x3xf64, #Row>) {
119-
%c0 = arith.constant 0 : index
120-
%cu = arith.constant -1 : index
121-
%fu = arith.constant 99.0 : f64
122-
%p0 = sparse_tensor.positions %arg0 { level = 0 : index } : tensor<4x3xf64, #Row> to memref<?xindex>
123-
%i0 = sparse_tensor.coordinates %arg0 { level = 0 : index } : tensor<4x3xf64, #Row> to memref<?xindex>
124-
%v = sparse_tensor.values %arg0 : tensor<4x3xf64, #Row> to memref<?xf64>
125-
%vp0 = vector.transfer_read %p0[%c0], %cu: memref<?xindex>, vector<2xindex>
126-
vector.print %vp0 : vector<2xindex>
127-
%vi0 = vector.transfer_read %i0[%c0], %cu: memref<?xindex>, vector<3xindex>
128-
vector.print %vi0 : vector<3xindex>
129-
%vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<9xf64>
130-
vector.print %vv : vector<9xf64>
131-
return
132-
}
133-
13451
//
13552
// Main driver. We test the contents of various sparse tensor
13653
// schemes when they are still empty and after a few insertions.
13754
//
138-
func.func @entry() {
55+
func.func @main() {
13956
%c0 = arith.constant 0 : index
14057
%c2 = arith.constant 2 : index
14158
%c3 = arith.constant 3 : index
@@ -147,98 +64,103 @@ module {
14764
//
14865
// Dense case.
14966
//
150-
// CHECK: ( 1, 0, 0, 0, 0, 0, 0, 0, 2, 3, 0, 4 )
67+
// CHECK: ---- Sparse Tensor ----
68+
// CHECK-NEXT: nse = 12
69+
// CHECK-NEXT: dim = ( 4, 3 )
70+
// CHECK-NEXT: lvl = ( 4, 3 )
71+
// CHECK-NEXT: values : ( 1, 0, 0, 0, 0, 0, 0, 0, 2, 3, 0, 4,
72+
// CHECK-NEXT: ----
15173
//
15274
%densea = tensor.empty() : tensor<4x3xf64, #Dense>
15375
%dense1 = sparse_tensor.insert %f1 into %densea[%c0, %c0] : tensor<4x3xf64, #Dense>
15476
%dense2 = sparse_tensor.insert %f2 into %dense1[%c2, %c2] : tensor<4x3xf64, #Dense>
15577
%dense3 = sparse_tensor.insert %f3 into %dense2[%c3, %c0] : tensor<4x3xf64, #Dense>
15678
%dense4 = sparse_tensor.insert %f4 into %dense3[%c3, %c2] : tensor<4x3xf64, #Dense>
15779
%densem = sparse_tensor.load %dense4 hasInserts : tensor<4x3xf64, #Dense>
158-
call @dump_dense(%densem) : (tensor<4x3xf64, #Dense>) -> ()
80+
sparse_tensor.print %densem : tensor<4x3xf64, #Dense>
15981

16082
//
16183
// COO case.
16284
//
163-
// CHECK-NEXT: ( 0, 4 )
164-
// CHECK-NEXT: ( 0, 2, 3, 3 )
165-
// CHECK-NEXT: ( 0, 2, 0, 2 )
166-
// CHECK-NEXT: ( 1, 2, 3, 4 )
85+
// CHECK-NEXT: ---- Sparse Tensor ----
86+
// CHECK-NEXT: nse = 4
87+
// CHECK-NEXT: dim = ( 4, 3 )
88+
// CHECK-NEXT: lvl = ( 4, 3 )
89+
// CHECK-NEXT: pos[0] : ( 0, 4,
90+
// CHECK-NEXT: crd[0] : ( 0, 2, 3, 3,
91+
// CHECK-NEXT: crd[1] : ( 0, 2, 0, 2,
92+
// CHECK-NEXT: values : ( 1, 2, 3, 4,
93+
// CHECK-NEXT: ----
16794
//
16895
%cooa = tensor.empty() : tensor<4x3xf64, #SortedCOO>
16996
%coo1 = sparse_tensor.insert %f1 into %cooa[%c0, %c0] : tensor<4x3xf64, #SortedCOO>
17097
%coo2 = sparse_tensor.insert %f2 into %coo1[%c2, %c2] : tensor<4x3xf64, #SortedCOO>
17198
%coo3 = sparse_tensor.insert %f3 into %coo2[%c3, %c0] : tensor<4x3xf64, #SortedCOO>
17299
%coo4 = sparse_tensor.insert %f4 into %coo3[%c3, %c2] : tensor<4x3xf64, #SortedCOO>
173100
%coom = sparse_tensor.load %coo4 hasInserts : tensor<4x3xf64, #SortedCOO>
174-
call @dump_coo(%coom) : (tensor<4x3xf64, #SortedCOO>) -> ()
101+
sparse_tensor.print %coom : tensor<4x3xf64, #SortedCOO>
175102

176103
//
177104
// CSR case.
178105
//
179-
// CHECK-NEXT: ( 0, 1, 1, 2, 4 )
180-
// CHECK-NEXT: ( 0, 2, 0, 2 )
181-
// CHECK-NEXT: ( 1, 2, 3, 4 )
106+
// CHECK-NEXT: ---- Sparse Tensor ----
107+
// CHECK-NEXT: nse = 4
108+
// CHECK-NEXT: dim = ( 4, 3 )
109+
// CHECK-NEXT: lvl = ( 4, 3 )
110+
// CHECK-NEXT: pos[1] : ( 0, 1, 1, 2, 4,
111+
// CHECK-NEXT: crd[1] : ( 0, 2, 0, 2,
112+
// CHECK-NEXT: values : ( 1, 2, 3, 4,
113+
// CHECK-NEXT: ----
182114
//
183115
%csra = tensor.empty() : tensor<4x3xf64, #CSR>
184116
%csr1 = sparse_tensor.insert %f1 into %csra[%c0, %c0] : tensor<4x3xf64, #CSR>
185117
%csr2 = sparse_tensor.insert %f2 into %csr1[%c2, %c2] : tensor<4x3xf64, #CSR>
186118
%csr3 = sparse_tensor.insert %f3 into %csr2[%c3, %c0] : tensor<4x3xf64, #CSR>
187119
%csr4 = sparse_tensor.insert %f4 into %csr3[%c3, %c2] : tensor<4x3xf64, #CSR>
188120
%csrm = sparse_tensor.load %csr4 hasInserts : tensor<4x3xf64, #CSR>
189-
call @dump_csr(%csrm) : (tensor<4x3xf64, #CSR>) -> ()
121+
sparse_tensor.print %csrm : tensor<4x3xf64, #CSR>
190122

191123
//
192124
// DCSR case.
193125
//
194-
// CHECK-NEXT: ( 0, 3 )
195-
// CHECK-NEXT: ( 0, 2, 3 )
196-
// CHECK-NEXT: ( 0, 1, 2, 4 )
197-
// CHECK-NEXT: ( 0, 2, 0, 2 )
198-
// CHECK-NEXT: ( 1, 2, 3, 4 )
126+
// CHECK-NEXT: ---- Sparse Tensor ----
127+
// CHECK-NEXT: nse = 4
128+
// CHECK-NEXT: dim = ( 4, 3 )
129+
// CHECK-NEXT: lvl = ( 4, 3 )
130+
// CHECK-NEXT: pos[0] : ( 0, 3,
131+
// CHECK-NEXT: crd[0] : ( 0, 2, 3,
132+
// CHECK-NEXT: pos[1] : ( 0, 1, 2, 4,
133+
// CHECK-NEXT: crd[1] : ( 0, 2, 0, 2,
134+
// CHECK-NEXT: values : ( 1, 2, 3, 4,
135+
// CHECK-NEXT: ----
199136
//
200137
%dcsra = tensor.empty() : tensor<4x3xf64, #DCSR>
201138
%dcsr1 = sparse_tensor.insert %f1 into %dcsra[%c0, %c0] : tensor<4x3xf64, #DCSR>
202139
%dcsr2 = sparse_tensor.insert %f2 into %dcsr1[%c2, %c2] : tensor<4x3xf64, #DCSR>
203140
%dcsr3 = sparse_tensor.insert %f3 into %dcsr2[%c3, %c0] : tensor<4x3xf64, #DCSR>
204141
%dcsr4 = sparse_tensor.insert %f4 into %dcsr3[%c3, %c2] : tensor<4x3xf64, #DCSR>
205142
%dcsrm = sparse_tensor.load %dcsr4 hasInserts : tensor<4x3xf64, #DCSR>
206-
call @dump_dcsr(%dcsrm) : (tensor<4x3xf64, #DCSR>) -> ()
143+
sparse_tensor.print %dcsrm : tensor<4x3xf64, #DCSR>
207144

208145
//
209146
// Row case.
210147
//
211-
// CHECK-NEXT: ( 0, 3 )
212-
// CHECK-NEXT: ( 0, 2, 3 )
213-
// CHECK-NEXT: ( 1, 0, 0, 0, 0, 2, 3, 0, 4 )
148+
// CHECK-NEXT: ---- Sparse Tensor ----
149+
// CHECK-NEXT: nse = 9
150+
// CHECK-NEXT: dim = ( 4, 3 )
151+
// CHECK-NEXT: lvl = ( 4, 3 )
152+
// CHECK-NEXT: pos[0] : ( 0, 3,
153+
// CHECK-NEXT: crd[0] : ( 0, 2, 3,
154+
// CHECK-NEXT: values : ( 1, 0, 0, 0, 0, 2, 3, 0, 4,
155+
// CHECK-NEXT: ----
214156
//
215157
%rowa = tensor.empty() : tensor<4x3xf64, #Row>
216158
%row1 = sparse_tensor.insert %f1 into %rowa[%c0, %c0] : tensor<4x3xf64, #Row>
217159
%row2 = sparse_tensor.insert %f2 into %row1[%c2, %c2] : tensor<4x3xf64, #Row>
218160
%row3 = sparse_tensor.insert %f3 into %row2[%c3, %c0] : tensor<4x3xf64, #Row>
219161
%row4 = sparse_tensor.insert %f4 into %row3[%c3, %c2] : tensor<4x3xf64, #Row>
220162
%rowm = sparse_tensor.load %row4 hasInserts : tensor<4x3xf64, #Row>
221-
call @dump_row(%rowm) : (tensor<4x3xf64, #Row>) -> ()
222-
223-
//
224-
// NOE sanity check.
225-
//
226-
// CHECK-NEXT: 12
227-
// CHECK-NEXT: 4
228-
// CHECK-NEXT: 4
229-
// CHECK-NEXT: 4
230-
// CHECK-NEXT: 9
231-
//
232-
%noe1 = sparse_tensor.number_of_entries %densem : tensor<4x3xf64, #Dense>
233-
%noe2 = sparse_tensor.number_of_entries %coom : tensor<4x3xf64, #SortedCOO>
234-
%noe3 = sparse_tensor.number_of_entries %csrm : tensor<4x3xf64, #CSR>
235-
%noe4 = sparse_tensor.number_of_entries %dcsrm : tensor<4x3xf64, #DCSR>
236-
%noe5 = sparse_tensor.number_of_entries %rowm : tensor<4x3xf64, #Row>
237-
vector.print %noe1 : index
238-
vector.print %noe2 : index
239-
vector.print %noe3 : index
240-
vector.print %noe4 : index
241-
vector.print %noe5 : index
163+
sparse_tensor.print %rowm : tensor<4x3xf64, #Row>
242164

243165
// Release resources.
244166
bufferization.dealloc_tensor %densem : tensor<4x3xf64, #Dense>

0 commit comments

Comments
 (0)