53 lines · plain
1// RUN: mlir-opt %s --lower-sparse-iteration-to-scf | FileCheck %s2// RUN: mlir-opt %s --sparse-space-collapse --lower-sparse-iteration-to-scf | FileCheck %s --check-prefix COLLAPSED3// RUN: mlir-opt %s --sparsification-and-bufferization="sparse-emit-strategy=sparse-iterator" | FileCheck %s --check-prefix COLLAPSED4 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_scf13// // deduplication14// CHECK: scf.while {{.*}} {15// CHECK: } do {16// CHECK: }17// CHECK: scf.while {{.*}} {18// CHECK: } do {19// // actual computation20// CHECK: scf.for {{.*}} {21// CHECK: arith.addi22// CHECK: }23// // deduplication24// CHECK: scf.while {{.*}} {25// CHECK: } do {26// CHECK: }27// CHECK: scf.yield28// CHECK: }29// CHECK: return30 31// COLLAPSED-LABEL: @sparse_iteration_to_scf32// COLLAPSED: %[[RET:.*]] = scf.for {{.*}} {33// COLLAPSED: %[[VAL:.*]] = arith.addi34// COLLAPSED: scf.yield %[[VAL]] : index35// COLLAPSED: }36// COLLAPSED: return %[[RET]] : index37func.func @sparse_iteration_to_scf(%sp : tensor<4x8xf32, #COO>) -> index {38 %i = arith.constant 0 : index39 %c1 = arith.constant 1 : index40 %l1 = sparse_tensor.extract_iteration_space %sp lvls = 041 : 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 = 144 : 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 : index47 sparse_tensor.yield %k : index48 }49 sparse_tensor.yield %r2 : index50 }51 return %r1 : index52}53