xref: /llvm-project/mlir/test/Dialect/SparseTensor/convert_sparse2dense.mlir (revision f82bee1367a1d612d688790b81c1c146ce99f2ea) 
1// RUN: mlir-opt %s --stage-sparse-ops --lower-sparse-ops-to-foreach --canonicalize --cse | FileCheck %s
2
3#SparseVector = #sparse_tensor.encoding<{
4  map = (d0) -> (d0 : compressed)
5}>
6
7#SparseMatrix = #sparse_tensor.encoding<{
8  map = (d0, d1) -> (d0 : dense, d1 : compressed)
9}>
10
11#SparseTensor = #sparse_tensor.encoding<{
12  map = (d0, d1, d2) -> (d2 : dense, d0 : compressed, d1 : compressed)
13}>
14
15// CHECK-LABEL:  func.func @sparse_convert_1d
16// CHECK-NOT:      sparse_tensor.reorder_coo
17// CHECK:          bufferization.alloc_tensor
18// CHECK:          linalg.fill
19// CHECK:          sparse_tensor.foreach
20// CHECK:            tensor.insert
21func.func @sparse_convert_1d(%arg0: tensor<13xi32, #SparseVector>) -> tensor<13xi32> {
22  %0 = sparse_tensor.convert %arg0 : tensor<13xi32, #SparseVector> to tensor<13xi32>
23  return %0 : tensor<13xi32>
24}
25
26// CHECK-LABEL:  func.func @sparse_convert_1d_dyn
27// CHECK-NOT:      sparse_tensor.reorder_coo
28// CHECK:          bufferization.alloc_tensor
29// CHECK:          linalg.fill
30// CHECK:          sparse_tensor.foreach
31// CHECK:            tensor.insert
32func.func @sparse_convert_1d_dyn(%arg0: tensor<?xi32, #SparseVector>) -> tensor<?xi32> {
33  %0 = sparse_tensor.convert %arg0 : tensor<?xi32, #SparseVector> to tensor<?xi32>
34  return %0 : tensor<?xi32>
35}
36
37// CHECK-LABEL:  func.func @sparse_convert_2d
38// CHECK-NOT:      sparse_tensor.reorder_coo
39// CHECK:          bufferization.alloc_tensor
40// CHECK:          linalg.fill
41// CHECK:          sparse_tensor.foreach
42// CHECK:            tensor.insert
43func.func @sparse_convert_2d(%arg0: tensor<2x4xf64, #SparseMatrix>) -> tensor<2x4xf64> {
44  %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64, #SparseMatrix> to tensor<2x4xf64>
45  return %0 : tensor<2x4xf64>
46}
47
48// CHECK-LABEL:  func.func @sparse_convert_2d_dyn
49// CHECK-NOT:      sparse_tensor.reorder_coo
50// CHECK:          bufferization.alloc_tensor
51// CHECK:          linalg.fill
52// CHECK:          sparse_tensor.foreach
53// CHECK:            tensor.insert
54func.func @sparse_convert_2d_dyn0(%arg0: tensor<?x4xf64, #SparseMatrix>) -> tensor<?x4xf64> {
55  %0 = sparse_tensor.convert %arg0 : tensor<?x4xf64, #SparseMatrix> to tensor<?x4xf64>
56  return %0 : tensor<?x4xf64>
57}
58
59// CHECK-LABEL:  func.func @sparse_convert_2d_dyn1
60// CHECK-NOT:      sparse_tensor.reorder_coo
61// CHECK:          bufferization.alloc_tensor
62// CHECK:          linalg.fill
63// CHECK:          sparse_tensor.foreach
64// CHECK:            tensor.insert
65func.func @sparse_convert_2d_dyn1(%arg0: tensor<2x?xf64, #SparseMatrix>) -> tensor<2x?xf64> {
66  %0 = sparse_tensor.convert %arg0 : tensor<2x?xf64, #SparseMatrix> to tensor<2x?xf64>
67  return %0 : tensor<2x?xf64>
68}
69
70// CHECK-LABEL:  func.func @sparse_convert_2d_dyn2
71// CHECK-NOT:      sparse_tensor.reorder_coo
72// CHECK:          bufferization.alloc_tensor
73// CHECK:          linalg.fill
74// CHECK:          sparse_tensor.foreach
75// CHECK:            tensor.insert
76func.func @sparse_convert_2d_dyn2(%arg0: tensor<?x?xf64, #SparseMatrix>) -> tensor<?x?xf64> {
77  %0 = sparse_tensor.convert %arg0 : tensor<?x?xf64, #SparseMatrix> to tensor<?x?xf64>
78  return %0 : tensor<?x?xf64>
79}
80
81// CHECK-LABEL:  func.func @sparse_convert_3d
82// CHECK-NOT:      sparse_tensor.reorder_coo
83// CHECK:          bufferization.alloc_tensor
84// CHECK:          linalg.fill
85// CHECK:          sparse_tensor.foreach
86// CHECK:            tensor.insert
87func.func @sparse_convert_3d(%arg0: tensor<2x3x4xf64, #SparseTensor>) -> tensor<2x3x4xf64> {
88  %0 = sparse_tensor.convert %arg0 : tensor<2x3x4xf64, #SparseTensor> to tensor<2x3x4xf64>
89  return %0 : tensor<2x3x4xf64>
90}
91