1// RUN: mlir-opt %s \
2// RUN:   --sparsification --sparse-tensor-conversion \
3// RUN:   --convert-linalg-to-loops --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-std-to-llvm | \
7// RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
8// RUN: TENSOR1="%mlir_integration_test_dir/data/zero.mtx" \
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!Filename = type !llvm.ptr<i8>
15
16#DenseMatrix = #sparse_tensor.encoding<{
17  dimLevelType = [ "dense", "dense" ],
18  dimOrdering = affine_map<(i,j) -> (i,j)>
19}>
20
21#SparseMatrix = #sparse_tensor.encoding<{
22  dimLevelType = [ "dense", "compressed" ],
23  dimOrdering = affine_map<(i,j) -> (i,j)>
24}>
25
26#trait_assign = {
27  indexing_maps = [
28    affine_map<(i,j) -> (i,j)>, // A
29    affine_map<(i,j) -> (i,j)>  // X (out)
30  ],
31  iterator_types = ["parallel", "parallel"],
32  doc = "X(i,j) = A(i,j)"
33}
34
35//
36// Integration test that demonstrates assigning a sparse tensor
37// to an all-dense annotated "sparse" tensor, which effectively
38// result in inserting the nonzero elements into a linearized array.
39//
40// Note that there is a subtle difference between a non-annotated
41// tensor and an all-dense annotated tensor. Both tensors are assumed
42// dense, but the former remains an n-dimensional memref whereas the
43// latter is linearized into a one-dimensional memref that is further
44// lowered into a storage scheme that is backed by the runtime support
45// library.
46module {
47  //
48  // A kernel that assigns elements from A to an initially zero X.
49  //
50  func @dense_output(%arga: tensor<?x?xf64, #SparseMatrix>,
51                     %argx: tensor<?x?xf64, #DenseMatrix>
52		     {linalg.inplaceable = true})
53       -> tensor<?x?xf64, #DenseMatrix> {
54    %0 = linalg.generic #trait_assign
55       ins(%arga: tensor<?x?xf64, #SparseMatrix>)
56      outs(%argx: tensor<?x?xf64, #DenseMatrix>) {
57      ^bb(%a: f64, %x: f64):
58        linalg.yield %a : f64
59    } -> tensor<?x?xf64, #DenseMatrix>
60    return %0 : tensor<?x?xf64, #DenseMatrix>
61  }
62
63  func private @getTensorFilename(index) -> (!Filename)
64
65  //
66  // Main driver that reads matrix from file and calls the kernel.
67  //
68  func @entry() {
69    %d0 = constant 0.0 : f64
70    %c0 = constant 0 : index
71    %c1 = constant 1 : index
72
73    // Read the sparse matrix from file, construct sparse storage.
74    %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
75    %a = sparse_tensor.new %fileName
76      : !llvm.ptr<i8> to tensor<?x?xf64, #SparseMatrix>
77
78    // Initialize all-dense annotated "sparse" matrix to all zeros.
79    %fileZero = call @getTensorFilename(%c1) : (index) -> (!Filename)
80    %x = sparse_tensor.new %fileZero
81      : !llvm.ptr<i8> to tensor<?x?xf64, #DenseMatrix>
82
83    // Call the kernel.
84    %0 = call @dense_output(%a, %x)
85      : (tensor<?x?xf64, #SparseMatrix>,
86         tensor<?x?xf64, #DenseMatrix>) -> tensor<?x?xf64, #DenseMatrix>
87
88    //
89    // Print the linearized 5x5 result for verification.
90    //
91    // CHECK: ( 1, 0, 0, 1.4, 0, 0, 2, 0, 0, 2.5, 0, 0, 3, 0, 0, 4.1, 0, 0, 4, 0, 0, 5.2, 0, 0, 5 )
92    //
93    %m = sparse_tensor.values %0
94      : tensor<?x?xf64, #DenseMatrix> to memref<?xf64>
95    %v = vector.load %m[%c0] : memref<?xf64>, vector<25xf64>
96    vector.print %v : vector<25xf64>
97
98    return
99  }
100}
101