1// RUN: mlir-opt %s --sparse-tensor-conversion | FileCheck %s
2
3#SparseVector = #sparse_tensor.encoding<{
4  dimLevelType = ["compressed"]
5}>
6
7#SparseVector64 = #sparse_tensor.encoding<{
8  dimLevelType = ["compressed"],
9  pointerBitWidth = 64,
10  indexBitWidth = 64
11}>
12
13#SparseVector32 = #sparse_tensor.encoding<{
14  dimLevelType = ["compressed"],
15  pointerBitWidth = 32,
16  indexBitWidth = 32
17}>
18
19#SparseMatrix = #sparse_tensor.encoding<{
20  dimLevelType = ["dense", "compressed"]
21}>
22
23// CHECK-LABEL: func @sparse_dim(
24//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
25//       CHECK: %[[C:.*]] = constant 0 : index
26//       CHECK: %[[D:.*]] = call @sparseDimSize(%[[A]], %[[C]])
27//       CHECK: return %[[D]] : index
28func @sparse_dim(%arg0: tensor<?xf64, #SparseVector>) -> index {
29  %c = constant 0 : index
30  %0 = memref.dim %arg0, %c : tensor<?xf64, #SparseVector>
31  return %0 : index
32}
33
34// CHECK-LABEL: func @sparse_new1d(
35//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
36//       CHECK: %[[D:.*]] = constant dense<true> : tensor<1xi1>
37//       CHECK: %[[C:.*]] = tensor.cast %[[D]] : tensor<1xi1> to tensor<?xi1>
38//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[A]], %[[C]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<i8>, tensor<?xi1>, i64, i64, i64) -> !llvm.ptr<i8>
39//       CHECK: return %[[T]] : !llvm.ptr<i8>
40func @sparse_new1d(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> {
41  %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector>
42  return %0 : tensor<128xf64, #SparseVector>
43}
44
45// CHECK-LABEL: func @sparse_new2d(
46//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
47//       CHECK: %[[D:.*]] = constant dense<[false, true]> : tensor<2xi1>
48//       CHECK: %[[C:.*]] = tensor.cast %[[D]] : tensor<2xi1> to tensor<?xi1>
49//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[A]], %[[C]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<i8>, tensor<?xi1>, i64, i64, i64) -> !llvm.ptr<i8>
50//       CHECK: return %[[T]] : !llvm.ptr<i8>
51func @sparse_new2d(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #SparseMatrix> {
52  %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #SparseMatrix>
53  return %0 : tensor<?x?xf32, #SparseMatrix>
54}
55
56// CHECK-LABEL: func @sparse_pointers(
57//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
58//       CHECK: %[[C:.*]] = constant 0 : index
59//       CHECK: %[[T:.*]] = call @sparsePointers(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
60//       CHECK: return %[[T]] : memref<?xindex>
61func @sparse_pointers(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> {
62  %c = constant 0 : index
63  %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex>
64  return %0 : memref<?xindex>
65}
66
67// CHECK-LABEL: func @sparse_pointers64(
68//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
69//       CHECK: %[[C:.*]] = constant 0 : index
70//       CHECK: %[[T:.*]] = call @sparsePointers64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi64>
71//       CHECK: return %[[T]] : memref<?xi64>
72func @sparse_pointers64(%arg0: tensor<128xf64, #SparseVector64>) -> memref<?xi64> {
73  %c = constant 0 : index
74  %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector64> to memref<?xi64>
75  return %0 : memref<?xi64>
76}
77
78// CHECK-LABEL: func @sparse_pointers32(
79//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
80//       CHECK: %[[C:.*]] = constant 0 : index
81//       CHECK: %[[T:.*]] = call @sparsePointers32(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi32>
82//       CHECK: return %[[T]] : memref<?xi32>
83func @sparse_pointers32(%arg0: tensor<128xf64, #SparseVector32>) -> memref<?xi32> {
84  %c = constant 0 : index
85  %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector32> to memref<?xi32>
86  return %0 : memref<?xi32>
87}
88
89// CHECK-LABEL: func @sparse_indices(
90//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
91//       CHECK: %[[C:.*]] = constant 0 : index
92//       CHECK: %[[T:.*]] = call @sparseIndices(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
93//       CHECK: return %[[T]] : memref<?xindex>
94func @sparse_indices(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> {
95  %c = constant 0 : index
96  %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex>
97  return %0 : memref<?xindex>
98}
99
100// CHECK-LABEL: func @sparse_indices64(
101//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
102//       CHECK: %[[C:.*]] = constant 0 : index
103//       CHECK: %[[T:.*]] = call @sparseIndices64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi64>
104//       CHECK: return %[[T]] : memref<?xi64>
105func @sparse_indices64(%arg0: tensor<128xf64, #SparseVector64>) -> memref<?xi64> {
106  %c = constant 0 : index
107  %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector64> to memref<?xi64>
108  return %0 : memref<?xi64>
109}
110
111// CHECK-LABEL: func @sparse_indices32(
112//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
113//       CHECK: %[[C:.*]] = constant 0 : index
114//       CHECK: %[[T:.*]] = call @sparseIndices32(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi32>
115//       CHECK: return %[[T]] : memref<?xi32>
116func @sparse_indices32(%arg0: tensor<128xf64, #SparseVector32>) -> memref<?xi32> {
117  %c = constant 0 : index
118  %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector32> to memref<?xi32>
119  return %0 : memref<?xi32>
120}
121
122// CHECK-LABEL: func @sparse_valuesf64(
123//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
124//       CHECK: %[[T:.*]] = call @sparseValuesF64(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf64>
125//       CHECK: return %[[T]] : memref<?xf64>
126func @sparse_valuesf64(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xf64> {
127  %0 = sparse_tensor.values %arg0 : tensor<128xf64, #SparseVector> to memref<?xf64>
128  return %0 : memref<?xf64>
129}
130
131// CHECK-LABEL: func @sparse_valuesf32(
132//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
133//       CHECK: %[[T:.*]] = call @sparseValuesF32(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf32>
134//       CHECK: return %[[T]] : memref<?xf32>
135func @sparse_valuesf32(%arg0: tensor<128xf32, #SparseVector>) -> memref<?xf32> {
136  %0 = sparse_tensor.values %arg0: tensor<128xf32, #SparseVector> to memref<?xf32>
137  return %0 : memref<?xf32>
138}
139
140// CHECK-LABEL: func @sparse_valuesi32(
141//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
142//       CHECK: %[[T:.*]] = call @sparseValuesI32(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi32>
143//       CHECK: return %[[T]] : memref<?xi32>
144func @sparse_valuesi32(%arg0: tensor<128xi32, #SparseVector>) -> memref<?xi32> {
145  %0 = sparse_tensor.values %arg0: tensor<128xi32, #SparseVector> to memref<?xi32>
146  return %0 : memref<?xi32>
147}
148
149// CHECK-LABEL: func @sparse_valuesi16(
150//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
151//       CHECK: %[[T:.*]] = call @sparseValuesI16(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi16>
152//       CHECK: return %[[T]] : memref<?xi16>
153func @sparse_valuesi16(%arg0: tensor<128xi16, #SparseVector>) -> memref<?xi16> {
154  %0 = sparse_tensor.values %arg0: tensor<128xi16, #SparseVector> to memref<?xi16>
155  return %0 : memref<?xi16>
156}
157
158// CHECK-LABEL: func @sparse_valuesi8(
159//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
160//       CHECK: %[[T:.*]] = call @sparseValuesI8(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi8>
161//       CHECK: return %[[T]] : memref<?xi8>
162func @sparse_valuesi8(%arg0: tensor<128xi8, #SparseVector>) -> memref<?xi8> {
163  %0 = sparse_tensor.values %arg0: tensor<128xi8, #SparseVector> to memref<?xi8>
164  return %0 : memref<?xi8>
165}
166