1// RUN: mlir-opt %s \
2// RUN:   --test-sparsification="lower" \
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//
14// Use descriptive names for opaque pointers.
15//
16!Filename     = type !llvm.ptr<i8>
17!SparseTensor = type !llvm.ptr<i8>
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  sparse = [
25    [ "S", "S" ], // A
26    [          ]  // x
27  ],
28  iterator_types = ["reduction", "reduction"],
29  doc = "x += A(i,j)"
30}
31
32//
33// Integration test that lowers a kernel annotated as sparse to
34// actual sparse code, initializes a matching sparse storage scheme
35// from file, and runs the resulting code with the JIT compiler.
36//
37module {
38  //
39  // The kernel expressed as an annotated Linalg op. The kernel
40  // sum reduces a matrix to a single scalar.
41  //
42  func @kernel_sum_reduce(%argA: !SparseTensor,
43                          %argx: tensor<f64>) -> tensor<f64> {
44    %arga = sparse_tensor.fromPtr %argA : !SparseTensor to tensor<?x?xf64>
45    %0 = linalg.generic #trait_sum_reduce
46      ins(%arga: tensor<?x?xf64>)
47      outs(%argx: tensor<f64>) {
48      ^bb(%a: f64, %x: f64):
49        %0 = addf %x, %a : f64
50        linalg.yield %0 : f64
51    } -> tensor<f64>
52    return %0 : tensor<f64>
53  }
54
55  //
56  // Runtime support library that is called directly from here.
57  //
58  func private @getTensorFilename(index) -> (!Filename)
59  func private @newSparseTensor(!Filename, memref<?xi1>, index, index, index) -> (!SparseTensor)
60  func private @delSparseTensor(!SparseTensor) -> ()
61
62  //
63  // Main driver that reads matrix from file and calls the sparse kernel.
64  //
65  func @entry() {
66    %d0 = constant 0.0 : f64
67    %c0 = constant 0 : index
68    %c1 = constant 1 : index
69    %c2 = constant 2 : index
70
71    // Mark both dimensions of the matrix as sparse and encode the
72    // storage scheme types (this must match the metadata in the
73    // trait and compiler switches).
74    %annotations = memref.alloc(%c2) : memref<?xi1>
75    %sparse = constant true
76    memref.store %sparse, %annotations[%c0] : memref<?xi1>
77    memref.store %sparse, %annotations[%c1] : memref<?xi1>
78    %i64 = constant 1 : index
79    %f64 = constant 1 : index
80
81    // Setup memory for a single reduction scalar,
82    // initialized to zero.
83    %xdata = memref.alloc() : memref<f64>
84    memref.store %d0, %xdata[] : memref<f64>
85    %x = memref.tensor_load %xdata : memref<f64>
86
87    // Read the sparse matrix from file, construct sparse storage
88    // according to <sparse,sparse> in memory, and call the kernel.
89    %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
90    %a = call @newSparseTensor(%fileName, %annotations, %i64, %i64, %f64)
91      : (!Filename, memref<?xi1>, index, index, index) -> (!SparseTensor)
92    %0 = call @kernel_sum_reduce(%a, %x)
93      : (!SparseTensor, tensor<f64>) -> tensor<f64>
94
95    // Print the result for verification.
96    //
97    // CHECK: 28.2
98    //
99    %m = memref.buffer_cast %0 : memref<f64>
100    %v = memref.load %m[] : memref<f64>
101    vector.print %v : f64
102
103    // Release the resources.
104    call @delSparseTensor(%a) : (!SparseTensor) -> ()
105    memref.dealloc %xdata : memref<f64>
106
107    return
108  }
109}
110