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  \
6// RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-std-to-llvm | \
7// RUN: mlir-cpu-runner \
8// RUN:  -e entry -entry-point-result=void  \
9// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
10// RUN: FileCheck %s
11
12//
13// Several common sparse storage schemes.
14//
15
16#Dense  = #sparse_tensor.encoding<{
17  dimLevelType = [ "dense", "dense" ]
18}>
19
20#CSR  = #sparse_tensor.encoding<{
21  dimLevelType = [ "dense", "compressed" ]
22}>
23
24#DCSR = #sparse_tensor.encoding<{
25  dimLevelType = [ "compressed", "compressed" ]
26}>
27
28#CSC = #sparse_tensor.encoding<{
29  dimLevelType = [ "dense", "compressed" ],
30  dimOrdering = affine_map<(i,j) -> (j,i)>
31}>
32
33#DCSC = #sparse_tensor.encoding<{
34  dimLevelType = [ "compressed", "compressed" ],
35  dimOrdering = affine_map<(i,j) -> (j,i)>
36}>
37
38//
39// Integration test that looks "under the hood" of sparse storage schemes.
40//
41module {
42  //
43  // Main driver that initializes a sparse tensor and inspects the sparse
44  // storage schemes in detail. Note that users of the MLIR sparse compiler
45  // are typically not concerned with such details, but the test ensures
46  // everything is working "under the hood".
47  //
48  func @entry() {
49    %c0 = constant 0 : index
50    %c1 = constant 1 : index
51    %d0 = constant 0.0 : f64
52
53    //
54    // Initialize a dense tensor.
55    //
56    %t = constant dense<[
57       [ 1.0,  0.0,  2.0,  0.0,  0.0,  0.0,  0.0,  3.0],
58       [ 0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0],
59       [ 0.0,  0.0,  4.0,  0.0,  0.0,  0.0,  0.0,  0.0],
60       [ 0.0,  0.0,  0.0,  5.0,  0.0,  0.0,  0.0,  0.0],
61       [ 0.0,  0.0,  0.0,  0.0,  6.0,  0.0,  0.0,  0.0],
62       [ 0.0,  7.0,  8.0,  0.0,  0.0,  0.0,  0.0,  9.0],
63       [ 0.0,  0.0, 10.0,  0.0,  0.0,  0.0, 11.0, 12.0],
64       [ 0.0, 13.0, 14.0,  0.0,  0.0,  0.0, 15.0, 16.0],
65       [ 0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0],
66       [ 0.0,  0.0,  0.0,  0.0,  0.0,  0.0, 17.0,  0.0]
67    ]> : tensor<10x8xf64>
68
69    //
70    // Convert dense tensor to various sparse tensors.
71    //
72    %0 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #Dense>
73    %1 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #CSR>
74    %2 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #DCSR>
75    %3 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #CSC>
76    %4 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #DCSC>
77
78    //
79    // Inspect storage scheme of Dense.
80    //
81    // CHECK:    ( 1, 0, 2, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0,
82    // CHECK-SAME: 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0,
83    // CHECK-SAME: 0, 0, 0, 0, 6, 0, 0, 0, 0, 7, 8, 0, 0, 0, 0, 9,
84    // CHECK-SAME: 0, 0, 10, 0, 0, 0, 11, 12, 0, 13, 14, 0, 0, 0, 15, 16,
85    // CHECK-SAME: 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 17, 0 )
86    //
87    %5 = sparse_tensor.values %0 : tensor<10x8xf64, #Dense> to memref<?xf64>
88    %6 = vector.transfer_read %5[%c0], %d0: memref<?xf64>, vector<80xf64>
89    vector.print %6 : vector<80xf64>
90
91    //
92    // Inspect storage scheme of CSR.
93    //
94    // pointers(1)
95    // indices(1)
96    // values
97    //
98    // CHECK: ( 0, 3, 3, 4, 5, 6, 9, 12, 16, 16, 17 )
99    // CHECK: ( 0, 2, 7, 2, 3, 4, 1, 2, 7, 2, 6, 7, 1, 2, 6, 7, 6 )
100    // CHECK: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17 )
101    //
102    %7 = sparse_tensor.pointers %1, %c1 : tensor<10x8xf64, #CSR> to memref<?xindex>
103    %8 = vector.transfer_read %7[%c0], %c0: memref<?xindex>, vector<11xindex>
104    vector.print %8 : vector<11xindex>
105    %9 = sparse_tensor.indices %1, %c1 : tensor<10x8xf64, #CSR> to memref<?xindex>
106    %10 = vector.transfer_read %9[%c0], %c0: memref<?xindex>, vector<17xindex>
107    vector.print %10 : vector<17xindex>
108    %11 = sparse_tensor.values %1 : tensor<10x8xf64, #CSR> to memref<?xf64>
109    %12 = vector.transfer_read %11[%c0], %d0: memref<?xf64>, vector<17xf64>
110    vector.print %12 : vector<17xf64>
111
112    //
113    // Inspect storage scheme of DCSR.
114    //
115    // pointers(0)
116    // indices(0)
117    // pointers(1)
118    // indices(1)
119    // values
120    //
121    // CHECK: ( 0, 8 )
122    // CHECK: ( 0, 2, 3, 4, 5, 6, 7, 9 )
123    // CHECK: ( 0, 3, 4, 5, 6, 9, 12, 16, 17 )
124    // CHECK: ( 0, 2, 7, 2, 3, 4, 1, 2, 7, 2, 6, 7, 1, 2, 6, 7, 6 )
125    // CHECK: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17 )
126    //
127    %13 = sparse_tensor.pointers %2, %c0 : tensor<10x8xf64, #DCSR> to memref<?xindex>
128    %14 = vector.transfer_read %13[%c0], %c0: memref<?xindex>, vector<2xindex>
129    vector.print %14 : vector<2xindex>
130    %15 = sparse_tensor.indices %2, %c0 : tensor<10x8xf64, #DCSR> to memref<?xindex>
131    %16 = vector.transfer_read %15[%c0], %c0: memref<?xindex>, vector<8xindex>
132    vector.print %16 : vector<8xindex>
133    %17 = sparse_tensor.pointers %2, %c1 : tensor<10x8xf64, #DCSR> to memref<?xindex>
134    %18 = vector.transfer_read %17[%c0], %c0: memref<?xindex>, vector<9xindex>
135    vector.print %18 : vector<9xindex>
136    %19 = sparse_tensor.indices %2, %c1 : tensor<10x8xf64, #DCSR> to memref<?xindex>
137    %20 = vector.transfer_read %19[%c0], %c0: memref<?xindex>, vector<17xindex>
138    vector.print %20 : vector<17xindex>
139    %21 = sparse_tensor.values %2 : tensor<10x8xf64, #DCSR> to memref<?xf64>
140    %22 = vector.transfer_read %21[%c0], %d0: memref<?xf64>, vector<17xf64>
141    vector.print %22 : vector<17xf64>
142
143    //
144    // Inspect storage scheme of CSC.
145    //
146    // pointers(1)
147    // indices(1)
148    // values
149    //
150    // CHECK: ( 0, 1, 3, 8, 9, 10, 10, 13, 17 )
151    // CHECK: ( 0, 5, 7, 0, 2, 5, 6, 7, 3, 4, 6, 7, 9, 0, 5, 6, 7 )
152    // CHECK: ( 1, 7, 13, 2, 4, 8, 10, 14, 5, 6, 11, 15, 17, 3, 9, 12, 16 )
153    //
154    %23 = sparse_tensor.pointers %3, %c1 : tensor<10x8xf64, #CSC> to memref<?xindex>
155    %24 = vector.transfer_read %23[%c0], %c0: memref<?xindex>, vector<9xindex>
156    vector.print %24 : vector<9xindex>
157    %25 = sparse_tensor.indices %3, %c1 : tensor<10x8xf64, #CSC> to memref<?xindex>
158    %26 = vector.transfer_read %25[%c0], %c0: memref<?xindex>, vector<17xindex>
159    vector.print %26 : vector<17xindex>
160    %27 = sparse_tensor.values %3 : tensor<10x8xf64, #CSC> to memref<?xf64>
161    %28 = vector.transfer_read %27[%c0], %d0: memref<?xf64>, vector<17xf64>
162    vector.print %28 : vector<17xf64>
163
164    //
165    // Inspect storage scheme of DCSC.
166    //
167    // pointers(0)
168    // indices(0)
169    // pointers(1)
170    // indices(1)
171    // values
172    //
173    // CHECK: ( 0, 7 )
174    // CHECK: ( 0, 1, 2, 3, 4, 6, 7 )
175    // CHECK: ( 0, 1, 3, 8, 9, 10, 13, 17 )
176    // CHECK: ( 0, 5, 7, 0, 2, 5, 6, 7, 3, 4, 6, 7, 9, 0, 5, 6, 7 )
177    // CHECK: ( 1, 7, 13, 2, 4, 8, 10, 14, 5, 6, 11, 15, 17, 3, 9, 12, 16 )
178    //
179    %29 = sparse_tensor.pointers %4, %c0 : tensor<10x8xf64, #DCSC> to memref<?xindex>
180    %30 = vector.transfer_read %29[%c0], %c0: memref<?xindex>, vector<2xindex>
181    vector.print %30 : vector<2xindex>
182    %31 = sparse_tensor.indices %4, %c0 : tensor<10x8xf64, #DCSC> to memref<?xindex>
183    %32 = vector.transfer_read %31[%c0], %c0: memref<?xindex>, vector<7xindex>
184    vector.print %32 : vector<7xindex>
185    %33 = sparse_tensor.pointers %4, %c1 : tensor<10x8xf64, #DCSC> to memref<?xindex>
186    %34 = vector.transfer_read %33[%c0], %c0: memref<?xindex>, vector<8xindex>
187    vector.print %34 : vector<8xindex>
188    %35 = sparse_tensor.indices %4, %c1 : tensor<10x8xf64, #DCSC> to memref<?xindex>
189    %36 = vector.transfer_read %35[%c0], %c0: memref<?xindex>, vector<17xindex>
190    vector.print %36 : vector<17xindex>
191    %37 = sparse_tensor.values %4 : tensor<10x8xf64, #DCSC> to memref<?xf64>
192    %38 = vector.transfer_read %37[%c0], %d0: memref<?xf64>, vector<17xf64>
193    vector.print %38 : vector<17xf64>
194
195    return
196  }
197}
198