[Mlir-commits] [mlir] ebca222 - [mlir] Check 'iter_args' in 'isLoopParallel' utility

[Diego Caballero]

Author: Diego Caballero
Date: 2021-02-25T18:12:34+02:00
New Revision: ebca222b65cb847f7bf4ee3da1dd7e2df35d0338

URL: https://github.com/llvm/llvm-project/commit/ebca222b65cb847f7bf4ee3da1dd7e2df35d0338
DIFF: https://github.com/llvm/llvm-project/commit/ebca222b65cb847f7bf4ee3da1dd7e2df35d0338.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Analysis/Utils.cpp
    mlir/test/Dialect/Affine/parallelize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp
index 383a6587bbef..c3c63f3a7f91 100644
--- a/mlir/lib/Analysis/Utils.cpp
+++ b/mlir/lib/Analysis/Utils.cpp
@@ -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 {

diff  --git a/mlir/test/Dialect/Affine/parallelize.mlir b/mlir/test/Dialect/Affine/parallelize.mlir
index 3f7a79daf91d..5ca2f4652827 100644
--- a/mlir/test/Dialect/Affine/parallelize.mlir
+++ b/mlir/test/Dialect/Affine/parallelize.mlir
@@ -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
+}