1// RUN: mlir-opt %s --sparse-tensor-conversion --canonicalize --cse | FileCheck %s 2 3#DenseVector = #sparse_tensor.encoding<{ 4 dimLevelType = ["dense"] 5}> 6 7#SparseVector = #sparse_tensor.encoding<{ 8 dimLevelType = ["compressed"] 9}> 10 11#SparseVector64 = #sparse_tensor.encoding<{ 12 dimLevelType = ["compressed"], 13 pointerBitWidth = 64, 14 indexBitWidth = 64 15}> 16 17#SparseVector32 = #sparse_tensor.encoding<{ 18 dimLevelType = ["compressed"], 19 pointerBitWidth = 32, 20 indexBitWidth = 32 21}> 22 23#SparseMatrix = #sparse_tensor.encoding<{ 24 dimLevelType = ["dense", "compressed"] 25}> 26 27#SparseTensor = #sparse_tensor.encoding<{ 28 dimLevelType = ["dense", "compressed", "compressed"], 29 dimOrdering = affine_map<(i,j,k) -> (k,i,j)> 30}> 31 32// CHECK-LABEL: func @sparse_dim1d( 33// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 34// CHECK: %[[C:.*]] = constant 0 : index 35// CHECK: %[[D:.*]] = call @sparseDimSize(%[[A]], %[[C]]) 36// CHECK: return %[[D]] : index 37func @sparse_dim1d(%arg0: tensor<?xf64, #SparseVector>) -> index { 38 %c = constant 0 : index 39 %0 = tensor.dim %arg0, %c : tensor<?xf64, #SparseVector> 40 return %0 : index 41} 42 43// CHECK-LABEL: func @sparse_dim3d( 44// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 45// CHECK: %[[C:.*]] = constant 2 : index 46// CHECK: %[[D:.*]] = call @sparseDimSize(%[[A]], %[[C]]) 47// CHECK: return %[[D]] : index 48func @sparse_dim3d(%arg0: tensor<?x?x?xf64, #SparseTensor>) -> index { 49 // Querying for dimension 1 in the tensor type needs to be 50 // permuted into querying for dimension 2 in the stored sparse 51 // tensor scheme, since the latter honors the dimOrdering. 52 %c = constant 1 : index 53 %0 = tensor.dim %arg0, %c : tensor<?x?x?xf64, #SparseTensor> 54 return %0 : index 55} 56 57// CHECK-LABEL: func @sparse_dim3d_const( 58// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 59// CHECK: %[[C:.*]] = constant 20 : index 60// CHECK: return %[[C]] : index 61func @sparse_dim3d_const(%arg0: tensor<10x20x30xf64, #SparseTensor>) -> index { 62 // Querying for dimension 1 in the tensor type can be directly 63 // folded into the right value (even though it corresponds 64 // to dimension 2 in the stored sparse tensor scheme). 65 %c = constant 1 : index 66 %0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #SparseTensor> 67 return %0 : index 68} 69 70// CHECK-LABEL: func @sparse_new1d( 71// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 72// CHECK-DAG: %[[U:.*]] = constant dense<1> : tensor<1xi8> 73// CHECK-DAG: %[[V:.*]] = constant dense<128> : tensor<1xi64> 74// CHECK-DAG: %[[W:.*]] = constant dense<0> : tensor<1xi64> 75// CHECK-DAG: %[[X:.*]] = tensor.cast %[[U]] : tensor<1xi8> to tensor<?xi8> 76// CHECK-DAG: %[[Y:.*]] = tensor.cast %[[V]] : tensor<1xi64> to tensor<?xi64> 77// CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<1xi64> to tensor<?xi64> 78// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]]) 79// CHECK: return %[[T]] : !llvm.ptr<i8> 80func @sparse_new1d(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> { 81 %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector> 82 return %0 : tensor<128xf64, #SparseVector> 83} 84 85// CHECK-LABEL: func @sparse_new2d( 86// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 87// CHECK-DAG: %[[U:.*]] = constant dense<[0, 1]> : tensor<2xi8> 88// CHECK-DAG: %[[V:.*]] = constant dense<0> : tensor<2xi64> 89// CHECK-DAG: %[[W:.*]] = constant dense<[0, 1]> : tensor<2xi64> 90// CHECK-DAG: %[[X:.*]] = tensor.cast %[[U]] : tensor<2xi8> to tensor<?xi8> 91// CHECK-DAG: %[[Y:.*]] = tensor.cast %[[V]] : tensor<2xi64> to tensor<?xi64> 92// CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<2xi64> to tensor<?xi64> 93// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]]) 94// CHECK: return %[[T]] : !llvm.ptr<i8> 95func @sparse_new2d(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #SparseMatrix> { 96 %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #SparseMatrix> 97 return %0 : tensor<?x?xf32, #SparseMatrix> 98} 99 100// CHECK-LABEL: func @sparse_new3d( 101// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 102// CHECK-DAG: %[[U:.*]] = constant dense<[0, 1, 1]> : tensor<3xi8> 103// CHECK-DAG: %[[V:.*]] = constant dense<0> : tensor<3xi64> 104// CHECK-DAG: %[[W:.*]] = constant dense<[1, 2, 0]> : tensor<3xi64> 105// CHECK-DAG: %[[X:.*]] = tensor.cast %[[U]] : tensor<3xi8> to tensor<?xi8> 106// CHECK-DAG: %[[Y:.*]] = tensor.cast %[[V]] : tensor<3xi64> to tensor<?xi64> 107// CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<3xi64> to tensor<?xi64> 108// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]]) 109// CHECK: return %[[T]] : !llvm.ptr<i8> 110func @sparse_new3d(%arg0: !llvm.ptr<i8>) -> tensor<?x?x?xf32, #SparseTensor> { 111 %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?x?xf32, #SparseTensor> 112 return %0 : tensor<?x?x?xf32, #SparseTensor> 113} 114 115// CHECK-LABEL: func @sparse_convert( 116// CHECK-SAME: %[[A:.*]]: tensor<2x4xf64>) -> !llvm.ptr<i8> 117// CHECK-DAG: %[[U:.*]] = constant dense<[0, 1]> : tensor<2xi8> 118// CHECK-DAG: %[[V:.*]] = constant dense<[2, 4]> : tensor<2xi64> 119// CHECK-DAG: %[[W:.*]] = constant dense<[0, 1]> : tensor<2xi64> 120// CHECK: %[[C:.*]] = memref.buffer_cast %arg0 : memref<2x4xf64> 121// CHECK: %[[M:.*]] = memref.cast %[[C]] : memref<2x4xf64> to memref<*xf64> 122// CHECK: %[[C:.*]] = call @getPtrF64(%[[M]]) : (memref<*xf64>) -> !llvm.ptr<i8> 123// CHECK-DAG: %[[X:.*]] = tensor.cast %[[U]] : tensor<2xi8> to tensor<?xi8> 124// CHECK-DAG: %[[Y:.*]] = tensor.cast %[[V]] : tensor<2xi64> to tensor<?xi64> 125// CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<2xi64> to tensor<?xi64> 126// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]]) 127// CHECK: return %[[T]] : !llvm.ptr<i8> 128func @sparse_convert(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #SparseMatrix> { 129 %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64> to tensor<2x4xf64, #SparseMatrix> 130 return %0 : tensor<2x4xf64, #SparseMatrix> 131} 132 133// CHECK-LABEL: func @sparse_pointers( 134// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 135// CHECK: %[[C:.*]] = constant 0 : index 136// CHECK: %[[T:.*]] = call @sparsePointers(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> 137// CHECK: return %[[T]] : memref<?xindex> 138func @sparse_pointers(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> { 139 %c = constant 0 : index 140 %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex> 141 return %0 : memref<?xindex> 142} 143 144// CHECK-LABEL: func @sparse_pointers64( 145// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 146// CHECK: %[[C:.*]] = constant 0 : index 147// CHECK: %[[T:.*]] = call @sparsePointers64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi64> 148// CHECK: return %[[T]] : memref<?xi64> 149func @sparse_pointers64(%arg0: tensor<128xf64, #SparseVector64>) -> memref<?xi64> { 150 %c = constant 0 : index 151 %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector64> to memref<?xi64> 152 return %0 : memref<?xi64> 153} 154 155// CHECK-LABEL: func @sparse_pointers32( 156// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 157// CHECK: %[[C:.*]] = constant 0 : index 158// CHECK: %[[T:.*]] = call @sparsePointers32(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi32> 159// CHECK: return %[[T]] : memref<?xi32> 160func @sparse_pointers32(%arg0: tensor<128xf64, #SparseVector32>) -> memref<?xi32> { 161 %c = constant 0 : index 162 %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector32> to memref<?xi32> 163 return %0 : memref<?xi32> 164} 165 166// CHECK-LABEL: func @sparse_indices( 167// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 168// CHECK: %[[C:.*]] = constant 0 : index 169// CHECK: %[[T:.*]] = call @sparseIndices(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> 170// CHECK: return %[[T]] : memref<?xindex> 171func @sparse_indices(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> { 172 %c = constant 0 : index 173 %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex> 174 return %0 : memref<?xindex> 175} 176 177// CHECK-LABEL: func @sparse_indices64( 178// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 179// CHECK: %[[C:.*]] = constant 0 : index 180// CHECK: %[[T:.*]] = call @sparseIndices64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi64> 181// CHECK: return %[[T]] : memref<?xi64> 182func @sparse_indices64(%arg0: tensor<128xf64, #SparseVector64>) -> memref<?xi64> { 183 %c = constant 0 : index 184 %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector64> to memref<?xi64> 185 return %0 : memref<?xi64> 186} 187 188// CHECK-LABEL: func @sparse_indices32( 189// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 190// CHECK: %[[C:.*]] = constant 0 : index 191// CHECK: %[[T:.*]] = call @sparseIndices32(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi32> 192// CHECK: return %[[T]] : memref<?xi32> 193func @sparse_indices32(%arg0: tensor<128xf64, #SparseVector32>) -> memref<?xi32> { 194 %c = constant 0 : index 195 %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector32> to memref<?xi32> 196 return %0 : memref<?xi32> 197} 198 199// CHECK-LABEL: func @sparse_valuesf64( 200// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 201// CHECK: %[[T:.*]] = call @sparseValuesF64(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf64> 202// CHECK: return %[[T]] : memref<?xf64> 203func @sparse_valuesf64(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xf64> { 204 %0 = sparse_tensor.values %arg0 : tensor<128xf64, #SparseVector> to memref<?xf64> 205 return %0 : memref<?xf64> 206} 207 208// CHECK-LABEL: func @sparse_valuesf32( 209// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 210// CHECK: %[[T:.*]] = call @sparseValuesF32(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf32> 211// CHECK: return %[[T]] : memref<?xf32> 212func @sparse_valuesf32(%arg0: tensor<128xf32, #SparseVector>) -> memref<?xf32> { 213 %0 = sparse_tensor.values %arg0: tensor<128xf32, #SparseVector> to memref<?xf32> 214 return %0 : memref<?xf32> 215} 216 217// CHECK-LABEL: func @sparse_valuesi32( 218// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 219// CHECK: %[[T:.*]] = call @sparseValuesI32(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi32> 220// CHECK: return %[[T]] : memref<?xi32> 221func @sparse_valuesi32(%arg0: tensor<128xi32, #SparseVector>) -> memref<?xi32> { 222 %0 = sparse_tensor.values %arg0: tensor<128xi32, #SparseVector> to memref<?xi32> 223 return %0 : memref<?xi32> 224} 225 226// CHECK-LABEL: func @sparse_valuesi16( 227// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 228// CHECK: %[[T:.*]] = call @sparseValuesI16(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi16> 229// CHECK: return %[[T]] : memref<?xi16> 230func @sparse_valuesi16(%arg0: tensor<128xi16, #SparseVector>) -> memref<?xi16> { 231 %0 = sparse_tensor.values %arg0: tensor<128xi16, #SparseVector> to memref<?xi16> 232 return %0 : memref<?xi16> 233} 234 235// CHECK-LABEL: func @sparse_valuesi8( 236// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 237// CHECK: %[[T:.*]] = call @sparseValuesI8(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi8> 238// CHECK: return %[[T]] : memref<?xi8> 239func @sparse_valuesi8(%arg0: tensor<128xi8, #SparseVector>) -> memref<?xi8> { 240 %0 = sparse_tensor.values %arg0: tensor<128xi8, #SparseVector> to memref<?xi8> 241 return %0 : memref<?xi8> 242} 243 244// CHECK-LABEL: func @sparse_reconstruct_1( 245// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8> 246// CHECK: return %[[A]] : !llvm.ptr<i8> 247func @sparse_reconstruct_1(%arg0: tensor<128xf32, #DenseVector> {linalg.inplaceable = true}) -> tensor<128xf32, #DenseVector> { 248 %0 = sparse_tensor.values %arg0 : tensor<128xf32, #DenseVector> to memref<?xf32> 249 %1 = sparse_tensor.tensor %0 : memref<?xf32> to tensor<128xf32, #DenseVector> 250 return %1 : tensor<128xf32, #DenseVector> 251} 252 253// CHECK-LABEL: func @sparse_reconstruct_n( 254// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8> 255// CHECK: return %[[A]] : !llvm.ptr<i8> 256func @sparse_reconstruct_n(%arg0: tensor<128xf32, #SparseVector> {linalg.inplaceable = true}) -> tensor<128xf32, #SparseVector> { 257 %c = constant 0 : index 258 %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf32, #SparseVector> to memref<?xindex> 259 %1 = sparse_tensor.indices %arg0, %c : tensor<128xf32, #SparseVector> to memref<?xindex> 260 %2 = sparse_tensor.values %arg0 : tensor<128xf32, #SparseVector> to memref<?xf32> 261 %3 = sparse_tensor.tensor %0, %1, %2 : memref<?xindex>, memref<?xindex>, memref<?xf32> to tensor<128xf32, #SparseVector> 262 return %3 : tensor<128xf32, #SparseVector> 263} 264