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