1// RUN: mlir-opt %s --sparse-tensor-conversion --canonicalize --cse | 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#SparseTensor = #sparse_tensor.encoding<{ 24 dimLevelType = ["dense", "compressed", "compressed"], 25 dimOrdering = affine_map<(i,j,k) -> (k,i,j)> 26}> 27 28// CHECK-LABEL: func @sparse_dim1d( 29// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 30// CHECK: %[[C:.*]] = arith.constant 0 : index 31// CHECK: %[[D:.*]] = call @sparseDimSize(%[[A]], %[[C]]) 32// CHECK: return %[[D]] : index 33func.func @sparse_dim1d(%arg0: tensor<?xf64, #SparseVector>) -> index { 34 %c = arith.constant 0 : index 35 %0 = tensor.dim %arg0, %c : tensor<?xf64, #SparseVector> 36 return %0 : index 37} 38 39// CHECK-LABEL: func @sparse_dim3d( 40// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 41// CHECK: %[[C:.*]] = arith.constant 2 : index 42// CHECK: %[[D:.*]] = call @sparseDimSize(%[[A]], %[[C]]) 43// CHECK: return %[[D]] : index 44func.func @sparse_dim3d(%arg0: tensor<?x?x?xf64, #SparseTensor>) -> index { 45 // Querying for dimension 1 in the tensor type needs to be 46 // permuted into querying for dimension 2 in the stored sparse 47 // tensor scheme, since the latter honors the dimOrdering. 48 %c = arith.constant 1 : index 49 %0 = tensor.dim %arg0, %c : tensor<?x?x?xf64, #SparseTensor> 50 return %0 : index 51} 52 53// CHECK-LABEL: func @sparse_dim3d_const( 54// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 55// CHECK: %[[C:.*]] = arith.constant 20 : index 56// CHECK: return %[[C]] : index 57func.func @sparse_dim3d_const(%arg0: tensor<10x20x30xf64, #SparseTensor>) -> index { 58 // Querying for dimension 1 in the tensor type can be directly 59 // folded into the right value (even though it corresponds 60 // to dimension 2 in the stored sparse tensor scheme). 61 %c = arith.constant 1 : index 62 %0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #SparseTensor> 63 return %0 : index 64} 65 66// CHECK-LABEL: func @sparse_new1d( 67// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 68// CHECK-DAG: %[[FromFile:.*]] = arith.constant 1 : i32 69// CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<1xi8> 70// CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<1xindex> 71// CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<1xindex> 72// CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<1xi8> to memref<?xi8> 73// CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xindex> to memref<?xindex> 74// CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xindex> to memref<?xindex> 75// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromFile]], %[[A]]) 76// CHECK: return %[[T]] : !llvm.ptr<i8> 77func.func @sparse_new1d(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> { 78 %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector> 79 return %0 : tensor<128xf64, #SparseVector> 80} 81 82// CHECK-LABEL: func @sparse_new2d( 83// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 84// CHECK-DAG: %[[FromFile:.*]] = arith.constant 1 : i32 85// CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<2xi8> 86// CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<2xindex> 87// CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<2xindex> 88// CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<2xi8> to memref<?xi8> 89// CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xindex> to memref<?xindex> 90// CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xindex> to memref<?xindex> 91// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromFile]], %[[A]]) 92// CHECK: return %[[T]] : !llvm.ptr<i8> 93func.func @sparse_new2d(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #SparseMatrix> { 94 %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #SparseMatrix> 95 return %0 : tensor<?x?xf32, #SparseMatrix> 96} 97 98// CHECK-LABEL: func @sparse_new3d( 99// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 100// CHECK-DAG: %[[FromFile:.*]] = arith.constant 1 : i32 101// CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<3xi8> 102// CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<3xindex> 103// CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<3xindex> 104// CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<3xi8> to memref<?xi8> 105// CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<3xindex> to memref<?xindex> 106// CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<3xindex> to memref<?xindex> 107// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromFile]], %[[A]]) 108// CHECK: return %[[T]] : !llvm.ptr<i8> 109func.func @sparse_new3d(%arg0: !llvm.ptr<i8>) -> tensor<?x?x?xf32, #SparseTensor> { 110 %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?x?xf32, #SparseTensor> 111 return %0 : tensor<?x?x?xf32, #SparseTensor> 112} 113 114// CHECK-LABEL: func @sparse_init( 115// CHECK-SAME: %[[I:.*]]: index, 116// CHECK-SAME: %[[J:.*]]: index) -> !llvm.ptr<i8> 117// CHECK-DAG: %[[Empty:.*]] = arith.constant 0 : i32 118// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 119// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 120// CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<2xi8> 121// CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<2xindex> 122// CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<2xindex> 123// CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<2xi8> to memref<?xi8> 124// CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xindex> to memref<?xindex> 125// CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xindex> to memref<?xindex> 126// CHECK-DAG: memref.store %[[I]], %[[Q]][%[[C0]]] : memref<2xindex> 127// CHECK-DAG: memref.store %[[J]], %[[Q]][%[[C1]]] : memref<2xindex> 128// CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr<i8> 129// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[Empty]], %[[NP]]) 130// CHECK: return %[[T]] : !llvm.ptr<i8> 131func.func @sparse_init(%arg0: index, %arg1: index) -> tensor<?x?xf64, #SparseMatrix> { 132 %0 = sparse_tensor.init [%arg0, %arg1] : tensor<?x?xf64, #SparseMatrix> 133 return %0 : tensor<?x?xf64, #SparseMatrix> 134} 135 136// CHECK-LABEL: func @sparse_release( 137// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 138// CHECK: call @delSparseTensor(%[[A]]) : (!llvm.ptr<i8>) -> () 139// CHECK: return 140func.func @sparse_release(%arg0: tensor<128xf64, #SparseVector>) { 141 sparse_tensor.release %arg0 : tensor<128xf64, #SparseVector> 142 return 143} 144 145// CHECK-LABEL: func @sparse_nop_convert( 146// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 147// CHECK: return %[[A]] : !llvm.ptr<i8> 148func.func @sparse_nop_convert(%arg0: tensor<64xf32, #SparseVector>) -> tensor<64xf32, #SparseVector> { 149 %0 = sparse_tensor.convert %arg0 : tensor<64xf32, #SparseVector> to tensor<64xf32, #SparseVector> 150 return %0 : tensor<64xf32, #SparseVector> 151} 152 153// CHECK-LABEL: func @sparse_hidden_nop_cast( 154// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 155// CHECK: return %[[A]] : !llvm.ptr<i8> 156func.func @sparse_hidden_nop_cast(%arg0: tensor<32xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> { 157 %0 = sparse_tensor.convert %arg0 : tensor<32xf32, #SparseVector> to tensor<?xf32, #SparseVector> 158 return %0 : tensor<?xf32, #SparseVector> 159} 160 161// CHECK-LABEL: func @sparse_nop_cast( 162// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> 163// CHECK: return %[[A]] : !llvm.ptr<i8> 164func.func @sparse_nop_cast(%arg0: tensor<64xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> { 165 %0 = tensor.cast %arg0 : tensor<64xf32, #SparseVector> to tensor<?xf32, #SparseVector> 166 return %0 : tensor<?xf32, #SparseVector> 167} 168 169// CHECK-LABEL: func @sparse_convert_1d( 170// CHECK-SAME: %[[A:.*]]: tensor<?xi32>) -> !llvm.ptr<i8> 171// CHECK-DAG: %[[EmptyCOO:.*]] = arith.constant 4 : i32 172// CHECK-DAG: %[[FromCOO:.*]] = arith.constant 2 : i32 173// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 174// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 175// CHECK-DAG: %[[U:.*]] = tensor.dim %[[A]], %[[C0]] : tensor<?xi32> 176// CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<1xi8> 177// CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<1xindex> 178// CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<1xindex> 179// CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<1xi8> to memref<?xi8> 180// CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xindex> to memref<?xindex> 181// CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xindex> to memref<?xindex> 182// CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr<i8> 183// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[EmptyCOO]], %[[NP]]) 184// CHECK: %[[M:.*]] = memref.alloca() : memref<1xindex> 185// CHECK: %[[T:.*]] = memref.cast %[[M]] : memref<1xindex> to memref<?xindex> 186// CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[U]] step %[[C1]] { 187// CHECK: %[[E:.*]] = tensor.extract %[[A]][%[[I]]] : tensor<?xi32> 188// CHECK: memref.store %[[I]], %[[M]][%[[C0]]] : memref<1xindex> 189// CHECK: call @addEltI32(%[[C]], %[[E]], %[[T]], %[[Z]]) 190// CHECK: } 191// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromCOO]], %[[C]]) 192// CHECK: call @delSparseTensorCOOI32(%[[C]]) 193// CHECK: return %[[T]] : !llvm.ptr<i8> 194func.func @sparse_convert_1d(%arg0: tensor<?xi32>) -> tensor<?xi32, #SparseVector> { 195 %0 = sparse_tensor.convert %arg0 : tensor<?xi32> to tensor<?xi32, #SparseVector> 196 return %0 : tensor<?xi32, #SparseVector> 197} 198 199// CHECK-LABEL: func @sparse_convert_1d_ss( 200// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 201// CHECK-DAG: %[[ToCOO:.*]] = arith.constant 5 : i32 202// CHECK-DAG: %[[FromCOO:.*]] = arith.constant 2 : i32 203// CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<1xi8> 204// CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<1xindex> 205// CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<1xindex> 206// CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<1xi8> to memref<?xi8> 207// CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xindex> to memref<?xindex> 208// CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xindex> to memref<?xindex> 209// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ToCOO]], %[[A]]) 210// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromCOO]], %[[C]]) 211// CHECK: call @delSparseTensorCOOF32(%[[C]]) 212// CHECK: return %[[T]] : !llvm.ptr<i8> 213func.func @sparse_convert_1d_ss(%arg0: tensor<?xf32, #SparseVector64>) -> tensor<?xf32, #SparseVector32> { 214 %0 = sparse_tensor.convert %arg0 : tensor<?xf32, #SparseVector64> to tensor<?xf32, #SparseVector32> 215 return %0 : tensor<?xf32, #SparseVector32> 216} 217 218// CHECK-LABEL: func @sparse_convert_2d( 219// CHECK-SAME: %[[A:.*]]: tensor<2x4xf64>) -> !llvm.ptr<i8> 220// CHECK-DAG: %[[EmptyCOO:.*]] = arith.constant 4 : i32 221// CHECK-DAG: %[[FromCOO:.*]] = arith.constant 2 : i32 222// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 223// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 224// CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<2xi8> 225// CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<2xindex> 226// CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<2xindex> 227// CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<2xi8> to memref<?xi8> 228// CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xindex> to memref<?xindex> 229// CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xindex> to memref<?xindex> 230// CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr<i8> 231// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[EmptyCOO]], %[[NP]]) 232// CHECK: %[[M:.*]] = memref.alloca() : memref<2xindex> 233// CHECK: %[[T:.*]] = memref.cast %[[M]] : memref<2xindex> to memref<?xindex> 234// CHECK: scf.for %[[I:.*]] = %[[C0]] to %{{.*}} step %[[C1]] { 235// CHECK: scf.for %[[J:.*]] = %[[C0]] to %{{.*}} step %[[C1]] { 236// CHECK: %[[E:.*]] = tensor.extract %[[A]][%[[I]], %[[J]]] : tensor<2x4xf64> 237// CHECK: memref.store %[[I]], %[[M]][%[[C0]]] : memref<2xindex> 238// CHECK: memref.store %[[J]], %[[M]][%[[C1]]] : memref<2xindex> 239// CHECK: call @addEltF64(%[[C]], %[[E]], %[[T]], %[[Z]]) 240// CHECK: } 241// CHECK: } 242// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromCOO]], %[[C]]) 243// CHECK: call @delSparseTensorCOOF64(%[[C]]) 244// CHECK: return %[[T]] : !llvm.ptr<i8> 245func.func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #SparseMatrix> { 246 %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64> to tensor<2x4xf64, #SparseMatrix> 247 return %0 : tensor<2x4xf64, #SparseMatrix> 248} 249 250// CHECK-LABEL: func @sparse_constant() -> !llvm.ptr<i8> { 251// CHECK-DAG: %[[EmptyCOO:.*]] = arith.constant 4 : i32 252// CHECK-DAG: %[[FromCOO:.*]] = arith.constant 2 : i32 253// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 254// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 255// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index 256// CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<2xi8> 257// CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<2xindex> 258// CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<2xindex> 259// CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<2xi8> to memref<?xi8> 260// CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xindex> to memref<?xindex> 261// CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xindex> to memref<?xindex> 262// CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr<i8> 263// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[EmptyCOO]], %[[NP]]) 264// CHECK: %[[M:.*]] = memref.alloca() : memref<2xindex> 265// CHECK: %[[N:.*]] = memref.cast %[[M]] : memref<2xindex> to memref<?xindex> 266// CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[C2]] step %[[C1]] { 267// CHECK: memref.store %{{.*}}, %[[M]][%[[C0]]] : memref<2xindex> 268// CHECK: memref.store %{{.*}}, %[[M]][%[[C1]]] : memref<2xindex> 269// CHECK: %[[V:.*]] = tensor.extract %{{.*}}[%[[I]]] : tensor<2xf32> 270// CHECK: call @addEltF32(%{{.*}}, %[[V]], %[[N]], %{{.*}}) 271// CHECK: } 272// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromCOO]], %[[C]]) 273// CHECK: call @delSparseTensorCOOF32(%[[C]]) 274// CHECK: return %[[T]] : !llvm.ptr<i8> 275func.func @sparse_constant() -> tensor<8x7xf32, #SparseMatrix>{ 276 // Initialize a tensor. 277 %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32> 278 // Convert the tensor to a sparse tensor. 279 %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #SparseMatrix> 280 return %1 : tensor<8x7xf32, #SparseMatrix> 281} 282 283// CHECK-LABEL: func @sparse_convert_3d( 284// CHECK-SAME: %[[A:.*]]: tensor<?x?x?xf64>) -> !llvm.ptr<i8> 285// CHECK-DAG: %[[EmptyCOO:.*]] = arith.constant 4 : i32 286// CHECK-DAG: %[[FromCOO:.*]] = arith.constant 2 : i32 287// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 288// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 289// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index 290// CHECK-DAG: %[[U1:.*]] = tensor.dim %[[A]], %[[C0]] : tensor<?x?x?xf64> 291// CHECK-DAG: %[[U2:.*]] = tensor.dim %[[A]], %[[C1]] : tensor<?x?x?xf64> 292// CHECK-DAG: %[[U3:.*]] = tensor.dim %[[A]], %[[C2]] : tensor<?x?x?xf64> 293// CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<3xi8> 294// CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<3xindex> 295// CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<3xindex> 296// CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<3xi8> to memref<?xi8> 297// CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<3xindex> to memref<?xindex> 298// CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<3xindex> to memref<?xindex> 299// CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr<i8> 300// CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[EmptyCOO]], %[[NP]]) 301// CHECK: %[[M:.*]] = memref.alloca() : memref<3xindex> 302// CHECK: %[[N:.*]] = memref.cast %[[M]] : memref<3xindex> to memref<?xindex> 303// CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[U1]] step %[[C1]] { 304// CHECK: scf.for %[[J:.*]] = %[[C0]] to %[[U2]] step %[[C1]] { 305// CHECK: scf.for %[[K:.*]] = %[[C0]] to %[[U3]] step %[[C1]] { 306// CHECK: %[[E:.*]] = tensor.extract %[[A]][%[[I]], %[[J]], %[[K]]] : tensor<?x?x?xf64> 307// CHECK: memref.store %[[I]], %[[M]][%[[C0]]] : memref<3xindex> 308// CHECK: memref.store %[[J]], %[[M]][%[[C1]]] : memref<3xindex> 309// CHECK: memref.store %[[K]], %[[M]][%[[C2]]] : memref<3xindex> 310// CHECK: call @addEltF64(%[[C]], %[[E]], %[[N]], %[[Z]]) 311// CHECK: } 312// CHECK: } 313// CHECK: } 314// CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromCOO]], %[[C]]) 315// CHECK: call @delSparseTensorCOOF64(%[[C]]) 316// CHECK: return %[[T]] : !llvm.ptr<i8> 317func.func @sparse_convert_3d(%arg0: tensor<?x?x?xf64>) -> tensor<?x?x?xf64, #SparseTensor> { 318 %0 = sparse_tensor.convert %arg0 : tensor<?x?x?xf64> to tensor<?x?x?xf64, #SparseTensor> 319 return %0 : tensor<?x?x?xf64, #SparseTensor> 320} 321 322// CHECK-LABEL: func @sparse_pointers( 323// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 324// CHECK: %[[C:.*]] = arith.constant 0 : index 325// CHECK: %[[T:.*]] = call @sparsePointers(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> 326// CHECK: return %[[T]] : memref<?xindex> 327func.func @sparse_pointers(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> { 328 %c = arith.constant 0 : index 329 %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex> 330 return %0 : memref<?xindex> 331} 332 333// CHECK-LABEL: func @sparse_pointers64( 334// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 335// CHECK: %[[C:.*]] = arith.constant 0 : index 336// CHECK: %[[T:.*]] = call @sparsePointers64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi64> 337// CHECK: return %[[T]] : memref<?xi64> 338func.func @sparse_pointers64(%arg0: tensor<128xf64, #SparseVector64>) -> memref<?xi64> { 339 %c = arith.constant 0 : index 340 %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector64> to memref<?xi64> 341 return %0 : memref<?xi64> 342} 343 344// CHECK-LABEL: func @sparse_pointers32( 345// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 346// CHECK: %[[C:.*]] = arith.constant 0 : index 347// CHECK: %[[T:.*]] = call @sparsePointers32(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi32> 348// CHECK: return %[[T]] : memref<?xi32> 349func.func @sparse_pointers32(%arg0: tensor<128xf64, #SparseVector32>) -> memref<?xi32> { 350 %c = arith.constant 0 : index 351 %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector32> to memref<?xi32> 352 return %0 : memref<?xi32> 353} 354 355// CHECK-LABEL: func @sparse_indices( 356// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 357// CHECK: %[[C:.*]] = arith.constant 0 : index 358// CHECK: %[[T:.*]] = call @sparseIndices(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> 359// CHECK: return %[[T]] : memref<?xindex> 360func.func @sparse_indices(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> { 361 %c = arith.constant 0 : index 362 %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex> 363 return %0 : memref<?xindex> 364} 365 366// CHECK-LABEL: func @sparse_indices64( 367// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 368// CHECK: %[[C:.*]] = arith.constant 0 : index 369// CHECK: %[[T:.*]] = call @sparseIndices64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi64> 370// CHECK: return %[[T]] : memref<?xi64> 371func.func @sparse_indices64(%arg0: tensor<128xf64, #SparseVector64>) -> memref<?xi64> { 372 %c = arith.constant 0 : index 373 %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector64> to memref<?xi64> 374 return %0 : memref<?xi64> 375} 376 377// CHECK-LABEL: func @sparse_indices32( 378// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 379// CHECK: %[[C:.*]] = arith.constant 0 : index 380// CHECK: %[[T:.*]] = call @sparseIndices32(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi32> 381// CHECK: return %[[T]] : memref<?xi32> 382func.func @sparse_indices32(%arg0: tensor<128xf64, #SparseVector32>) -> memref<?xi32> { 383 %c = arith.constant 0 : index 384 %0 = sparse_tensor.indices %arg0, %c : tensor<128xf64, #SparseVector32> to memref<?xi32> 385 return %0 : memref<?xi32> 386} 387 388// CHECK-LABEL: func @sparse_valuesf64( 389// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 390// CHECK: %[[T:.*]] = call @sparseValuesF64(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf64> 391// CHECK: return %[[T]] : memref<?xf64> 392func.func @sparse_valuesf64(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xf64> { 393 %0 = sparse_tensor.values %arg0 : tensor<128xf64, #SparseVector> to memref<?xf64> 394 return %0 : memref<?xf64> 395} 396 397// CHECK-LABEL: func @sparse_valuesf32( 398// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 399// CHECK: %[[T:.*]] = call @sparseValuesF32(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf32> 400// CHECK: return %[[T]] : memref<?xf32> 401func.func @sparse_valuesf32(%arg0: tensor<128xf32, #SparseVector>) -> memref<?xf32> { 402 %0 = sparse_tensor.values %arg0: tensor<128xf32, #SparseVector> to memref<?xf32> 403 return %0 : memref<?xf32> 404} 405 406// CHECK-LABEL: func @sparse_valuesi32( 407// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 408// CHECK: %[[T:.*]] = call @sparseValuesI32(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi32> 409// CHECK: return %[[T]] : memref<?xi32> 410func.func @sparse_valuesi32(%arg0: tensor<128xi32, #SparseVector>) -> memref<?xi32> { 411 %0 = sparse_tensor.values %arg0: tensor<128xi32, #SparseVector> to memref<?xi32> 412 return %0 : memref<?xi32> 413} 414 415// CHECK-LABEL: func @sparse_valuesi16( 416// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 417// CHECK: %[[T:.*]] = call @sparseValuesI16(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi16> 418// CHECK: return %[[T]] : memref<?xi16> 419func.func @sparse_valuesi16(%arg0: tensor<128xi16, #SparseVector>) -> memref<?xi16> { 420 %0 = sparse_tensor.values %arg0: tensor<128xi16, #SparseVector> to memref<?xi16> 421 return %0 : memref<?xi16> 422} 423 424// CHECK-LABEL: func @sparse_valuesi8( 425// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) 426// CHECK: %[[T:.*]] = call @sparseValuesI8(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi8> 427// CHECK: return %[[T]] : memref<?xi8> 428func.func @sparse_valuesi8(%arg0: tensor<128xi8, #SparseVector>) -> memref<?xi8> { 429 %0 = sparse_tensor.values %arg0: tensor<128xi8, #SparseVector> to memref<?xi8> 430 return %0 : memref<?xi8> 431} 432 433// CHECK-LABEL: func @sparse_reconstruct( 434// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8> 435// CHECK: return %[[A]] : !llvm.ptr<i8> 436func.func @sparse_reconstruct(%arg0: tensor<128xf32, #SparseVector>) -> tensor<128xf32, #SparseVector> { 437 %0 = sparse_tensor.load %arg0 : tensor<128xf32, #SparseVector> 438 return %0 : tensor<128xf32, #SparseVector> 439} 440 441// CHECK-LABEL: func @sparse_reconstruct_ins( 442// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8> 443// CHECK: call @endInsert(%[[A]]) : (!llvm.ptr<i8>) -> () 444// CHECK: return %[[A]] : !llvm.ptr<i8> 445func.func @sparse_reconstruct_ins(%arg0: tensor<128xf32, #SparseVector>) -> tensor<128xf32, #SparseVector> { 446 %0 = sparse_tensor.load %arg0 hasInserts : tensor<128xf32, #SparseVector> 447 return %0 : tensor<128xf32, #SparseVector> 448} 449 450// CHECK-LABEL: func @sparse_insert( 451// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>, 452// CHECK-SAME: %[[B:.*]]: memref<?xindex>, 453// CHECK-SAME: %[[C:.*]]: f32) { 454// CHECK: call @lexInsertF32(%[[A]], %[[B]], %[[C]]) : (!llvm.ptr<i8>, memref<?xindex>, f32) -> () 455// CHECK: return 456func.func @sparse_insert(%arg0: tensor<128xf32, #SparseVector>, 457 %arg1: memref<?xindex>, 458 %arg2: f32) { 459 sparse_tensor.lex_insert %arg0, %arg1, %arg2 : tensor<128xf32, #SparseVector>, memref<?xindex>, f32 460 return 461} 462 463// CHECK-LABEL: func @sparse_expansion() 464// CHECK: %[[S:.*]] = call @sparseDimSize 465// CHECK: %[[A:.*]] = memref.alloc(%[[S]]) : memref<?xf64> 466// CHECK: %[[B:.*]] = memref.alloc(%[[S]]) : memref<?xi1> 467// CHECK: %[[C:.*]] = memref.alloc(%[[S]]) : memref<?xindex> 468// CHECK-DAG: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<?xf64>) 469// CHECK-DAG: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>) 470// CHECK: return %[[C]] : memref<?xindex> 471func.func @sparse_expansion() -> memref<?xindex> { 472 %c = arith.constant 8 : index 473 %0 = sparse_tensor.init [%c, %c] : tensor<8x8xf64, #SparseMatrix> 474 %values, %filled, %added, %count = sparse_tensor.expand %0 475 : tensor<8x8xf64, #SparseMatrix> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index 476 return %added : memref<?xindex> 477} 478 479// CHECK-LABEL: func @sparse_compression( 480// CHECK-SAME: %[[A:.*0]]: !llvm.ptr<i8>, 481// CHECK-SAME: %[[B:.*1]]: memref<?xindex>, 482// CHECK-SAME: %[[C:.*2]]: memref<?xf64>, 483// CHECK-SAME: %[[D:.*3]]: memref<?xi1>, 484// CHECK-SAME: %[[E:.*4]]: memref<?xindex>, 485// CHECK: call @expInsertF64(%[[A]], 486// CHECK-DAG: memref.dealloc %[[C]] : memref<?xf64> 487// CHECK-DAG: memref.dealloc %[[D]] : memref<?xi1> 488// CHECK-DAG: memref.dealloc %[[E]] : memref<?xindex> 489// CHECK: return 490func.func @sparse_compression(%arg0: tensor<8x8xf64, #SparseMatrix>, 491 %arg1: memref<?xindex>, %arg2: memref<?xf64>, %arg3: memref<?xi1>, 492 %arg4: memref<?xindex>, %arg5: index) { 493 sparse_tensor.compress %arg0, %arg1, %arg2, %arg3, %arg4, %arg5 494 : tensor<8x8xf64, #SparseMatrix>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index 495 return 496} 497 498// CHECK-LABEL: func @sparse_out1( 499// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>, 500// CHECK-SAME: %[[B:.*]]: !llvm.ptr<i8>) 501// CHECK-DAG: %[[ToCOO:.*]] = arith.constant 5 : i32 502// CHECK-DAG: %[[Sort:.*]] = arith.constant false 503// CHECK: %[[COO:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[ToCOO]], %[[A]]) 504// CHECK: call @outSparseTensorF64(%[[COO]], %[[B]], %[[Sort]]) : (!llvm.ptr<i8>, !llvm.ptr<i8>, i1) -> () 505// CHECK: call @delSparseTensorCOOF64(%[[COO]]) 506// CHECK: return 507func.func @sparse_out1(%arg0: tensor<?x?xf64, #SparseMatrix>, %arg1: !llvm.ptr<i8>) { 508 sparse_tensor.out %arg0, %arg1 : tensor<?x?xf64, #SparseMatrix>, !llvm.ptr<i8> 509 return 510} 511 512// CHECK-LABEL: func @sparse_out2( 513// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>, 514// CHECK-SAME: %[[B:.*]]: !llvm.ptr<i8>) 515// CHECK-DAG: %[[ToCOO:.*]] = arith.constant 5 : i32 516// CHECK-DAG: %[[Sort:.*]] = arith.constant true 517// CHECK: %[[COO:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[ToCOO]], %[[A]]) 518// CHECK: call @outSparseTensorF32(%[[COO]], %[[B]], %[[Sort]]) : (!llvm.ptr<i8>, !llvm.ptr<i8>, i1) -> () 519// CHECK: call @delSparseTensorCOOF32(%[[COO]]) 520// CHECK: return 521func.func @sparse_out2(%arg0: tensor<?x?x?xf32, #SparseTensor>, %arg1: !llvm.ptr<i8>) { 522 sparse_tensor.out %arg0, %arg1 : tensor<?x?x?xf32, #SparseTensor>, !llvm.ptr<i8> 523 return 524} 525