1// RUN: mlir-opt %s --lower-sparse-ops-to-foreach --lower-sparse-foreach-to-scf --sparse-reinterpret-map --sparse-tensor-codegen --canonicalize -cse | FileCheck %s 2 3#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }> 4 5#SparseVector = #sparse_tensor.encoding<{ 6 map = (d0) -> (d0 : compressed), 7 crdWidth = 64, 8 posWidth = 32 9}> 10 11#Dense2D = #sparse_tensor.encoding<{ 12 map = (d0, d1) -> (d0 : dense, d1 : dense), 13 crdWidth = 64, 14 posWidth = 32 15}> 16 17#Row = #sparse_tensor.encoding<{ 18 map = (d0, d1) -> (d0 : compressed, d1 : dense), 19 crdWidth = 64, 20 posWidth = 32 21}> 22 23#CSR = #sparse_tensor.encoding<{ 24 map = (d0, d1) -> (d0 : dense, d1 : compressed), 25 crdWidth = 64, 26 posWidth = 32 27}> 28 29#UCSR = #sparse_tensor.encoding<{ 30 map = (d0, d1) -> (d0 : dense, d1 : compressed(nonordered)) 31}> 32 33#CSC = #sparse_tensor.encoding<{ 34 map = (d0, d1) -> (d1 : dense, d0 : compressed) 35}> 36 37#BCSR = #sparse_tensor.encoding<{ 38 map = (d0, d1, d2, d3) -> (d0: batch, d1: batch, d2 : dense, d3 : compressed) 39}> 40 41#DCSR = #sparse_tensor.encoding<{ 42 map = (d0, d1) -> (d0 : compressed, d1 : compressed), 43 crdWidth = 64, 44 posWidth = 32 45}> 46 47#Dense3D = #sparse_tensor.encoding<{ 48 map = (d0, d1, d2) -> (d2 : dense, d0 : dense, d1 : dense) 49}> 50 51#Coo = #sparse_tensor.encoding<{ 52 map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton) 53}> 54 55#SoACOO = #sparse_tensor.encoding<{ 56 map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton(soa)) 57}> 58 59#CooPNo = #sparse_tensor.encoding<{ 60 map = (d0, d1) -> (d1 : compressed(nonunique), d0 : singleton(nonordered)) 61}> 62 63#ccoo = #sparse_tensor.encoding<{ 64 map = (d0, d1, d2) -> (d0 : compressed, d1 : compressed(nonunique), d2 : singleton) 65}> 66 67// CHECK-LABEL: func @sparse_nop( 68// CHECK-SAME: %[[A0:.*]]: memref<?xi32>, 69// CHECK-SAME: %[[A1:.*]]: memref<?xi64>, 70// CHECK-SAME: %[[A2:.*]]: memref<?xf64>, 71// CHECK-SAME: %[[A3:.*]]: !sparse_tensor.storage_specifier 72// CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] : 73// CHECK-SAME: memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier 74func.func @sparse_nop(%arg0: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { 75 return %arg0 : tensor<?xf64, #SparseVector> 76} 77 78// CHECK-LABEL: func @sparse_nop_aos_coo( 79// CHECK-SAME: %[[POS:.*0]]: memref<?xindex>, 80// CHECK-SAME: %[[AoS_CRD:.*1]]: memref<?xindex>, 81// CHECK-SAME: %[[VAL:.*]]: memref<?xf64>, 82// CHECK-SAME: %[[A3:.*]]: !sparse_tensor.storage_specifier 83// CHECK: return %[[POS]], %[[AoS_CRD]], %[[VAL]], %[[A3]] 84func.func @sparse_nop_aos_coo(%arg0: tensor<?x?xf64, #Coo>) -> tensor<?x?xf64, #Coo> { 85 return %arg0 : tensor<?x?xf64, #Coo> 86} 87 88// CHECK-LABEL: func @sparse_nop_soa_coo( 89// CHECK-SAME: %[[POS:.*0]]: memref<?xindex>, 90// CHECK-SAME: %[[SoA_CRD_0:.*1]]: memref<?xindex>, 91// CHECK-SAME: %[[SoA_CRD_1:.*2]]: memref<?xindex>, 92// CHECK-SAME: %[[VAL:.*]]: memref<?xf64>, 93// CHECK-SAME: %[[A3:.*]]: !sparse_tensor.storage_specifier 94// CHECK: return %[[POS]], %[[SoA_CRD_0]], %[[SoA_CRD_1]], %[[VAL]], %[[A3]] 95func.func @sparse_nop_soa_coo(%arg0: tensor<?x?xf64, #SoACOO>) -> tensor<?x?xf64, #SoACOO> { 96 return %arg0 : tensor<?x?xf64, #SoACOO> 97} 98 99 100// CHECK-LABEL: func @sparse_nop_multi_ret( 101// CHECK-SAME: %[[A0:.*0]]: memref<?xi32>, 102// CHECK-SAME: %[[A1:.*1]]: memref<?xi64>, 103// CHECK-SAME: %[[A2:.*2]]: memref<?xf64>, 104// CHECK-SAME: %[[A3:.*3]]: !sparse_tensor.storage_specifier 105// CHECK-SAME: %[[A4:.*4]]: memref<?xi32>, 106// CHECK-SAME: %[[A5:.*5]]: memref<?xi64>, 107// CHECK-SAME: %[[A6:.*6]]: memref<?xf64>, 108// CHECK-SAME: %[[A7:.*7]]: !sparse_tensor.storage_specifier 109// CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A7]] : 110// CHECK-SAME: memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier 111// CHECK-SAME: memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier 112func.func @sparse_nop_multi_ret(%arg0: tensor<?xf64, #SparseVector>, 113 %arg1: tensor<?xf64, #SparseVector>) -> 114 (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) { 115 return %arg0, %arg1 : tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector> 116} 117 118// CHECK-LABEL: func @sparse_nop_call( 119// CHECK-SAME: %[[A0:.*0]]: memref<?xi32>, 120// CHECK-SAME: %[[A1:.*1]]: memref<?xi64>, 121// CHECK-SAME: %[[A2:.*2]]: memref<?xf64>, 122// CHECK-SAME: %[[A3:.*3]]: !sparse_tensor.storage_specifier 123// CHECK-SAME: %[[A4:.*4]]: memref<?xi32>, 124// CHECK-SAME: %[[A5:.*5]]: memref<?xi64>, 125// CHECK-SAME: %[[A6:.*6]]: memref<?xf64>, 126// CHECK-SAME: %[[A7:.*7]]: !sparse_tensor.storage_specifier 127// CHECK: %[[T:.*]]:8 = call @sparse_nop_multi_ret(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A7]]) 128// CHECK: return %[[T]]#0, %[[T]]#1, %[[T]]#2, %[[T]]#3, %[[T]]#4, %[[T]]#5, %[[T]]#6, %[[T]]#7 : 129// CHECK-SAME: memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier 130// CHECK-SAME: memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier 131func.func @sparse_nop_call(%arg0: tensor<?xf64, #SparseVector>, 132 %arg1: tensor<?xf64, #SparseVector>) -> 133 (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) { 134 %1, %2 = call @sparse_nop_multi_ret(%arg0, %arg1) : 135 (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) -> 136 (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) 137 return %1, %2: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector> 138} 139 140// CHECK-LABEL: func @sparse_nop_cast( 141// CHECK-SAME: %[[A0:.*]]: memref<?xi32>, 142// CHECK-SAME: %[[A1:.*]]: memref<?xi64>, 143// CHECK-SAME: %[[A2:.*]]: memref<?xf32>, 144// CHECK-SAME: %[[A3:.*]]: !sparse_tensor.storage_specifier 145// CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] : 146func.func @sparse_nop_cast(%arg0: tensor<64xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> { 147 %0 = tensor.cast %arg0 : tensor<64xf32, #SparseVector> to tensor<?xf32, #SparseVector> 148 return %0 : tensor<?xf32, #SparseVector> 149} 150 151// CHECK-LABEL: func @sparse_nop_cast_3d( 152// CHECK-SAME: %[[A0:.*]]: memref<?xf32>, 153// CHECK-SAME: %[[A1:.*]]: !sparse_tensor.storage_specifier 154// CHECK: return %[[A0]], %[[A1]] : 155// CHECK-SAME: memref<?xf32>, !sparse_tensor.storage_specifier 156func.func @sparse_nop_cast_3d(%arg0: tensor<10x20x30xf32, #Dense3D>) -> tensor<?x?x?xf32, #Dense3D> { 157 %0 = tensor.cast %arg0 : tensor<10x20x30xf32, #Dense3D> to tensor<?x?x?xf32, #Dense3D> 158 return %0 : tensor<?x?x?xf32, #Dense3D> 159} 160 161// CHECK-LABEL: func @sparse_dense_2d( 162// CHECK-SAME: %[[A0:.*]]: memref<?xf64>, 163// CHECK-SAME: %[[A1:.*]]: !sparse_tensor.storage_specifier 164// CHECK: return 165func.func @sparse_dense_2d(%arg0: tensor<?x?xf64, #Dense2D>) { 166 return 167} 168 169// CHECK-LABEL: func @sparse_row( 170// CHECK-SAME: %[[A0:.*]]: memref<?xi32>, 171// CHECK-SAME: %[[A1:.*]]: memref<?xi64>, 172// CHECK-SAME: %[[A2:.*]]: memref<?xf64>, 173// CHECK-SAME: %[[A3:.*]]: !sparse_tensor.storage_specifier 174// CHECK: return 175func.func @sparse_row(%arg0: tensor<?x?xf64, #Row>) { 176 return 177} 178 179// CHECK-LABEL: func @sparse_csr( 180// CHECK-SAME: %[[A0:.*]]: memref<?xi32>, 181// CHECK-SAME: %[[A1:.*]]: memref<?xi64>, 182// CHECK-SAME: %[[A2:.*]]: memref<?xf64>, 183// CHECK-SAME: %[[A3:.*]]: !sparse_tensor.storage_specifier 184// CHECK: return 185func.func @sparse_csr(%arg0: tensor<?x?xf64, #CSR>) { 186 return 187} 188 189// CHECK-LABEL: func @sparse_bcsr_0( 190// CHECK-SAME: %[[A1:.*0]]: memref<?x2x?xindex>, 191// CHECK-SAME: %[[A2:.*1]]: memref<?x2x?xindex>, 192// CHECK-SAME: %[[A3:.*]]: memref<?x2x?xf64>, 193// CHECK-SAME: %[[A4:.*]]: !sparse_tensor.storage_specifier 194// CHECK: return 195func.func @sparse_bcsr_0(%arg0: tensor<?x2x?x?xf64, #BCSR>) { 196 return 197} 198 199// CHECK-LABEL: func @sparse_bcsr_1( 200// CHECK-SAME: %[[A1:.*0]]: memref<?x?x?xindex>, 201// CHECK-SAME: %[[A2:.*1]]: memref<?x?x?xindex>, 202// CHECK-SAME: %[[A3:.*]]: memref<?x?x?xf64>, 203// CHECK-SAME: %[[A4:.*]]: !sparse_tensor.storage_specifier 204// CHECK: return 205func.func @sparse_bcsr_1(%arg0: tensor<?x?x?x?xf64, #BCSR>) { 206 return 207} 208 209// CHECK-LABEL: func @sparse_bcsr_2( 210// CHECK-SAME: %[[A1:.*0]]: memref<18x6x?xindex>, 211// CHECK-SAME: %[[A2:.*1]]: memref<18x6x?xindex>, 212// CHECK-SAME: %[[A3:.*]]: memref<18x6x?xf64>, 213// CHECK-SAME: %[[A4:.*]]: !sparse_tensor.storage_specifier 214// CHECK: return 215func.func @sparse_bcsr_2(%arg0: tensor<18x6x4x2xf64, #BCSR>) { 216 return 217} 218 219// CHECK-LABEL: func @sparse_dcsr( 220// CHECK-SAME: %[[A0:.*0]]: memref<?xi32>, 221// CHECK-SAME: %[[A1:.*1]]: memref<?xi64>, 222// CHECK-SAME: %[[A2:.*2]]: memref<?xi32>, 223// CHECK-SAME: %[[A3:.*3]]: memref<?xi64>, 224// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>, 225// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier 226// CHECK: return 227func.func @sparse_dcsr(%arg0: tensor<?x?xf64, #DCSR>) { 228 return 229} 230 231// 232// Querying for dimension 1 in the tensor type can immediately 233// fold using the original static dimension sizes. 234// 235// CHECK-LABEL: func @sparse_dense_3d( 236// CHECK-SAME: %[[A0:.*]]: memref<?xf64>, 237// CHECK-SAME: %[[A1:.*]]: !sparse_tensor.storage_specifier 238// CHECK: %[[C:.*]] = arith.constant 20 : index 239// CHECK: return %[[C]] : index 240func.func @sparse_dense_3d(%arg0: tensor<10x20x30xf64, #Dense3D>) -> index { 241 %c = arith.constant 1 : index 242 %0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #Dense3D> 243 return %0 : index 244} 245 246// 247// Querying for dimension 1 in the tensor type needs to be permuted 248// into querying for dimension 2 in the stored sparse tensor scheme, 249// since the latter honors the dimToLvl mapping. 250// 251// CHECK-LABEL: func @sparse_dense_3d_dyn( 252// CHECK-SAME: %[[A0:.*]]: memref<?xf64>, 253// CHECK-SAME: %[[A1:.*]]: !sparse_tensor.storage_specifier 254// CHECK: %[[A2:.*]] = sparse_tensor.storage_specifier.get %[[A1]] lvl_sz at 2 255// CHECK: return %[[A2]] : index 256func.func @sparse_dense_3d_dyn(%arg0: tensor<?x?x?xf64, #Dense3D>) -> index { 257 %c = arith.constant 1 : index 258 %0 = tensor.dim %arg0, %c : tensor<?x?x?xf64, #Dense3D> 259 return %0 : index 260} 261 262// CHECK-LABEL: func @sparse_positions_dcsr( 263// CHECK-SAME: %[[A0:.*0]]: memref<?xi32>, 264// CHECK-SAME: %[[A1:.*1]]: memref<?xi64>, 265// CHECK-SAME: %[[A2:.*2]]: memref<?xi32>, 266// CHECK-SAME: %[[A3:.*3]]: memref<?xi64>, 267// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>, 268// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier 269// CHECK: %[[S:.*]] = sparse_tensor.storage_specifier.get %[[A5]] pos_mem_sz at 1 270// CHECK: %[[V:.*]] = memref.subview %[[A2]][0] [%[[S]]] [1] 271// CHECK: return %[[V]] : memref<?xi32> 272func.func @sparse_positions_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi32> { 273 %0 = sparse_tensor.positions %arg0 { level = 1 : index } : tensor<?x?xf64, #DCSR> to memref<?xi32> 274 return %0 : memref<?xi32> 275} 276 277// CHECK-LABEL: func @sparse_indices_dcsr( 278// CHECK-SAME: %[[A0:.*0]]: memref<?xi32>, 279// CHECK-SAME: %[[A1:.*1]]: memref<?xi64>, 280// CHECK-SAME: %[[A2:.*2]]: memref<?xi32>, 281// CHECK-SAME: %[[A3:.*3]]: memref<?xi64>, 282// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>, 283// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier 284// CHECK: %[[S:.*]] = sparse_tensor.storage_specifier.get %[[A5]] crd_mem_sz at 1 285// CHECK: %[[V:.*]] = memref.subview %[[A3]][0] [%[[S]]] [1] 286// CHECK: return %[[V]] : memref<?xi64> 287func.func @sparse_indices_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi64> { 288 %0 = sparse_tensor.coordinates %arg0 { level = 1 : index } : tensor<?x?xf64, #DCSR> to memref<?xi64> 289 return %0 : memref<?xi64> 290} 291 292// CHECK-LABEL: func @sparse_values_dcsr( 293// CHECK-SAME: %[[A0:.*0]]: memref<?xi32>, 294// CHECK-SAME: %[[A1:.*1]]: memref<?xi64>, 295// CHECK-SAME: %[[A2:.*2]]: memref<?xi32>, 296// CHECK-SAME: %[[A3:.*3]]: memref<?xi64>, 297// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>, 298// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier 299// CHECK: %[[S:.*]] = sparse_tensor.storage_specifier.get %[[A5]] val_mem_sz 300// CHECK: %[[V:.*]] = memref.subview %[[A4]][0] [%[[S]]] [1] 301// CHECK: return %[[V]] : memref<?xf64> 302func.func @sparse_values_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xf64> { 303 %0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64> 304 return %0 : memref<?xf64> 305} 306 307// CHECK-LABEL: func.func @sparse_values_coo( 308// CHECK-SAME: %[[A0:.*0]]: memref<?xindex>, 309// CHECK-SAME: %[[A1:.*1]]: memref<?xindex>, 310// CHECK-SAME: %[[A2:.*2]]: memref<?xindex>, 311// CHECK-SAME: %[[A3:.*3]]: memref<?xindex>, 312// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>, 313// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier 314// CHECK: %[[S:.*]] = sparse_tensor.storage_specifier.get %[[A5]] val_mem_sz 315// CHECK: %[[V:.*]] = memref.subview %[[A4]][0] [%[[S]]] [1] 316// CHECK: return %[[V]] : memref<?xf64> 317func.func @sparse_values_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xf64> { 318 %0 = sparse_tensor.values %arg0 : tensor<?x?x?xf64, #ccoo> to memref<?xf64> 319 return %0 : memref<?xf64> 320} 321 322// CHECK-LABEL: func.func @sparse_indices_coo( 323// CHECK-SAME: %[[A0:.*0]]: memref<?xindex>, 324// CHECK-SAME: %[[A1:.*1]]: memref<?xindex>, 325// CHECK-SAME: %[[A2:.*2]]: memref<?xindex>, 326// CHECK-SAME: %[[A3:.*3]]: memref<?xindex>, 327// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>, 328// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier 329// CHECK: %[[C2:.*]] = arith.constant 2 : index 330// CHECK: %[[S0:.*]] = sparse_tensor.storage_specifier.get %[[A5]] crd_mem_sz at 1 331// CHECK: %[[S2:.*]] = arith.divui %[[S0]], %[[C2]] : index 332// CHECK: %[[R1:.*]] = memref.subview %[[A3]][0] {{\[}}%[[S2]]] [2] : memref<?xindex> to memref<?xindex, strided<[2]>> 333// CHECK: %[[R2:.*]] = memref.cast %[[R1]] : memref<?xindex, strided<[2]>> to memref<?xindex, strided<[?], offset: ?>> 334// CHECK: return %[[R2]] : memref<?xindex, strided<[?], offset: ?>> 335func.func @sparse_indices_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xindex, strided<[?], offset: ?>> { 336 %0 = sparse_tensor.coordinates %arg0 { level = 1 : index } : tensor<?x?x?xf64, #ccoo> to memref<?xindex, strided<[?], offset: ?>> 337 return %0 : memref<?xindex, strided<[?], offset: ?>> 338} 339 340// CHECK-LABEL: func.func @sparse_indices_buffer_coo( 341// CHECK-SAME: %[[A0:.*0]]: memref<?xindex>, 342// CHECK-SAME: %[[A1:.*1]]: memref<?xindex>, 343// CHECK-SAME: %[[A2:.*2]]: memref<?xindex>, 344// CHECK-SAME: %[[A3:.*3]]: memref<?xindex>, 345// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>, 346// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier 347// CHECK: %[[S:.*]] = sparse_tensor.storage_specifier.get %[[A5]] crd_mem_sz at 1 348// CHECK: %[[V:.*]] = memref.subview %[[A3]][0] [%[[S]]] [1] 349// CHECK: return %[[V]] : memref<?xindex> 350func.func @sparse_indices_buffer_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xindex> { 351 %0 = sparse_tensor.coordinates_buffer %arg0 : tensor<?x?x?xf64, #ccoo> to memref<?xindex> 352 return %0 : memref<?xindex> 353} 354 355// CHECK-LABEL: func @sparse_noe( 356// CHECK-SAME: %[[A0:.*]]: memref<?xi32>, 357// CHECK-SAME: %[[A1:.*]]: memref<?xi64>, 358// CHECK-SAME: %[[A2:.*]]: memref<?xf64>, 359// CHECK-SAME: %[[A3:.*]]: !sparse_tensor.storage_specifier 360// CHECK: %[[NOE:.*]] = sparse_tensor.storage_specifier.get %[[A3]] val_mem_sz 361// CHECK: return %[[NOE]] : index 362func.func @sparse_noe(%arg0: tensor<128xf64, #SparseVector>) -> index { 363 %0 = sparse_tensor.number_of_entries %arg0 : tensor<128xf64, #SparseVector> 364 return %0 : index 365} 366 367// CHECK-LABEL: func @sparse_dealloc_csr( 368// CHECK-SAME: %[[A0:.*]]: memref<?xi32>, 369// CHECK-SAME: %[[A1:.*]]: memref<?xi64>, 370// CHECK-SAME: %[[A2:.*]]: memref<?xf64>, 371// CHECK-SAME: %[[A3:.*]]: !sparse_tensor.storage_specifier 372// CHECK: memref.dealloc %[[A0]] : memref<?xi32> 373// CHECK: memref.dealloc %[[A1]] : memref<?xi64> 374// CHECK: memref.dealloc %[[A2]] : memref<?xf64> 375// CHECK: return 376func.func @sparse_dealloc_csr(%arg0: tensor<?x?xf64, #CSR>) { 377 bufferization.dealloc_tensor %arg0 : tensor<?x?xf64, #CSR> 378 return 379} 380 381// CHECK-LABEL: func.func @sparse_alloc_csc( 382// CHECK-SAME: %[[A0:.*]]: index) -> (memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier 383// CHECK-DAG: %[[A1:.*]] = arith.constant 10 : index 384// CHECK-DAG: %[[A2:.*]] = arith.constant 0 : index 385// CHECK: %[[A3:.*]] = memref.alloc() : memref<16xindex> 386// CHECK: %[[A4:.*]] = memref.cast %[[A3]] : memref<16xindex> to memref<?xindex> 387// CHECK: %[[A5:.*]] = memref.alloc() : memref<16xindex> 388// CHECK: %[[A6:.*]] = memref.cast %[[A5]] : memref<16xindex> to memref<?xindex> 389// CHECK: %[[A7:.*]] = memref.alloc() : memref<16xf64> 390// CHECK: %[[A8:.*]] = memref.cast %[[A7]] : memref<16xf64> to memref<?xf64> 391// CHECK: %[[A9:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier 392// CHECK: %[[A11:.*]] = sparse_tensor.storage_specifier.set %[[A9]] lvl_sz at 0 with %[[A0]] : !sparse_tensor.storage_specifier 393// CHECK: %[[A12:.*]] = sparse_tensor.storage_specifier.set %[[A11]] lvl_sz at 1 with %[[A1]] : !sparse_tensor.storage_specifier 394// CHECK: %[[A14:.*]] = sparse_tensor.storage_specifier.get %[[A12]] pos_mem_sz at 1 : !sparse_tensor.storage_specifier 395// CHECK: %[[A15:.*]], %[[A17:.*]] = sparse_tensor.push_back %[[A14]], %[[A4]], %[[A2]] : index, memref<?xindex>, index 396// CHECK: %[[A18:.*]] = sparse_tensor.storage_specifier.set %[[A12]] pos_mem_sz at 1 with %[[A17]] : !sparse_tensor.storage_specifier 397// CHECK: %[[A23:.*]], %[[A25:.*]] = sparse_tensor.push_back %[[A17]], %[[A15]], %[[A2]], %[[A0]] : index, memref<?xindex>, index, index 398// CHECK: %[[A26:.*]] = sparse_tensor.storage_specifier.set %[[A18]] pos_mem_sz at 1 with %[[A25]] : !sparse_tensor.storage_specifier 399// CHECK: return %[[A23]], %[[A6]], %[[A8]], %[[A26]] : memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier 400func.func @sparse_alloc_csc(%arg0: index) -> tensor<10x?xf64, #CSC> { 401 %0 = tensor.empty(%arg0) : tensor<10x?xf64, #CSC> 402 %1 = sparse_tensor.load %0 : tensor<10x?xf64, #CSC> 403 return %1 : tensor<10x?xf64, #CSC> 404} 405 406// CHECK-LABEL: func.func @sparse_alloc_3d() -> (memref<?xf64>, !sparse_tensor.storage_specifier 407// CHECK-DAG: %[[A0:.*]] = arith.constant 6000 : index 408// CHECK-DAG: %[[A1:.*]] = arith.constant 20 : index 409// CHECK-DAG: %[[A2:.*]] = arith.constant 10 : index 410// CHECK-DAG: %[[A3:.*]] = arith.constant 30 : index 411// CHECK-DAG: %[[A4:.*]] = arith.constant 0.000000e+00 : f64 412// CHECK: %[[A5:.*]] = memref.alloc() : memref<6000xf64> 413// CHECK: %[[A6:.*]] = memref.cast %[[A5]] : memref<6000xf64> to memref<?xf64> 414// CHECK: %[[A7:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier 415// CHECK: %[[A8:.*]] = sparse_tensor.storage_specifier.set %[[A7]] lvl_sz at 0 with %[[A3]] : !sparse_tensor.storage_specifier 416// CHECK: %[[A9:.*]] = sparse_tensor.storage_specifier.set %[[A8]] lvl_sz at 1 with %[[A2]] : !sparse_tensor.storage_specifier 417// CHECK: %[[A10:.*]] = sparse_tensor.storage_specifier.set %[[A9]] lvl_sz at 2 with %[[A1]] : !sparse_tensor.storage_specifier 418// CHECK: %[[A12:.*]] = sparse_tensor.storage_specifier.get %[[A10]] val_mem_sz : !sparse_tensor.storage_specifier 419// CHECK: %[[A15:.*]], %[[A14:.*]] = sparse_tensor.push_back %[[A12]], %[[A6]], %[[A4]], %[[A0]] : index, memref<?xf64>, f64, index 420// CHECK: %[[A16:.*]] = sparse_tensor.storage_specifier.set %[[A10]] val_mem_sz with %[[A14]] : !sparse_tensor.storage_specifier 421// CHECK: return %[[A15]], %[[A16]] : memref<?xf64>, !sparse_tensor.storage_specifier 422func.func @sparse_alloc_3d() -> tensor<10x20x30xf64, #Dense3D> { 423 %0 = tensor.empty() : tensor<10x20x30xf64, #Dense3D> 424 %1 = sparse_tensor.load %0 : tensor<10x20x30xf64, #Dense3D> 425 return %1 : tensor<10x20x30xf64, #Dense3D> 426} 427 428// CHECK-LABEL: func.func @sparse_expansion1() 429// CHECK: %[[A:.*]] = memref.alloc() : memref<8xf64> 430// CHECK: %[[B:.*]] = memref.alloc() : memref<8xi1> 431// CHECK: %[[C:.*]] = memref.alloc() : memref<8xindex> 432// CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<8xindex> to memref<?xindex> 433// CHECK-DAG: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<8xf64>) 434// CHECK-DAG: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<8xi1>) 435// CHECK: return %[[D]] : memref<?xindex> 436func.func @sparse_expansion1() -> memref<?xindex> { 437 %0 = tensor.empty() : tensor<4x8xf64, #CSR> 438 %values, %filled, %added, %count = sparse_tensor.expand %0 439 : tensor<4x8xf64, #CSR> to memref<?xf64>, memref<?xi1>, memref<?xindex> 440 return %added : memref<?xindex> 441} 442 443// CHECK-LABEL: func.func @sparse_expansion2() 444// CHECK: %[[A:.*]] = memref.alloc() : memref<4xf64> 445// CHECK: %[[B:.*]] = memref.alloc() : memref<4xi1> 446// CHECK: %[[C:.*]] = memref.alloc() : memref<4xindex> 447// CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<4xindex> to memref<?xindex> 448// CHECK-DAG: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<4xf64>) 449// CHECK-DAG: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<4xi1>) 450// CHECK: return %[[D]] : memref<?xindex> 451func.func @sparse_expansion2() -> memref<?xindex> { 452 %0 = tensor.empty() : tensor<4x8xf64, #CSC> 453 %values, %filled, %added, %count = sparse_tensor.expand %0 454 : tensor<4x8xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex> 455 return %added : memref<?xindex> 456} 457 458// CHECK-LABEL: func.func @sparse_expansion3( 459// CHECK-SAME: %[[D0:.*]]: index, 460// CHECK-SAME: %{{.*}}: index) -> memref<?xindex> { 461// CHECK: %[[V:.*]] = memref.alloc(%[[D0]]) : memref<?xf64> 462// CHECK: %[[B:.*]] = memref.alloc(%[[D0]]) : memref<?xi1> 463// CHECK: %[[D:.*]] = memref.alloc(%[[D0]]) : memref<?xindex> 464// CHECK: linalg.fill ins(%{{.*}} : f64) outs(%[[V]] : memref<?xf64>) 465// CHECK: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>) 466// CHECK: return %[[D]] : memref<?xindex> 467func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref<?xindex> { 468 %0 = tensor.empty(%arg0, %arg1) : tensor<?x?xf64, #CSC> 469 %values, %filled, %added, %count = sparse_tensor.expand %0 470 : tensor<?x?xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex> 471 return %added : memref<?xindex> 472} 473 474// CHECK-LABEL: func.func private @_insert_compressed_100_f64_0_0( 475// CHECK-SAME: %[[A1:.*0]]: memref<?xindex>, 476// CHECK-SAME: %[[A2:.*1]]: memref<?xindex>, 477// CHECK-SAME: %[[A3:.*2]]: memref<?xf64>, 478// CHECK-SAME: %[[A4:.*3]]: !sparse_tensor.storage_specifier 479// CHECK-SAME: %[[A5:.*4]]: index, 480// CHECK-SAME: %[[A6:.*5]]: f64) 481// 482// CHECK-LABEL: func.func @sparse_compression_1d( 483// CHECK-SAME: %[[A0:.*0]]: memref<?xindex>, 484// CHECK-SAME: %[[A1:.*1]]: memref<?xindex>, 485// CHECK-SAME: %[[A2:.*2]]: memref<?xf64>, 486// CHECK-SAME: %[[A3:.*3]]: !sparse_tensor.storage_specifier 487// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>, 488// CHECK-SAME: %[[A5:.*5]]: memref<?xi1>, 489// CHECK-SAME: %[[A6:.*6]]: memref<?xindex>, 490// CHECK-SAME: %[[A7:.*7]]: index) -> (memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier 491// CHECK-DAG: %[[A8:.*]] = arith.constant false 492// CHECK-DAG: %[[A9:.*]] = arith.constant 0.000000e+00 : f64 493// CHECK-DAG: %[[A10:.*]] = arith.constant 1 : index 494// CHECK-DAG: %[[A11:.*]] = arith.constant 0 : index 495// CHECK: sparse_tensor.sort hybrid_quick_sort %[[A7]], %[[A6]] 496// CHECK: %[[A12:.*]]:4 = scf.for %[[A13:.*]] = %[[A11]] to %[[A7]] step %[[A10]] iter_args(%[[A14:.*]] = %[[A0]], %[[A15:.*]] = %[[A1]], %[[A16:.*]] = %[[A2]], %[[A17:.*]] = %[[A3]]) 497// CHECK: %[[A18:.*]] = memref.load %[[A6]]{{\[}}%[[A13]]] : memref<?xindex> 498// CHECK: %[[A19:.*]] = memref.load %[[A4]]{{\[}}%[[A18]]] : memref<?xf64> 499// CHECK: %[[A20:.*]]:4 = func.call @_insert_compressed_100_f64_0_0(%[[A14]], %[[A15]], %[[A16]], %[[A17]], %[[A18]], %[[A19]]) 500// CHECK: memref.store %[[A9]], %[[A4]]{{\[}}%[[A18]]] : memref<?xf64> 501// CHECK: memref.store %[[A8]], %[[A5]]{{\[}}%[[A18]]] : memref<?xi1> 502// CHECK: scf.yield %[[A20]]#0, %[[A20]]#1, %[[A20]]#2, %[[A20]]#3 503// CHECK: } 504// CHECK: memref.dealloc %[[A4]] : memref<?xf64> 505// CHECK: memref.dealloc %[[A5]] : memref<?xi1> 506// CHECK: memref.dealloc %[[A6]] : memref<?xindex> 507// CHECK: return %[[A21:.*]]#0, %[[A21]]#1, %[[A21]]#2, %[[A21]]#3 : memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier 508func.func @sparse_compression_1d(%tensor: tensor<100xf64, #SV>, 509 %values: memref<?xf64>, 510 %filled: memref<?xi1>, 511 %added: memref<?xindex>, 512 %count: index) -> tensor<100xf64, #SV> { 513 %0 = sparse_tensor.compress %values, %filled, %added, %count into %tensor[] 514 : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<100xf64, #SV> 515 %1 = sparse_tensor.load %0 hasInserts : tensor<100xf64, #SV> 516 return %1 : tensor<100xf64, #SV> 517} 518 519// CHECK-LABEL: func.func private @_insert_dense_compressed_8_8_f64_64_32( 520// CHECK-SAME: %[[A1:.*0]]: memref<?xi32>, 521// CHECK-SAME: %[[A2:.*1]]: memref<?xi64>, 522// CHECK-SAME: %[[A3:.*2]]: memref<?xf64>, 523// CHECK-SAME: %[[A4:.*3]]: !sparse_tensor.storage_specifier 524// CHECK-SAME: %[[A5:.*4]]: index, 525// CHECK-SAME: %[[A6:.*5]]: index, 526// CHECK-SAME: %[[A7:.*6]]: f64) 527// 528// CHECK-LABEL: func.func @sparse_compression( 529// CHECK-SAME: %[[A0:.*0]]: memref<?xi32>, 530// CHECK-SAME: %[[A1:.*1]]: memref<?xi64>, 531// CHECK-SAME: %[[A2:.*2]]: memref<?xf64>, 532// CHECK-SAME: %[[A3:.*3]]: !sparse_tensor.storage_specifier 533// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>, 534// CHECK-SAME: %[[A5:.*5]]: memref<?xi1>, 535// CHECK-SAME: %[[A6:.*6]]: memref<?xindex>, 536// CHECK-SAME: %[[A7:.*7]]: index, 537// CHECK-SAME: %[[A8:.*8]]: index) -> (memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier 538// CHECK: %[[A9:.*]] = arith.constant 0 : i32 539// CHECK: %[[A10:.*]] = arith.constant false 540// CHECK: %[[A11:.*]] = arith.constant 0.000000e+00 : f64 541// CHECK: %[[A12:.*]] = arith.constant 1 : index 542// CHECK: %[[A13:.*]] = arith.constant 0 : index 543// CHECK: sparse_tensor.sort hybrid_quick_sort %[[A7]], %[[A6]] 544// CHECK: %[[A14:.*]]:4 = scf.for %[[A15:.*]] = %[[A13]] to %[[A7]] step %[[A12]] iter_args(%[[A16:.*]] = %[[A0]], %[[A17:.*]] = %[[A1]], %[[A18:.*]] = %[[A2]], %[[A19:.*]] = %[[A3]]) -> (memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier 545// CHECK: %[[A20:.*]] = memref.load %[[A6]]{{\[}}%[[A15]]] : memref<?xindex> 546// CHECK: %[[A21:.*]] = memref.load %[[A4]]{{\[}}%[[A20]]] : memref<?xf64> 547// CHECK: %[[A22:.*]]:4 = func.call @_insert_dense_compressed_8_8_f64_64_32(%[[A16]], %[[A17]], %[[A18]], %[[A19]], %[[A8]], %[[A20]], %[[A21]]) : (memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier 548// CHECK: memref.store %[[A11]], %[[A4]]{{\[}}%[[A20]]] : memref<?xf64> 549// CHECK: memref.store %[[A10]], %[[A5]]{{\[}}%[[A20]]] : memref<?xi1> 550// CHECK: scf.yield %[[A22]]#0, %[[A22]]#1, %[[A22]]#2, %[[A22]]#3 : memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier 551// CHECK: } 552// CHECK: memref.dealloc %[[A4]] : memref<?xf64> 553// CHECK: memref.dealloc %[[A5]] : memref<?xi1> 554// CHECK: memref.dealloc %[[A6]] : memref<?xindex> 555// CHECK: %[[A25:.*]] = sparse_tensor.storage_specifier.get %[[A24:.*]]#3 pos_mem_sz at 1 : !sparse_tensor.storage_specifier 556// CHECK: %[[A26:.*]] = memref.load %[[A24]]#0{{\[}}%[[A13]]] : memref<?xi32> 557// CHECK: %[[A27:.*]] = scf.for %[[A28:.*]] = %[[A12]] to %[[A25]] step %[[A12]] iter_args(%[[A29:.*]] = %[[A26]]) -> (i32) { 558// CHECK: %[[A30:.*]] = memref.load %[[A24]]#0{{\[}}%[[A28]]] : memref<?xi32> 559// CHECK: %[[A31:.*]] = arith.cmpi eq, %[[A30]], %[[A9]] : i32 560// CHECK: %[[A32:.*]] = arith.select %[[A31]], %[[A29]], %[[A30]] : i32 561// CHECK: scf.if %[[A31]] { 562// CHECK: memref.store %[[A29]], %[[A24]]#0{{\[}}%[[A28]]] : memref<?xi32> 563// CHECK: } 564// CHECK: scf.yield %[[A32]] : i32 565// CHECK: } 566// CHECK: return %[[A24]]#0, %[[A24]]#1, %[[A24]]#2, %[[A24]]#3 : memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier 567func.func @sparse_compression(%tensor: tensor<8x8xf64, #CSR>, 568 %values: memref<?xf64>, 569 %filled: memref<?xi1>, 570 %added: memref<?xindex>, 571 %count: index, 572 %i: index) -> tensor<8x8xf64, #CSR> { 573 %0 = sparse_tensor.compress %values, %filled, %added, %count into %tensor[%i] 574 : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<8x8xf64, #CSR> 575 %1 = sparse_tensor.load %0 hasInserts : tensor<8x8xf64, #CSR> 576 return %1 : tensor<8x8xf64, #CSR> 577} 578 579// CHECK-LABEL: func.func private @_insert_dense_compressed_nonordered_8_8_f64_0_0( 580// CHECK-SAME: %[[A1:.*0]]: memref<?xindex>, 581// CHECK-SAME: %[[A2:.*1]]: memref<?xindex>, 582// CHECK-SAME: %[[A3:.*2]]: memref<?xf64>, 583// CHECK-SAME: %[[A4:.*3]]: !sparse_tensor.storage_specifier 584// CHECK-SAME: %[[A5:.*4]]: index, 585// CHECK-SAME: %[[A6:.*5]]: index, 586// CHECK-SAME: %[[A7:.*6]]: f64) 587// 588// CHECK-LABEL: func.func @sparse_compression_unordered( 589// CHECK-SAME: %[[A0:.*0]]: memref<?xindex>, 590// CHECK-SAME: %[[A1:.*1]]: memref<?xindex>, 591// CHECK-SAME: %[[A2:.*2]]: memref<?xf64>, 592// CHECK-SAME: %[[A3:.*3]]: !sparse_tensor.storage_specifier 593// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>, 594// CHECK-SAME: %[[A5:.*5]]: memref<?xi1>, 595// CHECK-SAME: %[[A6:.*6]]: memref<?xindex>, 596// CHECK-SAME: %[[A7:.*7]]: index, 597// CHECK-SAME: %[[A8:.*8]]: index) -> (memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier 598// CHECK: %[[A9:.*]] = arith.constant false 599// CHECK: %[[A10:.*]] = arith.constant 0.000000e+00 : f64 600// CHECK: %[[A11:.*]] = arith.constant 0 : index 601// CHECK: %[[A12:.*]] = arith.constant 1 : index 602// CHECK: %[[A13:.*]]:4 = scf.for %[[A14:.*]] = %[[A11]] to %[[A7]] step %[[A12]] iter_args(%[[A15:.*]] = %[[A0]], %[[A16:.*]] = %[[A1]], %[[A17:.*]] = %[[A2]], %[[A18:.*]] = %[[A3]]) -> (memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier 603// CHECK: %[[A19:.*]] = memref.load %[[A6]]{{\[}}%[[A14]]] : memref<?xindex> 604// CHECK: %[[A20:.*]] = memref.load %[[A4]]{{\[}}%[[A19]]] : memref<?xf64> 605// CHECK: %[[A21:.*]]:4 = func.call @_insert_dense_compressed_nonordered_8_8_f64_0_0(%[[A15]], %[[A16]], %[[A17]], %[[A18]], %[[A8]], %[[A19]], %[[A20]]) : (memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier 606// CHECK: memref.store %[[A10]], %[[A4]]{{\[}}%[[A19]]] : memref<?xf64> 607// CHECK: memref.store %[[A9]], %[[A5]]{{\[}}%[[A19]]] : memref<?xi1> 608// CHECK: scf.yield %[[A21]]#0, %[[A21]]#1, %[[A21]]#2, %[[A21]]#3 : memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier 609// CHECK: } 610// CHECK: memref.dealloc %[[A4]] : memref<?xf64> 611// CHECK: memref.dealloc %[[A5]] : memref<?xi1> 612// CHECK: memref.dealloc %[[A6]] : memref<?xindex> 613// CHECK: %[[A24:.*]] = sparse_tensor.storage_specifier.get %[[A23:.*]]#3 pos_mem_sz at 1 : !sparse_tensor.storage_specifier 614// CHECK: %[[A25:.*]] = memref.load %[[A23]]#0{{\[}}%[[A11]]] : memref<?xindex> 615// CHECK: %[[A26:.*]] = scf.for %[[A27:.*]] = %[[A12]] to %[[A24]] step %[[A12]] iter_args(%[[A28:.*]] = %[[A25]]) -> (index) { 616// CHECK: %[[A29:.*]] = memref.load %[[A23]]#0{{\[}}%[[A27]]] : memref<?xindex> 617// CHECK: %[[A30:.*]] = arith.cmpi eq, %[[A29]], %[[A11]] : index 618// CHECK: %[[A31:.*]] = arith.select %[[A30]], %[[A28]], %[[A29]] : index 619// CHECK: scf.if %[[A30]] { 620// CHECK: memref.store %[[A28]], %[[A23]]#0{{\[}}%[[A27]]] : memref<?xindex> 621// CHECK: } 622// CHECK: scf.yield %[[A31]] : index 623// CHECK: } 624// CHECK: return %[[A23]]#0, %[[A23]]#1, %[[A23]]#2, %[[A23]]#3 : memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier 625func.func @sparse_compression_unordered(%tensor: tensor<8x8xf64, #UCSR>, 626 %values: memref<?xf64>, 627 %filled: memref<?xi1>, 628 %added: memref<?xindex>, 629 %count: index, 630 %i: index) -> tensor<8x8xf64, #UCSR> { 631 %0 = sparse_tensor.compress %values, %filled, %added, %count into %tensor[%i] 632 : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<8x8xf64, #UCSR> 633 %1 = sparse_tensor.load %0 hasInserts : tensor<8x8xf64, #UCSR> 634 return %1 : tensor<8x8xf64, #UCSR> 635} 636 637// CHECK-LABEL: func.func private @_insert_compressed_128_f64_0_0( 638// CHECK-SAME: %[[A1:.*0]]: memref<?xindex>, 639// CHECK-SAME: %[[A2:.*1]]: memref<?xindex>, 640// CHECK-SAME: %[[A3:.*2]]: memref<?xf64>, 641// CHECK-SAME: %[[A4:.*3]]: !sparse_tensor.storage_specifier 642// CHECK-SAME: %[[A5:.*4]]: index, 643// CHECK-SAME: %[[A6:.*5]]: f64) 644// 645// CHECK-LABEL: func @sparse_insert( 646// CHECK-SAME: %[[A1:.*0]]: memref<?xindex>, 647// CHECK-SAME: %[[A2:.*1]]: memref<?xindex>, 648// CHECK-SAME: %[[A3:.*2]]: memref<?xf64>, 649// CHECK-SAME: %[[A4:.*3]]: !sparse_tensor.storage_specifier 650// CHECK-SAME: %[[A5:.*4]]: index, 651// CHECK-SAME: %[[A6:.*5]]: f64) 652// CHECK: %[[R:.*]]:4 = call @_insert_compressed_128_f64_0_0(%[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]]) 653// CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3 654func.func @sparse_insert(%arg0: tensor<128xf64, #SV>, %arg1: index, %arg2: f64) -> tensor<128xf64, #SV> { 655 %0 = tensor.insert %arg2 into %arg0[%arg1] : tensor<128xf64, #SV> 656 %1 = sparse_tensor.load %0 hasInserts : tensor<128xf64, #SV> 657 return %1 : tensor<128xf64, #SV> 658} 659 660// CHECK-LABEL: func.func private @_insert_compressed_128_f64_64_32( 661// CHECK-SAME: %[[A1:.*]]: memref<?xi32>, 662// CHECK-SAME: %[[A2:.*]]: memref<?xi64>, 663// CHECK-SAME: %[[A3:.*]]: memref<?xf64>, 664// CHECK-SAME: %[[A4:.*]]: !sparse_tensor.storage_specifier 665// CHECK-SAME: %[[A5:.*]]: index, 666// CHECK-SAME: %[[A6:.*]]: f64) 667// 668// CHECK-LABEL: func @sparse_insert_typed( 669// CHECK-SAME: %[[A1:.*]]: memref<?xi32>, 670// CHECK-SAME: %[[A2:.*]]: memref<?xi64>, 671// CHECK-SAME: %[[A3:.*]]: memref<?xf64>, 672// CHECK-SAME: %[[A4:.*]]: !sparse_tensor.storage_specifier 673// CHECK-SAME: %[[A5:.*]]: index, 674// CHECK-SAME: %[[A6:.*]]: f64) 675// CHECK: %[[R:.*]]:4 = call @_insert_compressed_128_f64_64_32(%[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]]) 676// CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3 677func.func @sparse_insert_typed(%arg0: tensor<128xf64, #SparseVector>, %arg1: index, %arg2: f64) -> tensor<128xf64, #SparseVector> { 678 %0 = tensor.insert %arg2 into %arg0[%arg1] : tensor<128xf64, #SparseVector> 679 %1 = sparse_tensor.load %0 hasInserts : tensor<128xf64, #SparseVector> 680 return %1 : tensor<128xf64, #SparseVector> 681} 682 683// CHECK-LABEL: func.func private @_insert_compressed_nonunique_singleton_5_6_f64_0_0( 684// CHECK-SAME: %[[A1:.*0]]: memref<?xindex>, 685// CHECK-SAME: %[[A2:.*1]]: memref<?xindex>, 686// CHECK-SAME: %[[A3:.*2]]: memref<?xf64>, 687// CHECK-SAME: %[[A4:.*3]]: !sparse_tensor.storage_specifier 688// CHECK-SAME: %[[A5:.*4]]: index, 689// CHECK-SAME: %[[A5:.*5]]: index, 690// CHECK-SAME: %[[A7:.*6]]: f64) 691// 692// CHECK-LABEL: func.func @sparse_insert_coo( 693// CHECK-SAME: %[[A0:.*0]]: memref<?xindex>, 694// CHECK-SAME: %[[A1:.*1]]: memref<?xindex>, 695// CHECK-SAME: %[[A2:.*2]]: memref<?xf64>, 696// CHECK-SAME: %[[A3:.*3]]: !sparse_tensor.storage_specifier 697// CHECK-SAME: %[[A4:.*4]]: index, 698// CHECK-SAME: %[[A5:.*5]]: f64) 699// CHECK: %[[R:.*]]:4 = call @_insert_compressed_nonunique_singleton_5_6_f64_0_0(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A4]], %[[A5]]) 700// CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3 701func.func @sparse_insert_coo(%arg0: tensor<5x6xf64, #Coo>, %arg1: index, %arg2: f64) -> tensor<5x6xf64, #Coo> { 702 %0 = tensor.insert %arg2 into %arg0[%arg1, %arg1] : tensor<5x6xf64, #Coo> 703 %1 = sparse_tensor.load %0 hasInserts : tensor<5x6xf64, #Coo> 704 return %1 : tensor<5x6xf64, #Coo> 705} 706 707// CHECK-LABEL: func.func @sparse_nop_convert( 708// CHECK-SAME: %[[A1:.*]]: memref<?xi32>, 709// CHECK-SAME: %[[A2:.*]]: memref<?xi64>, 710// CHECK-SAME: %[[A3:.*]]: memref<?xf32>, 711// CHECK-SAME: %[[A4:.*]]: !sparse_tensor.storage_specifier 712// CHECK: return %[[A1]], %[[A2]], %[[A3]], %[[A4]] : 713// CHECK-SAME: memref<?xi32>, memref<?xi64>, memref<?xf32>, !sparse_tensor.storage_specifier 714func.func @sparse_nop_convert(%arg0: tensor<32xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> { 715 %0 = sparse_tensor.convert %arg0 : tensor<32xf32, #SparseVector> to tensor<?xf32, #SparseVector> 716 return %0 : tensor<?xf32, #SparseVector> 717} 718 719// CHECK-LABEL: func.func @sparse_convert_element_type( 720// CHECK-SAME: %[[A1:.*]]: memref<?xi32>, 721// CHECK-SAME: %[[A2:.*]]: memref<?xi64>, 722// CHECK-SAME: %[[A3:.*]]: memref<?xf32>, 723// CHECK-SAME: %[[A4:.*]]: !sparse_tensor.storage_specifier 724// CHECK: scf.for 725// CHECK: %[[FValue:.*]] = memref.load 726// CHECK: %[[IValue:.*]] = arith.fptosi %[[FValue]] 727// CHECK: memref.store %[[IValue]] 728// CHECK: return %{{.*}}, %{{.*}}, %{{.*}}, %[[A4]] : 729// CHECK-SAME: memref<?xi32>, memref<?xi64>, memref<?xi32>, !sparse_tensor.storage_specifier 730func.func @sparse_convert_element_type(%arg0: tensor<32xf32, #SparseVector>) -> tensor<?xi32, #SparseVector> { 731 %0 = sparse_tensor.convert %arg0 : tensor<32xf32, #SparseVector> to tensor<?xi32, #SparseVector> 732 return %0 : tensor<?xi32, #SparseVector> 733} 734 735// CHECK-LABEL: func.func @sparse_new_coo( 736// CHECK-SAME: %[[A0:.*]]: !llvm.ptr) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier<#sparse{{[0-9]*}}>) { 737// CHECK-DAG: %[[VAL_1:.*]] = arith.constant false 738// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 2 : i32 739// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 1 : index 740// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index 741// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 2 : index 742// CHECK: %[[VAL_6:.*]] = memref.alloca() : memref<2xindex> 743// CHECK: %[[VAL_7:.*]] = memref.cast %[[VAL_6]] : memref<2xindex> to memref<?xindex> 744// CHECK: memref.store %[[VAL_4]], %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<2xindex> 745// CHECK: memref.store %[[VAL_4]], %[[VAL_6]]{{\[}}%[[VAL_3]]] : memref<2xindex> 746// CHECK: %[[VAL_8:.*]] = call @createCheckedSparseTensorReader(%[[A0]], %[[VAL_7]], %[[VAL_2]]) : (!llvm.ptr, memref<?xindex>, i32) -> !llvm.ptr 747// CHECK: %[[VAL_9:.*]] = call @getSparseTensorReaderDimSizes(%[[VAL_8]]) : (!llvm.ptr) -> memref<?xindex> 748// CHECK-DAG: %[[VAL_10:.*]] = call @getSparseTensorReaderNSE(%[[VAL_8]]) : (!llvm.ptr) -> index 749// CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex> 750// CHECK-DAG: %[[VAL_12:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_3]]] : memref<?xindex> 751// CHECK-DAG: %[[VAL_13:.*]] = arith.muli %[[VAL_10]], %[[VAL_5]] : index 752// CHECK-DAG: %[[VAL_14:.*]] = memref.alloc() : memref<2xindex> 753// CHECK-DAG: %[[VAL_15:.*]] = memref.cast %[[VAL_14]] : memref<2xindex> to memref<?xindex> 754// CHECK-DAG: %[[VAL_16:.*]] = memref.alloc(%[[VAL_13]]) : memref<?xindex> 755// CHECK-DAG: %[[VAL_17:.*]] = memref.alloc(%[[VAL_10]]) : memref<?xf32> 756// CHECK-DAG: %[[VAL_18:.*]] = sparse_tensor.storage_specifier.init 757// CHECK-DAG: %[[VAL_19:.*]] = sparse_tensor.storage_specifier.set %[[VAL_18]] lvl_sz at 0 with %[[VAL_11]] 758// CHECK-DAG: %[[VAL_20:.*]] = sparse_tensor.storage_specifier.get %[[VAL_19]] pos_mem_sz at 0 759// CHECK-DAG: %[[VAL_21:.*]], %[[VAL_22:.*]] = sparse_tensor.push_back %[[VAL_20]], %[[VAL_15]], %[[VAL_4]] 760// CHECK-DAG: %[[VAL_23:.*]] = sparse_tensor.storage_specifier.set %[[VAL_19]] pos_mem_sz at 0 with %[[VAL_22]] 761// CHECK-DAG: %[[VAL_24:.*]] = sparse_tensor.storage_specifier.set %[[VAL_23]] lvl_sz at 1 with %[[VAL_12]] 762// CHECK-DAG: %[[VAL_25:.*]], %[[VAL_26:.*]] = sparse_tensor.push_back %[[VAL_22]], %[[VAL_21]], %[[VAL_4]], %[[VAL_3]] 763// CHECK-DAG: %[[VAL_27:.*]] = sparse_tensor.storage_specifier.set %[[VAL_24]] pos_mem_sz at 0 with %[[VAL_26]] 764// CHECK-DAG: %[[VAL_28:.*]] = memref.alloca() : memref<2xindex> 765// CHECK-DAG: %[[VAL_29:.*]] = memref.cast %[[VAL_28]] : memref<2xindex> to memref<?xindex> 766// CHECK-DAG: memref.store %[[VAL_4]], %[[VAL_28]]{{\[}}%[[VAL_4]]] : memref<2xindex> 767// CHECK-DAG: memref.store %[[VAL_3]], %[[VAL_28]]{{\[}}%[[VAL_3]]] : memref<2xindex> 768// CHECK: %[[VAL_30:.*]] = call @getSparseTensorReaderReadToBuffers0F32(%[[VAL_8]], %[[VAL_29]], %[[VAL_29]], %[[VAL_16]], %[[VAL_17]]) : (!llvm.ptr, memref<?xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>) -> i1 769// CHECK: %[[VAL_31:.*]] = arith.cmpi eq, %[[VAL_30]], %[[VAL_1]] : i1 770// CHECK: scf.if %[[VAL_31]] { 771// CHECK: sparse_tensor.sort hybrid_quick_sort %[[VAL_10]], %[[VAL_16]] jointly %[[VAL_17]] 772// CHECK: } 773// CHECK: memref.store %[[VAL_10]], %[[VAL_25]]{{\[}}%[[VAL_3]]] : memref<?xindex> 774// CHECK: %[[VAL_32:.*]] = sparse_tensor.storage_specifier.set %[[VAL_27]] crd_mem_sz at 0 with %[[VAL_13]] 775// CHECK: %[[VAL_33:.*]] = sparse_tensor.storage_specifier.set %[[VAL_32]] val_mem_sz with %[[VAL_10]] 776// CHECK: call @delSparseTensorReader(%[[VAL_8]]) : (!llvm.ptr) -> () 777// CHECK: return %[[VAL_25]], %[[VAL_16]], %[[VAL_17]], %[[VAL_33]] 778func.func @sparse_new_coo(%arg0: !llvm.ptr) -> tensor<?x?xf32, #Coo> { 779 %0 = sparse_tensor.new %arg0 : !llvm.ptr to tensor<?x?xf32, #Coo> 780 return %0 : tensor<?x?xf32, #Coo> 781} 782 783// CHECK-LABEL: func.func @sparse_new_coo_permute_no( 784// CHECK-SAME: %[[A0:.*]]: !llvm.ptr) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier<#sparse{{[0-9]*}}>) { 785// CHECK-DAG: %[[VAL_1:.*]] = arith.constant 2 : i32 786// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 1 : index 787// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index 788// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 2 : index 789// CHECK: %[[VAL_5:.*]] = memref.alloca() : memref<2xindex> 790// CHECK: %[[VAL_6:.*]] = memref.cast %[[VAL_5]] : memref<2xindex> to memref<?xindex> 791// CHECK: memref.store %[[VAL_3]], %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<2xindex> 792// CHECK: memref.store %[[VAL_3]], %[[VAL_5]]{{\[}}%[[VAL_2]]] : memref<2xindex> 793// CHECK: %[[VAL_7:.*]] = call @createCheckedSparseTensorReader(%[[A0]], %[[VAL_6]], %[[VAL_1]]) : (!llvm.ptr, memref<?xindex>, i32) -> !llvm.ptr 794// CHECK-DAG: %[[VAL_8:.*]] = call @getSparseTensorReaderDimSizes(%[[VAL_7]]) : (!llvm.ptr) -> memref<?xindex> 795// CHECK-DAG: %[[VAL_9:.*]] = call @getSparseTensorReaderNSE(%[[VAL_7]]) : (!llvm.ptr) -> index 796// CHECK-DAG: %[[VAL_10:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex> 797// CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_2]]] : memref<?xindex> 798// CHECK-DAG: %[[VAL_12:.*]] = arith.muli %[[VAL_9]], %[[VAL_4]] : index 799// CHECK-DAG: %[[VAL_13:.*]] = memref.alloc() : memref<2xindex> 800// CHECK-DAG: %[[VAL_14:.*]] = memref.cast %[[VAL_13]] : memref<2xindex> to memref<?xindex> 801// CHECK-DAG: %[[VAL_15:.*]] = memref.alloc(%[[VAL_12]]) : memref<?xindex> 802// CHECK-DAG: %[[VAL_16:.*]] = memref.alloc(%[[VAL_9]]) : memref<?xf32> 803// CHECK-DAG: %[[VAL_17:.*]] = sparse_tensor.storage_specifier.init 804// CHECK-DAG: %[[VAL_18:.*]] = sparse_tensor.storage_specifier.set %[[VAL_17]] lvl_sz at 0 with %[[VAL_11]] 805// CHECK-DAG: %[[VAL_19:.*]] = sparse_tensor.storage_specifier.get %[[VAL_18]] pos_mem_sz at 0 806// CHECK-DAG: %[[VAL_20:.*]], %[[VAL_21:.*]] = sparse_tensor.push_back %[[VAL_19]], %[[VAL_14]], %[[VAL_3]] 807// CHECK-DAG: %[[VAL_22:.*]] = sparse_tensor.storage_specifier.set %[[VAL_18]] pos_mem_sz at 0 with %[[VAL_21]] 808// CHECK-DAG: %[[VAL_23:.*]] = sparse_tensor.storage_specifier.set %[[VAL_22]] lvl_sz at 1 with %[[VAL_10]] 809// CHECK-DAG: %[[VAL_24:.*]], %[[VAL_25:.*]] = sparse_tensor.push_back %[[VAL_21]], %[[VAL_20]], %[[VAL_3]], %[[VAL_2]] 810// CHECK-DAG: %[[VAL_26:.*]] = sparse_tensor.storage_specifier.set %[[VAL_23]] pos_mem_sz at 0 with %[[VAL_25]] 811// CHECK-DAG: %[[VAL_27:.*]] = memref.alloca() : memref<2xindex> 812// CHECK-DAG: %[[VAL_28:.*]] = memref.cast %[[VAL_27]] : memref<2xindex> to memref<?xindex> 813// CHECK-DAG: memref.store %[[VAL_2]], %[[VAL_27]]{{\[}}%[[VAL_3]]] : memref<2xindex> 814// CHECK-DAG: memref.store %[[VAL_3]], %[[VAL_27]]{{\[}}%[[VAL_2]]] : memref<2xindex> 815// CHECK-DAG: %[[VAL_29:.*]] = memref.alloca() : memref<2xindex> 816// CHECK-DAG: %[[VAL_30:.*]] = memref.cast %[[VAL_29]] : memref<2xindex> to memref<?xindex> 817// CHECK-DAG: memref.store %[[VAL_2]], %[[VAL_29]]{{\[}}%[[VAL_3]]] : memref<2xindex> 818// CHECK-DAG: memref.store %[[VAL_3]], %[[VAL_29]]{{\[}}%[[VAL_2]]] : memref<2xindex> 819// CHECK: %[[VAL_31:.*]] = call @getSparseTensorReaderReadToBuffers0F32(%[[VAL_7]], %[[VAL_28]], %[[VAL_30]], %[[VAL_15]], %[[VAL_16]]) : (!llvm.ptr, memref<?xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>) -> i1 820// CHECK: memref.store %[[VAL_9]], %[[VAL_24]]{{\[}}%[[VAL_2]]] : memref<?xindex> 821// CHECK: %[[VAL_32:.*]] = sparse_tensor.storage_specifier.set %[[VAL_26]] crd_mem_sz at 0 with %[[VAL_12]] 822// CHECK: %[[VAL_33:.*]] = sparse_tensor.storage_specifier.set %[[VAL_32]] val_mem_sz with %[[VAL_9]] 823// CHECK: call @delSparseTensorReader(%[[VAL_7]]) : (!llvm.ptr) -> () 824// CHECK: return %[[VAL_24]], %[[VAL_15]], %[[VAL_16]], %[[VAL_33]] 825func.func @sparse_new_coo_permute_no(%arg0: !llvm.ptr) -> tensor<?x?xf32, #CooPNo> { 826 %0 = sparse_tensor.new %arg0 : !llvm.ptr to tensor<?x?xf32, #CooPNo> 827 return %0 : tensor<?x?xf32, #CooPNo> 828} 829