[Mlir-commits] [mlir] [AffineParallelize] expose options when creating pass (PR #124959)

[Scott Manley]

https://github.com/rscottmanley created https://github.com/llvm/llvm-project/pull/124959

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

>From afface61313bd668bbf470554b75cb11d24caed5 Mon Sep 17 00:00:00 2001
From: Scott Manley <rscottmanley at gmail.com>
Date: Wed, 29 Jan 2025 09:34:53 -0800
Subject: [PATCH] [AffineParallelize] expose options when creating pass

Add a createAffineParallelizePass that takes the
AffineParallelizeOptions so it can be customized in a pass pipeline.
---
 mlir/include/mlir/Dialect/Affine/Passes.h           |  4 ++++
 .../Dialect/Affine/Transforms/AffineParallelize.cpp | 13 +++++++++++++
 2 files changed, 17 insertions(+)

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);
+}