[mlir] Check 'iter_args' in 'isLoopParallel' utility

Fix 'isLoopParallel' utility so that 'iter_args' is taken into account
and loops with loop-carried dependences are not classified as parallel.

Reviewed By: tungld, vinayaka-polymage

Differential Revision: https://reviews.llvm.org/D97347
This commit is contained in:
Diego Caballero
2021-02-25 17:56:51 +02:00
parent 203d5eeec5
commit ebca222b65
2 changed files with 31 additions and 0 deletions

View File

@@ -1173,6 +1173,12 @@ void mlir::getSequentialLoops(AffineForOp forOp,
/// Returns true if 'forOp' is parallel.
bool mlir::isLoopParallel(AffineForOp forOp) {
// Loop is not parallel if it has SSA loop-carried dependences.
// TODO: Conditionally support reductions and other loop-carried dependences
// that could be handled in the context of a parallel loop.
if (forOp.getNumIterOperands() > 0)
return false;
// Collect all load and store ops in loop nest rooted at 'forOp'.
SmallVector<Operation *, 8> loadAndStoreOpInsts;
auto walkResult = forOp.walk([&](Operation *opInst) -> WalkResult {

View File

@@ -159,4 +159,29 @@ func @max_nested(%m: memref<?x?xf32>, %lb0: index, %lb1: index,
return
}
// CHECK-LABEL: @unsupported_iter_args
func @unsupported_iter_args(%in: memref<10xf32>) {
%cst = constant 0.000000e+00 : f32
// CHECK-NOT: affine.parallel
%final_red = affine.for %i = 0 to 10 iter_args(%red_iter = %cst) -> (f32) {
%ld = affine.load %in[%i] : memref<10xf32>
%add = addf %red_iter, %ld : f32
affine.yield %add : f32
}
return
}
// CHECK-LABEL: @unsupported_nested_iter_args
func @unsupported_nested_iter_args(%in: memref<20x10xf32>) {
%cst = constant 0.000000e+00 : f32
// CHECK: affine.parallel
affine.for %i = 0 to 20 {
// CHECK: affine.for
%final_red = affine.for %j = 0 to 10 iter_args(%red_iter = %cst) -> (f32) {
%ld = affine.load %in[%i, %j] : memref<20x10xf32>
%add = addf %red_iter, %ld : f32
affine.yield %add : f32
}
}
return
}