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: mlir-cpu-runner \
9// RUN:  -e entry -entry-point-result=void  \
10// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
11// RUN: FileCheck %s
12
13!Filename = type !llvm.ptr<i8>
14
15#SparseMatrix = #sparse_tensor.encoding<{
16  dimLevelType = [ "compressed", "compressed" ]
17}>
18
19#trait_sum_reduce = {
20  indexing_maps = [
21    affine_map<(i,j) -> (i,j)>, // A
22    affine_map<(i,j) -> ()>     // x (out)
23  ],
24  iterator_types = ["reduction", "reduction"],
25  doc = "x += A(i,j)"
26}
27
28//
29// Integration test that lowers a kernel annotated as sparse to
30// actual sparse code, initializes a matching sparse storage scheme
31// from file, and runs the resulting code with the JIT compiler.
32//
33module {
34  //
35  // A kernel that sum-reduces a matrix to a single scalar.
36  //
37  func @kernel_sum_reduce(%arga: tensor<?x?xf64, #SparseMatrix>,
38                          %argx: tensor<f64>) -> tensor<f64> {
39    %0 = linalg.generic #trait_sum_reduce
40      ins(%arga: tensor<?x?xf64, #SparseMatrix>)
41      outs(%argx: tensor<f64>) {
42      ^bb(%a: f64, %x: f64):
43        %0 = addf %x, %a : f64
44        linalg.yield %0 : f64
45    } -> tensor<f64>
46    return %0 : tensor<f64>
47  }
48
49  func private @getTensorFilename(index) -> (!Filename)
50
51  //
52  // Main driver that reads matrix from file and calls the sparse kernel.
53  //
54  func @entry() {
55    %d0 = constant 0.0 : f64
56    %c0 = constant 0 : index
57
58    // Setup memory for a single reduction scalar,
59    // initialized to zero.
60    %xdata = memref.alloc() : memref<f64>
61    memref.store %d0, %xdata[] : memref<f64>
62    %x = memref.tensor_load %xdata : memref<f64>
63
64    // Read the sparse matrix from file, construct sparse storage.
65    %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
66    %a = sparse_tensor.new %fileName : !llvm.ptr<i8> to tensor<?x?xf64, #SparseMatrix>
67
68    // Call the kernel.
69    %0 = call @kernel_sum_reduce(%a, %x)
70      : (tensor<?x?xf64, #SparseMatrix>, tensor<f64>) -> tensor<f64>
71
72    // Print the result for verification.
73    //
74    // CHECK: 28.2
75    //
76    %m = memref.buffer_cast %0 : memref<f64>
77    %v = memref.load %m[] : memref<f64>
78    vector.print %v : f64
79
80    // Release the resources.
81    memref.dealloc %xdata : memref<f64>
82
83    return
84  }
85}
86