1// RUN: mlir-opt %s --sparse-tensor-conversion | 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_dim( 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_dim(%arg0: tensor<?xf64, #SparseVector>) -> index { 38 %c = constant 0 : index 39 %0 = memref.dim %arg0, %c : tensor<?xf64, #SparseVector> 40 return %0 : index 41} 42 43// CHECK-LABEL: func @sparse_new1d( 44// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 45// CHECK: %[[D:.*]] = constant dense<1> : tensor<1xi8> 46// CHECK: %[[C:.*]] = tensor.cast %[[D]] : tensor<1xi8> to tensor<?xi8> 47// CHECK: %[[P:.*]] = constant dense<0> : tensor<1xi64> 48// CHECK: %[[Q:.*]] = tensor.cast %[[P]] : tensor<1xi64> to tensor<?xi64> 49// CHECK: %[[T:.*]] = call @newSparseTensor(%[[A]], %[[C]], %[[Q]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<i8>, tensor<?xi8>, tensor<?xi64>, i64, i64, i64) -> !llvm.ptr<i8> 50// CHECK: return %[[T]] : !llvm.ptr<i8> 51func @sparse_new1d(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> { 52 %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector> 53 return %0 : tensor<128xf64, #SparseVector> 54} 55 56// CHECK-LABEL: func @sparse_new2d( 57// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 58// CHECK: %[[D:.*]] = constant dense<[0, 1]> : tensor<2xi8> 59// CHECK: %[[C:.*]] = tensor.cast %[[D]] : tensor<2xi8> to tensor<?xi8> 60// CHECK: %[[P:.*]] = constant dense<[0, 1]> : tensor<2xi64> 61// CHECK: %[[Q:.*]] = tensor.cast %[[P]] : tensor<2xi64> to tensor<?xi64> 62// CHECK: %[[T:.*]] = call @newSparseTensor(%[[A]], %[[C]], %[[Q]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<i8>, tensor<?xi8>, tensor<?xi64>, i64, i64, i64) -> !llvm.ptr<i8> 63// CHECK: return %[[T]] : !llvm.ptr<i8> 64func @sparse_new2d(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #SparseMatrix> { 65 %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #SparseMatrix> 66 return %0 : tensor<?x?xf32, #SparseMatrix> 67} 68 69// CHECK-LABEL: func @sparse_new3d( 70// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 71// CHECK: %[[D:.*]] = constant dense<[0, 1, 1]> : tensor<3xi8> 72// CHECK: %[[C:.*]] = tensor.cast %[[D]] : tensor<3xi8> to tensor<?xi8> 73// CHECK: %[[P:.*]] = constant dense<[1, 2, 0]> : tensor<3xi64> 74// CHECK: %[[Q:.*]] = tensor.cast %[[P]] : tensor<3xi64> to tensor<?xi64> 75// CHECK: %[[T:.*]] = call @newSparseTensor(%[[A]], %[[C]], %[[Q]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<i8>, tensor<?xi8>, tensor<?xi64>, i64, i64, i64) -> !llvm.ptr<i8> 76// CHECK: return %[[T]] : !llvm.ptr<i8> 77func @sparse_new3d(%arg0: !llvm.ptr<i8>) -> tensor<?x?x?xf32, #SparseTensor> { 78 %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?x?xf32, #SparseTensor> 79 return %0 : tensor<?x?x?xf32, #SparseTensor> 80} 81 82// CHECK-LABEL: func @sparse_pointers( 83// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 84// CHECK: %[[C:.*]] = constant 0 : index 85// CHECK: %[[T:.*]] = call @sparsePointers(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> 86// CHECK: return %[[T]] : memref<?xindex> 87func @sparse_pointers(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> { 88 %c = constant 0 : index 89 %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex> 90 return %0 : memref<?xindex> 91} 92 93// CHECK-LABEL: func @sparse_pointers64( 94// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 95// CHECK: %[[C:.*]] = constant 0 : index 96// CHECK: %[[T:.*]] = call @sparsePointers64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi64> 97// CHECK: return %[[T]] : memref<?xi64> 98func @sparse_pointers64(%arg0: tensor<128xf64, #SparseVector64>) -> memref<?xi64> { 99 %c = constant 0 : index 100 %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector64> to memref<?xi64> 101 return %0 : memref<?xi64> 102} 103 104// CHECK-LABEL: func @sparse_pointers32( 105// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 106// CHECK: %[[C:.*]] = constant 0 : index 107// CHECK: %[[T:.*]] = call @sparsePointers32(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi32> 108// CHECK: return %[[T]] : memref<?xi32> 109func @sparse_pointers32(%arg0: tensor<128xf64, #SparseVector32>) -> memref<?xi32> { 110 %c = constant 0 : index 111 %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector32> to memref<?xi32> 112 return %0 : memref<?xi32> 113} 114 115// CHECK-LABEL: func @sparse_indices( 116// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 117// CHECK: %[[C:.*]] = constant 0 : index 118// CHECK: %[[T:.*]] = call @sparseIndices(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> 119// CHECK: return %[[T]] : memref<?xindex> 120func @sparse_indices(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> { 121 %c = constant 0 : index 122 %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex> 123 return %0 : memref<?xindex> 124} 125 126// CHECK-LABEL: func @sparse_indices64( 127// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 128// CHECK: %[[C:.*]] = constant 0 : index 129// CHECK: %[[T:.*]] = call @sparseIndices64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi64> 130// CHECK: return %[[T]] : memref<?xi64> 131func @sparse_indices64(%arg0: tensor<128xf64, #SparseVector64>) -> memref<?xi64> { 132 %c = constant 0 : index 133 %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector64> to memref<?xi64> 134 return %0 : memref<?xi64> 135} 136 137// CHECK-LABEL: func @sparse_indices32( 138// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 139// CHECK: %[[C:.*]] = constant 0 : index 140// CHECK: %[[T:.*]] = call @sparseIndices32(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi32> 141// CHECK: return %[[T]] : memref<?xi32> 142func @sparse_indices32(%arg0: tensor<128xf64, #SparseVector32>) -> memref<?xi32> { 143 %c = constant 0 : index 144 %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector32> to memref<?xi32> 145 return %0 : memref<?xi32> 146} 147 148// CHECK-LABEL: func @sparse_valuesf64( 149// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 150// CHECK: %[[T:.*]] = call @sparseValuesF64(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf64> 151// CHECK: return %[[T]] : memref<?xf64> 152func @sparse_valuesf64(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xf64> { 153 %0 = sparse_tensor.values %arg0 : tensor<128xf64, #SparseVector> to memref<?xf64> 154 return %0 : memref<?xf64> 155} 156 157// CHECK-LABEL: func @sparse_valuesf32( 158// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 159// CHECK: %[[T:.*]] = call @sparseValuesF32(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf32> 160// CHECK: return %[[T]] : memref<?xf32> 161func @sparse_valuesf32(%arg0: tensor<128xf32, #SparseVector>) -> memref<?xf32> { 162 %0 = sparse_tensor.values %arg0: tensor<128xf32, #SparseVector> to memref<?xf32> 163 return %0 : memref<?xf32> 164} 165 166// CHECK-LABEL: func @sparse_valuesi32( 167// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 168// CHECK: %[[T:.*]] = call @sparseValuesI32(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi32> 169// CHECK: return %[[T]] : memref<?xi32> 170func @sparse_valuesi32(%arg0: tensor<128xi32, #SparseVector>) -> memref<?xi32> { 171 %0 = sparse_tensor.values %arg0: tensor<128xi32, #SparseVector> to memref<?xi32> 172 return %0 : memref<?xi32> 173} 174 175// CHECK-LABEL: func @sparse_valuesi16( 176// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 177// CHECK: %[[T:.*]] = call @sparseValuesI16(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi16> 178// CHECK: return %[[T]] : memref<?xi16> 179func @sparse_valuesi16(%arg0: tensor<128xi16, #SparseVector>) -> memref<?xi16> { 180 %0 = sparse_tensor.values %arg0: tensor<128xi16, #SparseVector> to memref<?xi16> 181 return %0 : memref<?xi16> 182} 183 184// CHECK-LABEL: func @sparse_valuesi8( 185// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 186// CHECK: %[[T:.*]] = call @sparseValuesI8(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi8> 187// CHECK: return %[[T]] : memref<?xi8> 188func @sparse_valuesi8(%arg0: tensor<128xi8, #SparseVector>) -> memref<?xi8> { 189 %0 = sparse_tensor.values %arg0: tensor<128xi8, #SparseVector> to memref<?xi8> 190 return %0 : memref<?xi8> 191} 192 193// CHECK-LABEL: func @sparse_reconstruct( 194// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8> 195// CHECK: return %[[A]] : !llvm.ptr<i8> 196func @sparse_reconstruct(%arg0: tensor<128xf32, #DenseVector> {linalg.inplaceable = true}) -> tensor<128xf32, #DenseVector> { 197 %0 = sparse_tensor.values %arg0 : tensor<128xf32, #DenseVector> to memref<?xf32> 198 %1 = sparse_tensor.tensor %0 : memref<?xf32> to tensor<128xf32, #DenseVector> 199 return %1 : tensor<128xf32, #DenseVector> 200} 201