From ebca222b65cb847f7bf4ee3da1dd7e2df35d0338 Mon Sep 17 00:00:00 2001
From: Diego Caballero <diego.caballero@intel.com>
Date: Thu, 25 Feb 2021 17:56:51 +0200
Subject: [PATCH] [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
---
 mlir/lib/Analysis/Utils.cpp               |  6 ++++++
 mlir/test/Dialect/Affine/parallelize.mlir | 25 +++++++++++++++++++++++
 2 files changed, 31 insertions(+)

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
+}