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

[dcaballe]

Summary

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

[tungld]

Hi Diego,

Does this patch handle iter_args in an inner loop? For example:

%cst = constant 0.000000e+00 : f32
affine.for %arg1 = 0 to 10 {
   %0 = affine.for %arg2 = 0 to 10 iter_args(%arg3 = %cst) -> (f32) {
       %1 = affine.load %arg0[%arg1, %arg2] : memref<10x10xf32>
       %2 = addf %arg3, %1 : f32
       affine.yield %2 : f32
    }
}

Thanks,
Tung.

mlir/test/Dialect/Affine/parallelize.mlir

177

How about using consecutive numbering for arg, and starting from %0 instead of %5?

[dcaballe]

Does this patch handle iter_args in an inner loop? For example:

Good point! I'll add a test for that but, yes, AffineParallelize pass is processing each loop independently so a parallel loop will be generated for the outer loop but not for the inner loop.

Thanks!

mlir/test/Dialect/Affine/parallelize.mlir

177

Thanks! I'll use more relevant names :)

[dcaballe]

Addressing feedback.

[vinayaka-polymage]

Looks good, thanks !

[tungld]

LGTM! Thanks!

mlir/lib/Analysis/Utils.cpp

forOp->walk([&](Operation *op) {
     if (auto innerFor = dyn_cast<AffineForOp>(op))
       if (!isLoopParallel(innerFor))
         sequentialLoops->insert(innerFor.getInductionVar());
   });
 }
 
 /// 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 {
     if (isa<AffineReadOpInterface, AffineWriteOpInterface>(opInst))
       loadAndStoreOpInsts.push_back(opInst);
     else if (!isa<AffineForOp, AffineYieldOp, AffineIfOp>(opInst) &&
              !MemoryEffectOpInterface::hasNoEffect(opInst))
       return WalkResult::interrupt();

mlir/test/Dialect/Affine/parallelize.mlir

affine.for %i = affine_map<(d0) -> (d0)>(%lb0) to affine_map<(d0) -> (d0)>(%ub0) {
     // MAX-NESTED: affine.for
     affine.for %j = affine_map<(d0) -> (d0)>(%lb1) to affine_map<(d0) -> (d0)>(%ub1) {
       affine.load %m[%i, %j] : memref<?x?xf32>
     }
   }
   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
+}