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