Author: Alex Zinenko Date: 2020-12-08T10:44:37+01:00 New Revision: 80766ecc65096deeb4ff6f03562dcad94c54b862
URL: https://github.com/llvm/llvm-project/commit/80766ecc65096deeb4ff6f03562dcad94c54b862 DIFF: https://github.com/llvm/llvm-project/commit/80766ecc65096deeb4ff6f03562dcad94c54b862.diff LOG: [mlir] Add an option to control the number of loops in affine parallelizer 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). Reviewed By: wsmoses, chelini Differential Revision: https://reviews.llvm.org/D92765 Added: Modified: mlir/include/mlir/Dialect/Affine/Passes.td mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp mlir/test/Dialect/Affine/parallelize.mlir Removed: ################################################################################ diff --git a/mlir/include/mlir/Dialect/Affine/Passes.td b/mlir/include/mlir/Dialect/Affine/Passes.td index ace272692d66..9f2aac7b685a 100644 --- a/mlir/include/mlir/Dialect/Affine/Passes.td +++ b/mlir/include/mlir/Dialect/Affine/Passes.td @@ -118,6 +118,11 @@ def AffineVectorize : FunctionPass<"affine-super-vectorize"> { 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"> { diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp index b3651e202245..d8511be326dd 100644 --- a/mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp @@ -36,13 +36,28 @@ struct AffineParallelize : public AffineParallelizeBase<AffineParallelize> { 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) - affineParallelize(loop); + + 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() { diff --git a/mlir/test/Dialect/Affine/parallelize.mlir b/mlir/test/Dialect/Affine/parallelize.mlir index cbc80a092e76..08aaa7f32fc8 100644 --- a/mlir/test/Dialect/Affine/parallelize.mlir +++ b/mlir/test/Dialect/Affine/parallelize.mlir @@ -1,4 +1,5 @@ // 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() { @@ -144,3 +145,18 @@ func @nested_for_with_minmax(%m: memref<?xf32>, %lb0: index, } 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 +} + + _______________________________________________ llvm-branch-commits mailing list llvm-branch-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
