1// RUN: mlir-opt %s \
2// RUN:   --sparsification="fast-output" --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  pointerBitWidth = 32,
18  indexBitWidth = 32
19}>
20
21#trait_sampled_dense_dense = {
22  indexing_maps = [
23    affine_map<(i,j,k) -> (i,j)>,  // S
24    affine_map<(i,j,k) -> (i,k)>,  // A
25    affine_map<(i,j,k) -> (k,j)>,  // B
26    affine_map<(i,j,k) -> (i,j)>   // X (out)
27  ],
28  iterator_types = ["parallel", "parallel", "reduction"],
29  doc = "X(i,j) += S(i,j) SUM_k A(i,k) B(k,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  // A kernel that computes a sampled matrix matrix multiplication.
40  //
41  func @sampled_dense_dense(%args: tensor<?x?xf32, #SparseMatrix>,
42                            %arga: tensor<?x?xf32>,
43                            %argb: tensor<?x?xf32>,
44                            %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
45    %0 = linalg.generic #trait_sampled_dense_dense
46      ins(%args, %arga, %argb: tensor<?x?xf32, #SparseMatrix>, tensor<?x?xf32>, tensor<?x?xf32>)
47      outs(%argx: tensor<?x?xf32>) {
48        ^bb(%s: f32, %a: f32, %b: f32, %x: f32):
49          %0 = mulf %a, %b : f32
50          %1 = mulf %s, %0 : f32
51          %2 = addf %x, %1 : f32
52          linalg.yield %2 : f32
53    } -> tensor<?x?xf32>
54    return %0 : tensor<?x?xf32>
55  }
56
57  func private @getTensorFilename(index) -> (!Filename)
58
59  //
60  // Main driver that reads matrix from file and calls the sparse kernel.
61  //
62  func @entry() {
63    %d0 = constant 0.0 : f32
64    %c0 = constant 0 : index
65    %c1 = constant 1 : index
66    %c5 = constant 5 : index
67    %c10 = constant 10 : index
68
69    // Setup memory for the dense matrices and initialize.
70    %adata = memref.alloc(%c5, %c10) : memref<?x?xf32>
71    %bdata = memref.alloc(%c10, %c5) : memref<?x?xf32>
72    %xdata = memref.alloc(%c5,  %c5) : memref<?x?xf32>
73    scf.for %i = %c0 to %c5 step %c1 {
74      scf.for %j = %c0 to %c5 step %c1 {
75        memref.store %d0, %xdata[%i, %j] : memref<?x?xf32>
76      }
77      %p = addi %i, %c1 : index
78      %q = index_cast %p : index to i32
79      %d = sitofp %q : i32 to f32
80      scf.for %j = %c0 to %c10 step %c1 {
81        memref.store %d, %adata[%i, %j] : memref<?x?xf32>
82        memref.store %d, %bdata[%j, %i] : memref<?x?xf32>
83      }
84    }
85    %a = memref.tensor_load %adata : memref<?x?xf32>
86    %b = memref.tensor_load %bdata : memref<?x?xf32>
87    %x = memref.tensor_load %xdata : memref<?x?xf32>
88
89    // Read the sparse matrix from file, construct sparse storage.
90    %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
91    %s = sparse_tensor.new %fileName : !llvm.ptr<i8> to tensor<?x?xf32, #SparseMatrix>
92
93    // Call the kernel.
94    %0 = call @sampled_dense_dense(%s, %a, %b, %x)
95       : (tensor<?x?xf32, #SparseMatrix>,
96          tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
97
98    // Print the result for verification.
99    //
100    // CHECK: ( 10, 0, 0, 56, 0 )
101    // CHECK: ( 0, 80, 0, 0, 250 )
102    // CHECK: ( 0, 0, 270, 0, 0 )
103    // CHECK: ( 164, 0, 0, 640, 0 )
104    // CHECK: ( 0, 520, 0, 0, 1250 )
105    //
106    %r = memref.buffer_cast %0 : memref<?x?xf32>
107    scf.for %i = %c0 to %c5 step %c1 {
108      %v = vector.transfer_read %r[%i, %c0], %d0: memref<?x?xf32>, vector<5xf32>
109      vector.print %v : vector<5xf32>
110    }
111
112    // Release the resources.
113    memref.dealloc %adata : memref<?x?xf32>
114    memref.dealloc %bdata : memref<?x?xf32>
115    memref.dealloc %xdata : memref<?x?xf32>
116
117    return
118  }
119}
120