[AffineParallelize] expose options when creating pass

[rscottmanley]

Add a createAffineParallelizePass() that takes AffineParallelizeOptions
so it can be customized in a pass pipeline.

[llvmbot]

@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: Scott Manley (rscottmanley)

Changes

Add a createAffineParallelizePass() that takes AffineParallelizeOptions
so it can be customized in a pass pipeline.

---

Full diff: https://github.com/llvm/llvm-project/pull/124959.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Affine/Passes.h (+4)
• (modified) mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp (+13)

diff --git a/mlir/include/mlir/Dialect/Affine/Passes.h b/mlir/include/mlir/Dialect/Affine/Passes.h
index e152101236dc7a9..2d2b5f461e813e3 100644
--- a/mlir/include/mlir/Dialect/Affine/Passes.h
+++ b/mlir/include/mlir/Dialect/Affine/Passes.h
@@ -48,6 +48,10 @@ createAffineLoopInvariantCodeMotionPass();
 /// ops.
 std::unique_ptr<OperationPass<func::FuncOp>> createAffineParallelizePass();
 
+/// Creates a pass to convert all parallel affine.for's into 1-d affine.parallel
+/// ops using the specified parallelize options.
+std::unique_ptr<OperationPass<func::FuncOp>> createAffineParallelizePass(const AffineParallelizeOptions &options);
+
 /// Apply normalization transformations to affine loop-like ops. If
 /// `promoteSingleIter` is true, single iteration loops are promoted (i.e., the
 /// loop is replaced by its loop body).
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp
index 64f2bc6e745c374..f841b83c5aecde0 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp
@@ -42,6 +42,14 @@ namespace {
 /// Convert all parallel affine.for op into 1-D affine.parallel op.
 struct AffineParallelize
     : public affine::impl::AffineParallelizeBase<AffineParallelize> {
+
+  AffineParallelize() = default;
+
+  explicit AffineParallelize(const AffineParallelizeOptions &options) {
+    maxNested = std::move(options.maxNested);
+    parallelReductions = std::move(options.parallelReductions);
+  }
+
   void runOnOperation() override;
 };
 
@@ -95,3 +103,8 @@ std::unique_ptr<OperationPass<func::FuncOp>>
 mlir::affine::createAffineParallelizePass() {
   return std::make_unique<AffineParallelize>();
 }
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+mlir::affine::createAffineParallelizePass(const AffineParallelizeOptions &options) {
+  return std::make_unique<AffineParallelize>(options);
+}

[dcaballe]

These constructors should have been auto-generated in AffineParallelizeBase. You probably just need to make it visible here with using AffineParallelizeBase<AffineParallelize>::AffineParallelizeBase...

[rscottmanley]

Agree -- will change. Thanks for the review Diego!

[dcaballe]

There is also a way to have this auto-generated. I think if you go to Affine/Passes.td and remove the line let constructor = "mlir::affine::createAffineParallelizePass()";, tablegen should generate the respective createAffineParallelizePass constructors.

[rscottmanley]

Ironically had done it this way before seeing how OneShotBufferize did it and then did it that way instead! I definitely prefer this way so I'll make that change.

[dcaballe]

Cool, thanks for modernizing this code!