1// RUN: mlir-opt %s --sparse-compiler | \
2// RUN: mlir-cpu-runner \
3// RUN:  -e entry -entry-point-result=void  \
4// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
5// RUN: FileCheck %s
6
7#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
8
9#trait_op = {
10  indexing_maps = [
11    affine_map<(i) -> (i)>,  // a (in)
12    affine_map<(i) -> (i)>   // x (out)
13  ],
14  iterator_types = ["parallel"],
15  doc = "x(i) = OP a(i)"
16}
17
18module {
19  func.func @cre(%arga: tensor<?xcomplex<f32>, #SparseVector>)
20                -> tensor<?xf32, #SparseVector> {
21    %c = arith.constant 0 : index
22    %d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector>
23    %xv = bufferization.alloc_tensor(%d) : tensor<?xf32, #SparseVector>
24    %0 = linalg.generic #trait_op
25       ins(%arga: tensor<?xcomplex<f32>, #SparseVector>)
26        outs(%xv: tensor<?xf32, #SparseVector>) {
27        ^bb(%a: complex<f32>, %x: f32):
28          %1 = complex.re %a : complex<f32>
29          linalg.yield %1 : f32
30    } -> tensor<?xf32, #SparseVector>
31    return %0 : tensor<?xf32, #SparseVector>
32  }
33
34  func.func @cim(%arga: tensor<?xcomplex<f32>, #SparseVector>)
35                -> tensor<?xf32, #SparseVector> {
36    %c = arith.constant 0 : index
37    %d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector>
38    %xv = bufferization.alloc_tensor(%d) : tensor<?xf32, #SparseVector>
39    %0 = linalg.generic #trait_op
40       ins(%arga: tensor<?xcomplex<f32>, #SparseVector>)
41        outs(%xv: tensor<?xf32, #SparseVector>) {
42        ^bb(%a: complex<f32>, %x: f32):
43          %1 = complex.im %a : complex<f32>
44          linalg.yield %1 : f32
45    } -> tensor<?xf32, #SparseVector>
46    return %0 : tensor<?xf32, #SparseVector>
47  }
48
49  func.func @dump(%arg0: tensor<?xf32, #SparseVector>) {
50    %c0 = arith.constant 0 : index
51    %d0 = arith.constant -1.0 : f32
52    %values = sparse_tensor.values %arg0 : tensor<?xf32, #SparseVector> to memref<?xf32>
53    %0 = vector.transfer_read %values[%c0], %d0: memref<?xf32>, vector<4xf32>
54    vector.print %0 : vector<4xf32>
55    %indices = sparse_tensor.indices %arg0, %c0 : tensor<?xf32, #SparseVector> to memref<?xindex>
56    %1 = vector.transfer_read %indices[%c0], %c0: memref<?xindex>, vector<4xindex>
57    vector.print %1 : vector<4xindex>
58    return
59  }
60
61  // Driver method to call and verify functions cim and cre.
62  func.func @entry() {
63    // Setup sparse vectors.
64    %v1 = arith.constant sparse<
65       [ [0], [20], [31] ],
66         [ (5.13, 2.0), (3.0, 4.0), (5.0, 6.0) ] > : tensor<32xcomplex<f32>>
67    %sv1 = sparse_tensor.convert %v1 : tensor<32xcomplex<f32>> to tensor<?xcomplex<f32>, #SparseVector>
68
69    // Call sparse vector kernels.
70    %0 = call @cre(%sv1)
71       : (tensor<?xcomplex<f32>, #SparseVector>) -> tensor<?xf32, #SparseVector>
72
73    %1 = call @cim(%sv1)
74       : (tensor<?xcomplex<f32>, #SparseVector>) -> tensor<?xf32, #SparseVector>
75
76    //
77    // Verify the results.
78    //
79    // CHECK: ( 5.13, 3, 5, -1 )
80    // CHECK-NEXT: ( 0, 20, 31, 0 )
81    // CHECK-NEXT: ( 2, 4, 6, -1 )
82    // CHECK-NEXT: ( 0, 20, 31, 0 )
83    //
84    call @dump(%0) : (tensor<?xf32, #SparseVector>) -> ()
85    call @dump(%1) : (tensor<?xf32, #SparseVector>) -> ()
86
87    // Release the resources.
88    bufferization.dealloc_tensor %sv1 : tensor<?xcomplex<f32>, #SparseVector>
89    bufferization.dealloc_tensor %0 : tensor<?xf32, #SparseVector>
90    bufferization.dealloc_tensor %1 : tensor<?xf32, #SparseVector>
91    return
92  }
93}
94