1// RUN: mlir-opt %s -sparsification                           | \
2// RUN:   FileCheck %s --check-prefix=CHECK-SPARSE
3// RUN: mlir-opt %s -sparsification -sparse-tensor-conversion | \
4// RUN:   FileCheck %s --check-prefix=CHECK-CONVERT
5
6#DCSC = #sparse_tensor.encoding<{
7  dimLevelType = [  "compressed", "compressed" ],
8  dimOrdering = affine_map<(i,j) -> (j,i)>
9}>
10
11#SV = #sparse_tensor.encoding<{
12  dimLevelType = [  "compressed" ]
13}>
14
15#rowsum = {
16  indexing_maps = [
17    affine_map<(i,j) -> (i,j)>, // A
18    affine_map<(i,j) -> (i)>    // x (out)
19  ],
20  iterator_types = ["parallel", "reduction"],
21  doc = "X(i) = SUM A(i,j)"
22}
23
24//
25// CHECK-SPARSE-LABEL: func @kernel(
26// CHECK-SPARSE: %[[A:.*]], %[[B:.*]], %[[C:.*]], %{{.*}} = sparse_tensor.expand
27// CHECK-SPARSE: scf.for
28// CHECK-SPARSE:   scf.for
29// CHECK-SPARSE: sparse_tensor.compress %{{.*}}, %{{.*}}, %[[A]], %[[B]], %[[C]]
30// CHECK-SPARSE: %[[RET:.*]] = sparse_tensor.load %{{.*}} hasInserts
31// CHECK-SPARSE: return %[[RET]]
32//
33// CHECK-CONVERT-LABEL: func @kernel(
34// CHECK-CONVERT: %{{.*}} = call @sparseDimSize
35// CHECK-CONVERT: %[[S:.*]] = call @sparseDimSize
36// CHECK-CONVERT: %[[A:.*]] = memref.alloc(%[[S]]) : memref<?xf64>
37// CHECK-CONVERT: %[[B:.*]] = memref.alloc(%[[S]]) : memref<?xi1>
38// CHECK-CONVERT: %[[C:.*]] = memref.alloc(%[[S]]) : memref<?xindex>
39// CHECK-CONVERT: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<?xf64>)
40// CHECK-CONVERT: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>)
41// CHECK-CONVERT: scf.for
42// CHECK-CONVERT:   scf.for
43// CHECK-CONVERT: call @expInsertF64
44// CHECK-CONVERT: memref.dealloc %[[A]] : memref<?xf64>
45// CHECK-CONVERT: memref.dealloc %[[B]] : memref<?xi1>
46// CHECK-CONVERT: memref.dealloc %[[C]] : memref<?xindex>
47// CHECK-CONVERT: call @endInsert
48//
49func.func @kernel(%arga: tensor<?x?xf64, #DCSC>) -> tensor<?xf64, #SV> {
50  %c0 = arith.constant 0 : index
51  %n = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSC>
52  %v = bufferization.alloc_tensor(%n) : tensor<?xf64, #SV>
53  %0 = linalg.generic #rowsum
54    ins(%arga: tensor<?x?xf64, #DCSC>)
55    outs(%v: tensor<?xf64, #SV>) {
56    ^bb(%a: f64, %x: f64):
57      %1 = arith.addf %x, %a : f64
58      linalg.yield %1 : f64
59  } -> tensor<?xf64, #SV>
60  return %0 : tensor<?xf64, #SV>
61}
62