1// RUN: mlir-opt %s --lower-sparse-iteration-to-scf | FileCheck %s 2// RUN: mlir-opt %s --sparse-space-collapse --lower-sparse-iteration-to-scf | FileCheck %s --check-prefix COLLAPSED 3// RUN: mlir-opt %s --sparsification-and-bufferization="sparse-emit-strategy=sparse-iterator" | FileCheck %s --check-prefix COLLAPSED 4 5#COO = #sparse_tensor.encoding<{ 6 map = (i, j) -> ( 7 i : compressed(nonunique), 8 j : singleton(soa) 9 ) 10}> 11 12// CHECK-LABEL: @sparse_iteration_to_scf 13// // deduplication 14// CHECK: scf.while {{.*}} { 15// CHECK: } do { 16// CHECK: } 17// CHECK: scf.while {{.*}} { 18// CHECK: } do { 19// // actual computation 20// CHECK: scf.for {{.*}} { 21// CHECK: arith.addi 22// CHECK: } 23// // deduplication 24// CHECK: scf.while {{.*}} { 25// CHECK: } do { 26// CHECK: } 27// CHECK: scf.yield 28// CHECK: } 29// CHECK: return 30 31// COLLAPSED-LABEL: @sparse_iteration_to_scf 32// COLLAPSED: %[[RET:.*]] = scf.for {{.*}} { 33// COLLAPSED: %[[VAL:.*]] = arith.addi 34// COLLAPSED: scf.yield %[[VAL]] : index 35// COLLAPSED: } 36// COLLAPSED: return %[[RET]] : index 37func.func @sparse_iteration_to_scf(%sp : tensor<4x8xf32, #COO>) -> index { 38 %i = arith.constant 0 : index 39 %c1 = arith.constant 1 : index 40 %l1 = sparse_tensor.extract_iteration_space %sp lvls = 0 41 : tensor<4x8xf32, #COO> -> !sparse_tensor.iter_space<#COO, lvls = 0> 42 %r1 = sparse_tensor.iterate %it1 in %l1 iter_args(%outer = %i): !sparse_tensor.iter_space<#COO, lvls = 0 to 1> -> index { 43 %l2 = sparse_tensor.extract_iteration_space %sp at %it1 lvls = 1 44 : tensor<4x8xf32, #COO>, !sparse_tensor.iterator<#COO, lvls = 0 to 1> -> !sparse_tensor.iter_space<#COO, lvls = 1> 45 %r2 = sparse_tensor.iterate %it2 in %l2 iter_args(%inner = %outer): !sparse_tensor.iter_space<#COO, lvls = 1 to 2> -> index { 46 %k = arith.addi %inner, %c1 : index 47 sparse_tensor.yield %k : index 48 } 49 sparse_tensor.yield %r2 : index 50 } 51 return %r1 : index 52} 53