Skip to content

Commit 813aaf3

Browse files
aartbikaartbik
authored
[mlir][sparse] stress test BSR (#72712)
I always enjoy a good stress test. This end-to-end integration test ensures the major ordering of both the block and within the block are correctly handled (giving row-row, row-col, col-row and col-row as options).
1 parent 635756e commit 813aaf3

File tree

1 file changed

+178
-0
lines changed

1 file changed

+178
-0
lines changed

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

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
//--------------------------------------------------------------------------------------------------
2+
// WHEN CREATING A NEW TEST, PLEASE JUST COPY & PASTE WITHOUT EDITS.
3+
//
4+
// Set-up that's shared across all tests in this directory. In principle, this
5+
// config could be moved to lit.local.cfg. However, there are downstream users that
6+
// do not use these LIT config files. Hence why this is kept inline.
7+
//
8+
// DEFINE: %{sparsifier_opts} = enable-runtime-library=true
9+
// DEFINE: %{sparsifier_opts_sve} = enable-arm-sve=true %{sparsifier_opts}
10+
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
11+
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
12+
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
13+
// DEFINE: %{run_opts} = -e main -entry-point-result=void
14+
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
15+
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
16+
//
17+
// DEFINE: %{env} =
18+
//--------------------------------------------------------------------------------------------------
19+
20+
// RUN: %{compile} | %{run} | FileCheck %s
21+
//
22+
// Do the same run, but now with direct IR generation.
23+
// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false
24+
// RUN: %{compile} | %{run} | FileCheck %s
25+
//
26+
// Do the same run, but now with direct IR generation and vectorization.
27+
// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false enable-buffer-initialization=true vl=2 reassociate-fp-reductions=true enable-index-optimizations=true
28+
// RUN: %{compile} | %{run} | FileCheck %s
29+
30+
#BSR_row_rowmajor = #sparse_tensor.encoding<{
31+
map = (i, j) ->
32+
( i floordiv 3 : dense
33+
, j floordiv 4 : compressed
34+
, i mod 3 : dense
35+
, j mod 4 : dense
36+
)
37+
}>
38+
39+
#BSR_row_colmajor = #sparse_tensor.encoding<{
40+
map = (i, j) ->
41+
( i floordiv 3 : dense
42+
, j floordiv 4 : compressed
43+
, j mod 4 : dense
44+
, i mod 3 : dense
45+
)
46+
}>
47+
48+
#BSR_col_rowmajor = #sparse_tensor.encoding<{
49+
map = (i, j) ->
50+
( j floordiv 4 : dense
51+
, i floordiv 3 : compressed
52+
, i mod 3 : dense
53+
, j mod 4 : dense
54+
)
55+
}>
56+
57+
#BSR_col_colmajor = #sparse_tensor.encoding<{
58+
map = (i, j) ->
59+
( j floordiv 4 : dense
60+
, i floordiv 3 : compressed
61+
, j mod 4 : dense
62+
, i mod 3 : dense
63+
)
64+
}>
65+
66+
//
67+
// Example 3x4 block storage of a 6x16 matrix:
68+
//
69+
// +---------+---------+---------+---------+
70+
// | 1 2 . . | . . . . | . . . . | . . . . |
71+
// | . . . . | . . . . | . . . . | . . . . |
72+
// | . . . 3 | . . . . | . . . . | . . . . |
73+
// +---------+---------+---------+---------+
74+
// | . . . . | . . . . | 4 5 . . | . . . . |
75+
// | . . . . | . . . . | . . . . | . . . . |
76+
// | . . . . | . . . . | . . 6 7 | . . . . |
77+
// +---------+---------+---------+---------+
78+
//
79+
// Storage for CSR block storage. Note that this essentially
80+
// provides CSR storage of 2x4 blocks with either row-major
81+
// or column-major storage within each 3x4 block of elements.
82+
//
83+
// positions[1] : 0 1 2
84+
// coordinates[1] : 0 2
85+
// values : 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3,
86+
// 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 6, 7 [row-major]
87+
//
88+
// 1, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3,
89+
// 4, 0, 0, 5, 0, 0, 0, 0, 6, 0, 0, 7 [col-major]
90+
//
91+
// Storage for CSC block storage. Note that this essentially
92+
// provides CSC storage of 4x2 blocks with either row-major
93+
// or column-major storage within each 3x4 block of elements.
94+
//
95+
// positions[1] : 0 1 1 2 2
96+
// coordinates[1] : 0 1
97+
// values : 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3,
98+
// 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 6, 7 [row-major]
99+
//
100+
// 1, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3,
101+
// 4, 0, 0, 5, 0, 0, 0, 0, 6, 0, 0, 7 [col-major]
102+
//
103+
module {
104+
105+
func.func @main() {
106+
%c0 = arith.constant 0 : index
107+
%f0 = arith.constant 0.0 : f64
108+
109+
%m = arith.constant sparse<
110+
[ [0, 0], [0, 1], [2, 3], [3, 8], [3, 9], [5, 10], [5, 11] ],
111+
[ 1., 2., 3., 4., 5., 6., 7.]
112+
> : tensor<6x16xf64>
113+
%s1 = sparse_tensor.convert %m : tensor<6x16xf64> to tensor<?x?xf64, #BSR_row_rowmajor>
114+
%s2 = sparse_tensor.convert %m : tensor<6x16xf64> to tensor<?x?xf64, #BSR_row_colmajor>
115+
%s3 = sparse_tensor.convert %m : tensor<6x16xf64> to tensor<?x?xf64, #BSR_col_rowmajor>
116+
%s4 = sparse_tensor.convert %m : tensor<6x16xf64> to tensor<?x?xf64, #BSR_col_colmajor>
117+
118+
// CHECK: ( 0, 1, 2 )
119+
// CHECK-NEXT: ( 0, 2 )
120+
// CHECK-NEXT: ( 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 6, 7 )
121+
%pos1 = sparse_tensor.positions %s1 {level = 1 : index } : tensor<?x?xf64, #BSR_row_rowmajor> to memref<?xindex>
122+
%vecp1 = vector.transfer_read %pos1[%c0], %c0 : memref<?xindex>, vector<3xindex>
123+
vector.print %vecp1 : vector<3xindex>
124+
%crd1 = sparse_tensor.coordinates %s1 {level = 1 : index } : tensor<?x?xf64, #BSR_row_rowmajor> to memref<?xindex>
125+
%vecc1 = vector.transfer_read %crd1[%c0], %c0 : memref<?xindex>, vector<2xindex>
126+
vector.print %vecc1 : vector<2xindex>
127+
%val1 = sparse_tensor.values %s1 : tensor<?x?xf64, #BSR_row_rowmajor> to memref<?xf64>
128+
%vecv1 = vector.transfer_read %val1[%c0], %f0 : memref<?xf64>, vector<24xf64>
129+
vector.print %vecv1 : vector<24xf64>
130+
131+
// CHECK-NEXT: ( 0, 1, 2 )
132+
// CHECK-NEXT: ( 0, 2 )
133+
// CHECK-NEXT: ( 1, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 0, 5, 0, 0, 0, 0, 6, 0, 0, 7 )
134+
%pos2 = sparse_tensor.positions %s2 {level = 1 : index } : tensor<?x?xf64, #BSR_row_colmajor> to memref<?xindex>
135+
%vecp2 = vector.transfer_read %pos2[%c0], %c0 : memref<?xindex>, vector<3xindex>
136+
vector.print %vecp2 : vector<3xindex>
137+
%crd2 = sparse_tensor.coordinates %s2 {level = 1 : index } : tensor<?x?xf64, #BSR_row_colmajor> to memref<?xindex>
138+
%vecc2 = vector.transfer_read %crd2[%c0], %c0 : memref<?xindex>, vector<2xindex>
139+
vector.print %vecc2 : vector<2xindex>
140+
%val2 = sparse_tensor.values %s2 : tensor<?x?xf64, #BSR_row_colmajor> to memref<?xf64>
141+
%vecv2 = vector.transfer_read %val2[%c0], %f0 : memref<?xf64>, vector<24xf64>
142+
vector.print %vecv2 : vector<24xf64>
143+
144+
// CHECK-NEXT: ( 0, 1, 1, 2, 2 )
145+
// CHECK-NEXT: ( 0, 1 )
146+
// CHECK-NEXT: ( 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 6, 7 )
147+
%pos3 = sparse_tensor.positions %s3 {level = 1 : index } : tensor<?x?xf64, #BSR_col_rowmajor> to memref<?xindex>
148+
%vecp3 = vector.transfer_read %pos3[%c0], %c0 : memref<?xindex>, vector<5xindex>
149+
vector.print %vecp3 : vector<5xindex>
150+
%crd3 = sparse_tensor.coordinates %s3 {level = 1 : index } : tensor<?x?xf64, #BSR_col_rowmajor> to memref<?xindex>
151+
%vecc3 = vector.transfer_read %crd3[%c0], %c0 : memref<?xindex>, vector<2xindex>
152+
vector.print %vecc3 : vector<2xindex>
153+
%val3 = sparse_tensor.values %s3 : tensor<?x?xf64, #BSR_col_rowmajor> to memref<?xf64>
154+
%vecv3 = vector.transfer_read %val3[%c0], %f0 : memref<?xf64>, vector<24xf64>
155+
vector.print %vecv3 : vector<24xf64>
156+
157+
// CHECK-NEXT: ( 0, 1, 1, 2, 2 )
158+
// CHECK-NEXT: ( 0, 1 )
159+
// CHECK-NEXT: ( 1, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 0, 5, 0, 0, 0, 0, 6, 0, 0, 7 )
160+
%pos4 = sparse_tensor.positions %s4 {level = 1 : index } : tensor<?x?xf64, #BSR_col_colmajor> to memref<?xindex>
161+
%vecp4 = vector.transfer_read %pos4[%c0], %c0 : memref<?xindex>, vector<5xindex>
162+
vector.print %vecp4 : vector<5xindex>
163+
%crd4 = sparse_tensor.coordinates %s4 {level = 1 : index } : tensor<?x?xf64, #BSR_col_colmajor> to memref<?xindex>
164+
%vecc4 = vector.transfer_read %crd4[%c0], %c0 : memref<?xindex>, vector<2xindex>
165+
vector.print %vecc4 : vector<2xindex>
166+
%val4 = sparse_tensor.values %s4 : tensor<?x?xf64, #BSR_col_colmajor> to memref<?xf64>
167+
%vecv4 = vector.transfer_read %val4[%c0], %f0 : memref<?xf64>, vector<24xf64>
168+
vector.print %vecv4 : vector<24xf64>
169+
170+
// Release the resources.
171+
bufferization.dealloc_tensor %s1: tensor<?x?xf64, #BSR_row_rowmajor>
172+
bufferization.dealloc_tensor %s2: tensor<?x?xf64, #BSR_row_colmajor>
173+
bufferization.dealloc_tensor %s3: tensor<?x?xf64, #BSR_col_rowmajor>
174+
bufferization.dealloc_tensor %s4: tensor<?x?xf64, #BSR_col_colmajor>
175+
176+
return
177+
}
178+
}

0 commit comments

Comments
 (0)