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