1// RUN: mlir-opt %s \
2// RUN:   --test-sparsification="lower ptr-type=2 ind-type=2 fast-output" \
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_sampled_dense_dense = {
20  indexing_maps = [
21    affine_map<(i,j,k) -> (i,j)>,  // S
22    affine_map<(i,j,k) -> (i,k)>,  // A
23    affine_map<(i,j,k) -> (k,j)>,  // B
24    affine_map<(i,j,k) -> (i,j)>   // X (out)
25  ],
26  sparse = [
27    [ "S", "S" ],  // S
28    [ "D", "D" ],  // A
29    [ "D", "D" ],  // B
30    [ "D", "D" ]   // X
31  ],
32  iterator_types = ["parallel", "parallel", "reduction"],
33  doc = "X(i,j) += S(i,j) SUM_k A(i,k) B(k,j)"
34}
35
36//
37// Integration test that lowers a kernel annotated as sparse to
38// actual sparse code, initializes a matching sparse storage scheme
39// from file, and runs the resulting code with the JIT compiler.
40//
41module {
42  //
43  // The kernel expressed as an annotated Linalg op. The kernel
44  // computes a sampled matrix matrix multiplication.
45  //
46  func @sampled_dense_dense(%argS: !SparseTensor,
47                            %arga: tensor<?x?xf32>,
48                            %argb: tensor<?x?xf32>,
49                            %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
50    %args = sparse_tensor.fromPtr %argS : !SparseTensor to tensor<?x?xf32>
51    %0 = linalg.generic #trait_sampled_dense_dense
52      ins(%args, %arga, %argb: tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>)
53      outs(%argx: tensor<?x?xf32>) {
54        ^bb(%s: f32, %a: f32, %b: f32, %x: f32):
55          %0 = mulf %a, %b : f32
56          %1 = mulf %s, %0 : f32
57          %2 = addf %x, %1 : f32
58          linalg.yield %2 : f32
59    } -> tensor<?x?xf32>
60    return %0 : tensor<?x?xf32>
61  }
62
63  //
64  // Runtime support library that is called directly from here.
65  //
66  func private @getTensorFilename(index) -> (!Filename)
67  func private @newSparseTensor(!Filename, memref<?xi1>, index, index, index) -> (!SparseTensor)
68  func private @delSparseTensor(!SparseTensor) -> ()
69
70  //
71  // Main driver that reads matrix from file and calls the sparse kernel.
72  //
73  func @entry() {
74    %d0 = constant 0.0 : f32
75    %c0 = constant 0 : index
76    %c1 = constant 1 : index
77    %c2 = constant 2 : index
78    %c5 = constant 5 : index
79    %c10 = constant 10 : index
80
81    // Mark both dimensions of the matrix as sparse and encode the
82    // storage scheme types (this must match the metadata in the
83    // trait and compiler switches).
84    %annotations = memref.alloc(%c2) : memref<?xi1>
85    %sparse = constant true
86    memref.store %sparse, %annotations[%c0] : memref<?xi1>
87    memref.store %sparse, %annotations[%c1] : memref<?xi1>
88    %i32 = constant 2 : index
89    %f32 = constant 2 : index
90
91    // Setup memory for the dense matrices and initialize.
92    %adata = memref.alloc(%c5, %c10) : memref<?x?xf32>
93    %bdata = memref.alloc(%c10, %c5) : memref<?x?xf32>
94    %xdata = memref.alloc(%c5,  %c5) : memref<?x?xf32>
95    scf.for %i = %c0 to %c5 step %c1 {
96      scf.for %j = %c0 to %c5 step %c1 {
97        memref.store %d0, %xdata[%i, %j] : memref<?x?xf32>
98      }
99      %p = addi %i, %c1 : index
100      %q = index_cast %p : index to i32
101      %d = sitofp %q : i32 to f32
102      scf.for %j = %c0 to %c10 step %c1 {
103        memref.store %d, %adata[%i, %j] : memref<?x?xf32>
104        memref.store %d, %bdata[%j, %i] : memref<?x?xf32>
105      }
106    }
107    %a = memref.tensor_load %adata : memref<?x?xf32>
108    %b = memref.tensor_load %bdata : memref<?x?xf32>
109    %x = memref.tensor_load %xdata : memref<?x?xf32>
110
111    // Read the sparse matrix from file, construct sparse storage
112    // according to <sparse,sparse> in memory, and call the kernel.
113    %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
114    %s = call @newSparseTensor(%fileName, %annotations, %i32, %i32, %f32)
115      : (!Filename, memref<?xi1>, index, index, index) -> (!SparseTensor)
116    %0 = call @sampled_dense_dense(%s, %a, %b, %x)
117       : (!SparseTensor, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
118
119    // Print the result for verification.
120    //
121    // CHECK: ( 10, 0, 0, 56, 0 )
122    // CHECK: ( 0, 80, 0, 0, 250 )
123    // CHECK: ( 0, 0, 270, 0, 0 )
124    // CHECK: ( 164, 0, 0, 640, 0 )
125    // CHECK: ( 0, 520, 0, 0, 1250 )
126    //
127    %r = memref.buffer_cast %0 : memref<?x?xf32>
128    scf.for %i = %c0 to %c5 step %c1 {
129      %v = vector.transfer_read %r[%i, %c0], %d0: memref<?x?xf32>, vector<5xf32>
130      vector.print %v : vector<5xf32>
131    }
132
133    // Release the resources.
134    call @delSparseTensor(%s) : (!SparseTensor) -> ()
135    memref.dealloc %adata : memref<?x?xf32>
136    memref.dealloc %bdata : memref<?x?xf32>
137    memref.dealloc %xdata : memref<?x?xf32>
138
139    return
140  }
141}
142