[mlir] Add an option to control the number of loops in affine parallelizer

[ftynse]

Summary

Add a pass option to control the number of nested parallel loops produced by
the parallelization passes. This is useful to build end-to-end passes targeting
systems that don't need multiple parallel dimensions (e.g., CPUs typically need
only one).

mlir/include/mlir/Dialect/Affine/Passes.td

ListOption<"fastestVaryingPattern", "test-fastest-varying", "int64_t",
                "a description and examples. This is used for testing purposes",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
   ];
 }
 
 def AffineParallelize : FunctionPass<"affine-parallelize"> {
   let summary = "Convert affine.for ops into 1-D affine.parallel";
   let constructor = "mlir::createAffineParallelizePass()";
+  let options = [
+    Option<"maxNested", "max-nested", "unsigned", /*default=*/"-1u",
+           "Maximum number of nested parallel loops to produce. "
+           "Defaults to unlimited (UINT_MAX).">,
+  ];
 }
 
 def AffineLoopNormalize : FunctionPass<"affine-loop-normalize"> {
   let summary = "Apply normalization transformations to affine loop-like ops";
   let constructor = "mlir::createAffineLoopNormalizePass()";
 }
 
 def SimplifyAffineStructures : FunctionPass<"simplify-affine-structures"> {
   let summary = "Simplify affine expressions in maps/sets and normalize "
                 "memrefs";
   let constructor = "mlir::createSimplifyAffineStructuresPass()";
 }
 
 #endif // MLIR_DIALECT_AFFINE_PASSES

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

 /// Convert all parallel affine.for op into 1-D affine.parallel op.
 struct AffineParallelize : public AffineParallelizeBase<AffineParallelize> {
   void runOnFunction() override;
 };
 } // namespace
 
 void AffineParallelize::runOnFunction() {
   FuncOp f = getFunction();
-  SmallVector<AffineForOp, 8> parallelizableLoops;
+
+  // 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<AffineForOp> parallelizableLoops;
   f.walk([&](AffineForOp loop) {
     if (isLoopParallel(loop))
-      parallelizableLoops.push_back(loop);
+      parallelizableLoops.push_front(loop);
   });
-  for (AffineForOp loop : parallelizableLoops)
+
+  for (AffineForOp loop : parallelizableLoops) {
+    unsigned numParentParallelOps = 0;
+    for (Operation *op = loop->getParentOp();
+         op != nullptr && !op->hasTrait<OpTrait::AffineScope>();
+         op = op->getParentOp()) {
+      if (isa<AffineParallelOp>(op))
+        ++numParentParallelOps;
+    }
+
+    if (numParentParallelOps < maxNested)
       affineParallelize(loop);
   }
+}
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createAffineParallelizePass() {
   return std::make_unique<AffineParallelize>();
 }

mlir/test/Dialect/Affine/parallelize.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -affine-parallelize| FileCheck %s
+// RUN: mlir-opt %s -allow-unregistered-dialect -affine-parallelize='max-nested=1' | FileCheck --check-prefix=MAX-NESTED %s
 
 // CHECK-LABEL:    func @reduce_window_max() {
 func @reduce_window_max() {
   %cst = constant 0.000000e+00 : f32
   %0 = alloc() : memref<1x8x8x64xf32>
   %1 = alloc() : memref<1x18x18x64xf32>
   affine.for %arg0 = 0 to 1 {
     affine.for %arg1 = 0 to 8 {
affine.for %j = 0 to 10 {
     // CHECK: affine.for
     affine.for %i = max affine_map<(d0, d1) -> (d0, d1)>(%lb0, %j)
             to min affine_map<(d0, d1) -> (d0, d1)>(%ub0, %ub1) {
       affine.load %m[%i] : memref<?xf32>
     }
   }
   return
 }
+
+// MAX-NESTED-LABEL: @max_nested
+func @max_nested(%m: memref<?x?xf32>, %lb0: index, %lb1: index,
+                 %ub0: index, %ub1: index) {
+  // MAX-NESTED: affine.parallel
+  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
+}
+
+