[MLIR] Fix affine parallelize pass.

[pashu123]

Summary

To control the number of outer parallel loops, we need to process the
outer loops first and hence pre-order walk fixes the issue.

[pashu123]

Added reviewer.

[bondhugula]

The commit summary is outdated/does not match the fix. Furthermore, a deque is not needed here any more.

mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp

52–53

Why is deque being used here? An std::vector would do.

[bondhugula]

mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp

53–56

Thanks for fixing this - looks good other than the fact that a deque isn't needed here.

[pashu123]

Updating the commit message and minor changes.

mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp

struct ParallelizationCandidate {
   AffineForOp loop;
   /// Desciprtors of reductions that can be parallelized in the loop.
   SmallVector<LoopReduction> reductions;
 };
 } // namespace
 
 void AffineParallelize::runOnFunction() {
   FuncOp f = getFunction();
 
-  // The walker proceeds in post-order, but we need to process outer loops first
-  // to control the number of outer parallel loops, so push candidate loops to
-  // the front of a deque.
-  std::deque<ParallelizationCandidate> parallelizableLoops;
-  f.walk([&](AffineForOp loop) {
+  // The walker proceeds in pre-order to process the outer loops first
+  // and control the number of outer parallel loops.
+  std::vector<ParallelizationCandidate> parallelizableLoops;
+  f.walk<WalkOrder::PreOrder>([&](AffineForOp loop) {
     SmallVector<LoopReduction> reductions;
     if (isLoopParallel(loop, parallelReductions ? &reductions : nullptr))
-      parallelizableLoops.emplace_back(loop, std::move(reductions));
+      parallelizableLoops.push_back({loop, std::move(reductions)});
   });
 
   for (const ParallelizationCandidate &candidate : parallelizableLoops) {
     unsigned numParentParallelOps = 0;
     AffineForOp loop = candidate.loop;
     for (Operation *op = loop->getParentOp();
          op != nullptr && !op->hasTrait<OpTrait::AffineScope>();
          op = op->getParentOp()) {

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
 }
 
+// MAX-NESTED-LABEL: @max_nested_1
+func @max_nested_1(%arg0: memref<4096x4096xf32>, %arg1: memref<4096x4096xf32>, %arg2: memref<4096x4096xf32>) {
+  %0 = memref.alloc() : memref<4096x4096xf32>
+  // MAX-NESTED: affine.parallel
+  affine.for %arg3 = 0 to 4096 {
+    // MAX-NESTED-NEXT: affine.for
+    affine.for %arg4 = 0 to 4096 {
+      // MAX-NESTED-NEXT: affine.for
+      affine.for %arg5 = 0 to 4096 {
+        %1 = affine.load %arg0[%arg3, %arg5] : memref<4096x4096xf32>
+        %2 = affine.load %arg1[%arg5, %arg4] : memref<4096x4096xf32>
+        %3 = affine.load %0[%arg3, %arg4] : memref<4096x4096xf32>
+        %4 = mulf %1, %2 : f32
+        %5 = addf %3, %4 : f32
+        affine.store %5, %0[%arg3, %arg4] : memref<4096x4096xf32>
+      }
+    }
+  }
+  return
+}
+
 // CHECK-LABEL: @iter_args
 // REDUCE-LABEL: @iter_args
 func @iter_args(%in: memref<10xf32>) {
   // REDUCE: %[[init:.*]] = constant
   %cst = constant 0.000000e+00 : f32
   // CHECK-NOT: affine.parallel
   // REDUCE: %[[reduced:.*]] = affine.parallel (%{{.*}}) = (0) to (10) reduce ("addf")
   %final_red = affine.for %i = 0 to 10 iter_args(%red_iter = %cst) -> (f32) {