1// RUN: mlir-opt %s --sparse-compiler | \
2// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
3// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
4// RUN: FileCheck %s
5
6#SparseVector = #sparse_tensor.encoding<{
7  dimLevelType = ["compressed"]
8}>
9
10#SparseMatrix = #sparse_tensor.encoding<{
11  dimLevelType = ["compressed", "compressed"]
12}>
13
14//
15// Test with various forms of the two most elementary reshape
16// operations: expand/collapse.
17//
18module {
19
20  func.func @expand_dense(%arg0: tensor<12xf64>) -> tensor<3x4xf64> {
21    %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64> into tensor<3x4xf64>
22    return %0 : tensor<3x4xf64>
23  }
24
25  func.func @expand_from_sparse(%arg0: tensor<12xf64, #SparseVector>) -> tensor<3x4xf64> {
26    %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64, #SparseVector> into tensor<3x4xf64>
27    return %0 : tensor<3x4xf64>
28  }
29
30  func.func @expand_to_sparse(%arg0: tensor<12xf64>) -> tensor<3x4xf64, #SparseMatrix> {
31    %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64> into tensor<3x4xf64, #SparseMatrix>
32    return %0 : tensor<3x4xf64, #SparseMatrix>
33  }
34
35// TODO: make this work
36//  func.func @expand_sparse2sparse(%arg0: tensor<12xf64, #SparseVector>) -> tensor<3x4xf64, #SparseMatrix> {
37//    %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64, #SparseVector> into tensor<3x4xf64, #SparseMatrix>
38//    return %0 : tensor<3x4xf64, #SparseMatrix>
39//  }
40
41  func.func @collapse_dense(%arg0: tensor<3x4xf64>) -> tensor<12xf64> {
42    %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64> into tensor<12xf64>
43    return %0 : tensor<12xf64>
44  }
45
46  func.func @collapse_from_sparse(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64> {
47    %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64, #SparseMatrix> into tensor<12xf64>
48    return %0 : tensor<12xf64>
49  }
50
51  func.func @collapse_to_sparse(%arg0: tensor<3x4xf64>) -> tensor<12xf64, #SparseVector> {
52    %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64> into tensor<12xf64, #SparseVector>
53    return %0 : tensor<12xf64, #SparseVector>
54  }
55
56// TODO: make this work
57//  func.func @collapse_sparse2sparse(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64, #SparseVector> {
58//    %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64, #SparseMatrix> into tensor<12xf64, #SparseVector>
59//    return %0 : tensor<12xf64, #SparseVector>
60//  }
61
62
63  //
64  // Main driver.
65  //
66  func.func @entry() {
67    %c0 = arith.constant 0 : index
68    %df = arith.constant -1.0 : f64
69
70    // Setup test vectors and matrices..
71    %v = arith.constant dense <[ 1.0, 2.0, 3.0,  4.0,  5.0,  6.0,
72                                 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]> : tensor<12xf64>
73    %m = arith.constant dense <[ [ 1.1,  1.2,  1.3,  1.4 ],
74                                 [ 2.1,  2.2,  2.3,  2.4 ],
75                                 [ 3.1,  3.2,  3.3,  3.4 ]]> : tensor<3x4xf64>
76    %sv = sparse_tensor.convert %v : tensor<12xf64> to tensor<12xf64, #SparseVector>
77    %sm = sparse_tensor.convert %m : tensor<3x4xf64> to tensor<3x4xf64, #SparseMatrix>
78
79
80    // Call the kernels.
81    %expand0 = call @expand_dense(%v) : (tensor<12xf64>) -> tensor<3x4xf64>
82    %expand1 = call @expand_from_sparse(%sv) : (tensor<12xf64, #SparseVector>) -> tensor<3x4xf64>
83    %expand2 = call @expand_to_sparse(%v) : (tensor<12xf64>) -> tensor<3x4xf64, #SparseMatrix>
84
85    %collapse0 = call @collapse_dense(%m) : (tensor<3x4xf64>) -> tensor<12xf64>
86    %collapse1 = call @collapse_from_sparse(%sm) : (tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64>
87    %collapse2 = call @collapse_to_sparse(%m) : (tensor<3x4xf64>) -> tensor<12xf64, #SparseVector>
88
89    //
90    // Verify result.
91    //
92    // CHECK:      ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) )
93    // CHECK-NEXT: ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) )
94    // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, -1, -1, -1, -1 )
95    // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4 )
96    // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4 )
97    // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 )
98    //
99    %m0 = vector.transfer_read %expand0[%c0, %c0], %df: tensor<3x4xf64>, vector<3x4xf64>
100    vector.print %m0 : vector<3x4xf64>
101    %m1 = vector.transfer_read %expand1[%c0, %c0], %df: tensor<3x4xf64>, vector<3x4xf64>
102    vector.print %m1 : vector<3x4xf64>
103    %a2 = sparse_tensor.values %expand2 : tensor<3x4xf64, #SparseMatrix> to memref<?xf64>
104    %m2 = vector.transfer_read %a2[%c0], %df: memref<?xf64>, vector<16xf64>
105    vector.print %m2 : vector<16xf64>
106
107    %v0 = vector.transfer_read %collapse0[%c0], %df: tensor<12xf64>, vector<12xf64>
108    vector.print %v0 : vector<12xf64>
109    %v1 = vector.transfer_read %collapse1[%c0], %df: tensor<12xf64>, vector<12xf64>
110    vector.print %v1 : vector<12xf64>
111    %b2 = sparse_tensor.values %collapse2 : tensor<12xf64, #SparseVector> to memref<?xf64>
112    %v2 = vector.transfer_read %b2[%c0], %df: memref<?xf64>, vector<16xf64>
113    vector.print %v2 : vector<16xf64>
114
115    // Release sparse resources.
116    sparse_tensor.release %sv : tensor<12xf64, #SparseVector>
117    sparse_tensor.release %sm : tensor<3x4xf64, #SparseMatrix>
118    sparse_tensor.release %expand2 : tensor<3x4xf64, #SparseMatrix>
119    sparse_tensor.release %collapse2 : tensor<12xf64, #SparseVector>
120
121    // Release dense resources.
122    %meme1 = bufferization.to_memref %expand1 : memref<3x4xf64>
123    memref.dealloc %meme1 : memref<3x4xf64>
124    %memc1 = bufferization.to_memref %collapse1 : memref<12xf64>
125    memref.dealloc %memc1 : memref<12xf64>
126
127    return
128  }
129}
130