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

[Scott Manley]

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

>From f5c04c883e405ba70f592afe43ab5fa7a11d814c 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          |  5 +++++
 .../Affine/Transforms/AffineParallelize.cpp        | 14 ++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Affine/Passes.h b/mlir/include/mlir/Dialect/Affine/Passes.h
index e152101236dc7a..4193382e53db17 100644
--- a/mlir/include/mlir/Dialect/Affine/Passes.h
+++ b/mlir/include/mlir/Dialect/Affine/Passes.h
@@ -48,6 +48,11 @@ 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 64f2bc6e745c37..319796f2b1d5f1 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,9 @@ 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);
+}