1// RUN: mlir-opt %s --sparse-compiler | \
2// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
3// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
4// RUN: FileCheck %s
5
6#SparseMatrix = #sparse_tensor.encoding<{
7  dimLevelType = ["compressed", "compressed"]
8}>
9
10#trait = {
11  indexing_maps = [
12    affine_map<(i,j) -> (i,j)>,  // A
13    affine_map<(i,j) -> (i,j)>   // X (out)
14  ],
15  iterator_types = ["parallel", "parallel"],
16  doc = "X(i,j) = A(i,j) * i * j"
17}
18
19module {
20
21  //
22  // Kernel that uses indices in the index notation.
23  //
24  func @sparse_index(%arga: tensor<3x4xi64, #SparseMatrix>)
25                         -> tensor<3x4xi64, #SparseMatrix> {
26    %d0 = arith.constant 3 : index
27    %d1 = arith.constant 4 : index
28    %init = sparse_tensor.init [%d0, %d1] : tensor<3x4xi64, #SparseMatrix>
29    %r = linalg.generic #trait
30        ins(%arga: tensor<3x4xi64, #SparseMatrix>)
31       outs(%init: tensor<3x4xi64, #SparseMatrix>) {
32        ^bb(%a: i64, %x: i64):
33          %i = linalg.index 0 : index
34          %j = linalg.index 1 : index
35          %ii = arith.index_cast %i : index to i64
36          %jj = arith.index_cast %j : index to i64
37          %m1 = arith.muli %ii, %a : i64
38          %m2 = arith.muli %jj, %m1 : i64
39          linalg.yield %m2 : i64
40    } -> tensor<3x4xi64, #SparseMatrix>
41    return %r : tensor<3x4xi64, #SparseMatrix>
42  }
43
44  //
45  // Main driver.
46  //
47  func @entry() {
48    %c0 = arith.constant 0 : index
49    %c1 = arith.constant 1 : index
50    %c4 = arith.constant 4 : index
51    %du = arith.constant -1 : i64
52
53    // Setup input "sparse" matrix.
54    %d = arith.constant dense <[
55       [ 1,  1,  1,  1 ],
56       [ 1,  1,  1,  1 ],
57       [ 1,  1,  1,  1 ]
58    ]> : tensor<3x4xi64>
59    %a = sparse_tensor.convert %d : tensor<3x4xi64> to tensor<3x4xi64, #SparseMatrix>
60
61    // Call the kernel.
62    %0 = call @sparse_index(%a) : (tensor<3x4xi64, #SparseMatrix>) -> tensor<3x4xi64, #SparseMatrix>
63
64    //
65    // Verify result.
66    //
67    // CHECK: ( ( 0, 0, 0, 0 ), ( 0, 1, 2, 3 ), ( 0, 2, 4, 6 ) )
68    //
69    %x = sparse_tensor.convert %0 : tensor<3x4xi64, #SparseMatrix> to tensor<3x4xi64>
70    %m = bufferization.to_memref %x : memref<3x4xi64>
71    %v = vector.transfer_read %m[%c0, %c0], %du: memref<3x4xi64>, vector<3x4xi64>
72    vector.print %v : vector<3x4xi64>
73
74    // Release resources.
75    sparse_tensor.release %a : tensor<3x4xi64, #SparseMatrix>
76    sparse_tensor.release %0 : tensor<3x4xi64, #SparseMatrix>
77    memref.dealloc %m : memref<3x4xi64>
78
79    return
80  }
81}
82