1// RUN: mlir-opt %s --sparse-compiler | \
2// RUN: mlir-cpu-runner \
3// RUN:  -e entry -entry-point-result=void  \
4// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
5// RUN: FileCheck %s
6
7#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
8
9//
10// Traits for 2-d tensor (aka matrix) operations.
11//
12#trait_scale = {
13  indexing_maps = [
14    affine_map<(i,j) -> (i,j)>,  // A (in)
15    affine_map<(i,j) -> (i,j)>   // X (out)
16  ],
17  iterator_types = ["parallel", "parallel"],
18  doc = "X(i,j) = A(i,j) * 2.0"
19}
20#trait_scale_inpl = {
21  indexing_maps = [
22    affine_map<(i,j) -> (i,j)>   // X (out)
23  ],
24  iterator_types = ["parallel", "parallel"],
25  doc = "X(i,j) *= 2.0"
26}
27#trait_op = {
28  indexing_maps = [
29    affine_map<(i,j) -> (i,j)>,  // A (in)
30    affine_map<(i,j) -> (i,j)>,  // B (in)
31    affine_map<(i,j) -> (i,j)>   // X (out)
32  ],
33  iterator_types = ["parallel", "parallel"],
34  doc = "X(i,j) = A(i,j) OP B(i,j)"
35}
36
37module {
38  // Scales a sparse matrix into a new sparse matrix.
39  func @matrix_scale(%arga: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
40    %s = arith.constant 2.0 : f64
41    %c0 = arith.constant 0 : index
42    %c1 = arith.constant 1 : index
43    %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR>
44    %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR>
45    %xm = sparse_tensor.init [%d0, %d1] : tensor<?x?xf64, #DCSR>
46    %0 = linalg.generic #trait_scale
47       ins(%arga: tensor<?x?xf64, #DCSR>)
48        outs(%xm: tensor<?x?xf64, #DCSR>) {
49        ^bb(%a: f64, %x: f64):
50          %1 = arith.mulf %a, %s : f64
51          linalg.yield %1 : f64
52    } -> tensor<?x?xf64, #DCSR>
53    return %0 : tensor<?x?xf64, #DCSR>
54  }
55
56  // Scales a sparse matrix in place.
57  func @matrix_scale_inplace(%argx: tensor<?x?xf64, #DCSR>
58                             {linalg.inplaceable = true}) -> tensor<?x?xf64, #DCSR> {
59    %s = arith.constant 2.0 : f64
60    %0 = linalg.generic #trait_scale_inpl
61      outs(%argx: tensor<?x?xf64, #DCSR>) {
62        ^bb(%x: f64):
63          %1 = arith.mulf %x, %s : f64
64          linalg.yield %1 : f64
65    } -> tensor<?x?xf64, #DCSR>
66    return %0 : tensor<?x?xf64, #DCSR>
67  }
68
69  // Adds two sparse matrices element-wise into a new sparse matrix.
70  func @matrix_add(%arga: tensor<?x?xf64, #DCSR>,
71                   %argb: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
72    %c0 = arith.constant 0 : index
73    %c1 = arith.constant 1 : index
74    %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR>
75    %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR>
76    %xv = sparse_tensor.init [%d0, %d1] : tensor<?x?xf64, #DCSR>
77    %0 = linalg.generic #trait_op
78       ins(%arga, %argb: tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>)
79        outs(%xv: tensor<?x?xf64, #DCSR>) {
80        ^bb(%a: f64, %b: f64, %x: f64):
81          %1 = arith.addf %a, %b : f64
82          linalg.yield %1 : f64
83    } -> tensor<?x?xf64, #DCSR>
84    return %0 : tensor<?x?xf64, #DCSR>
85  }
86
87  // Multiplies two sparse matrices element-wise into a new sparse matrix.
88  func @matrix_mul(%arga: tensor<?x?xf64, #DCSR>,
89                   %argb: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
90    %c0 = arith.constant 0 : index
91    %c1 = arith.constant 1 : index
92    %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR>
93    %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR>
94    %xv = sparse_tensor.init [%d0, %d1] : tensor<?x?xf64, #DCSR>
95    %0 = linalg.generic #trait_op
96       ins(%arga, %argb: tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>)
97        outs(%xv: tensor<?x?xf64, #DCSR>) {
98        ^bb(%a: f64, %b: f64, %x: f64):
99          %1 = arith.mulf %a, %b : f64
100          linalg.yield %1 : f64
101    } -> tensor<?x?xf64, #DCSR>
102    return %0 : tensor<?x?xf64, #DCSR>
103  }
104
105  // Dump a sparse matrix.
106  func @dump(%arg0: tensor<?x?xf64, #DCSR>) {
107    %d0 = arith.constant 0.0 : f64
108    %c0 = arith.constant 0 : index
109    %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #DCSR> to tensor<?x?xf64>
110    %0 = bufferization.to_memref %dm : memref<?x?xf64>
111    %1 = vector.transfer_read %0[%c0, %c0], %d0: memref<?x?xf64>, vector<4x8xf64>
112    vector.print %1 : vector<4x8xf64>
113    memref.dealloc %0 : memref<?x?xf64>
114    return
115  }
116
117  // Driver method to call and verify matrix kernels.
118  func @entry() {
119    %c0 = arith.constant 0 : index
120    %d1 = arith.constant 1.1 : f64
121
122    // Setup sparse matrices.
123    %m1 = arith.constant sparse<
124       [ [0,0], [0,1], [1,7], [2,2], [2,4], [2,7], [3,0], [3,2], [3,3] ],
125         [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ]
126    > : tensor<4x8xf64>
127    %m2 = arith.constant sparse<
128       [ [0,0], [0,7], [1,0], [1,6], [2,1], [2,7] ],
129         [6.0, 5.0, 4.0, 3.0, 2.0, 1.0 ]
130    > : tensor<4x8xf64>
131    %sm1 = sparse_tensor.convert %m1 : tensor<4x8xf64> to tensor<?x?xf64, #DCSR>
132    %sm2 = sparse_tensor.convert %m2 : tensor<4x8xf64> to tensor<?x?xf64, #DCSR>
133
134    // Call sparse vector kernels.
135    %0 = call @matrix_scale(%sm1)
136      : (tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR>
137    %1 = call @matrix_scale_inplace(%sm1)
138      : (tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR>
139    %2 = call @matrix_add(%sm1, %sm2)
140      : (tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR>
141    %3 = call @matrix_mul(%sm1, %sm2)
142      : (tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR>
143
144    //
145    // Verify the results.
146    //
147    // CHECK:      ( ( 2, 4, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 6 ), ( 0, 0, 8, 0, 10, 0, 0, 12 ), ( 14, 0, 16, 18, 0, 0, 0, 0 ) )
148    // CHECK-NEXT: ( ( 6, 0, 0, 0, 0, 0, 0, 5 ), ( 4, 0, 0, 0, 0, 0, 3, 0 ), ( 0, 2, 0, 0, 0, 0, 0, 1 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ) )
149    // CHECK-NEXT: ( ( 2, 4, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 6 ), ( 0, 0, 8, 0, 10, 0, 0, 12 ), ( 14, 0, 16, 18, 0, 0, 0, 0 ) )
150    // CHECK-NEXT: ( ( 2, 4, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 6 ), ( 0, 0, 8, 0, 10, 0, 0, 12 ), ( 14, 0, 16, 18, 0, 0, 0, 0 ) )
151    // CHECK-NEXT: ( ( 8, 4, 0, 0, 0, 0, 0, 5 ), ( 4, 0, 0, 0, 0, 0, 3, 6 ), ( 0, 2, 8, 0, 10, 0, 0, 13 ), ( 14, 0, 16, 18, 0, 0, 0, 0 ) )
152    // CHECK-NEXT: ( ( 12, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 12 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ) )
153    //
154    call @dump(%sm1) : (tensor<?x?xf64, #DCSR>) -> ()
155    call @dump(%sm2) : (tensor<?x?xf64, #DCSR>) -> ()
156    call @dump(%0) : (tensor<?x?xf64, #DCSR>) -> ()
157    call @dump(%1) : (tensor<?x?xf64, #DCSR>) -> ()
158    call @dump(%2) : (tensor<?x?xf64, #DCSR>) -> ()
159    call @dump(%3) : (tensor<?x?xf64, #DCSR>) -> ()
160
161    // Release the resources.
162    sparse_tensor.release %sm1 : tensor<?x?xf64, #DCSR>
163    sparse_tensor.release %sm2 : tensor<?x?xf64, #DCSR>
164    sparse_tensor.release %0 : tensor<?x?xf64, #DCSR>
165    sparse_tensor.release %2 : tensor<?x?xf64, #DCSR>
166    sparse_tensor.release %3 : tensor<?x?xf64, #DCSR>
167    return
168  }
169}
170