[MLIR] Pass options to ParallelLoopCollapsing constructor.

[rriddle]

mlir/include/mlir/Transforms/Passes.h

18

This looks wrong?

mlir/lib/Transforms/ParallelLoopCollapsing.cpp

22

This looks unrelated.

26

This constructor should be unnecessary.

[mehdi_amini]

Please elaborate a complete description for the revision.

mlir/include/mlir/Transforms/Passes.h

63

Comment needs an update

[pifon2a]

It looks like the pass has to be reworked a bit, so I am closing this one.

mlir/include/mlir/Transforms/Passes.h

 // transformation library.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_PASSES_H
 #define MLIR_TRANSFORMS_PASSES_H
 
 #include "mlir/Support/LLVM.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/ArrayRef.h"
+
 #include <functional>
 #include <limits>
 
 namespace mlir {
 
 class AffineForOp;
 class FuncOp;
 class ModuleOp;

 /// primitives).
 std::unique_ptr<OperationPass<FuncOp>> createLowerAffinePass();
 
 /// Creates a pass that transforms perfectly nested loops with independent
 /// bounds into a single loop.
 std::unique_ptr<OperationPass<FuncOp>> createLoopCoalescingPass();
 
 /// Creates a pass that transforms a single ParallelLoop over N induction
 /// variables into another ParallelLoop over less than N induction variables.
-std::unique_ptr<Pass> createParallelLoopCollapsingPass();
+std::unique_ptr<Pass>
+createParallelLoopCollapsingPass(ArrayRef<unsigned> collapsedIndices0 = {},
+                                 ArrayRef<unsigned> collapsedIndices1 = {},
+                                 ArrayRef<unsigned> collapsedIndices2 = {});
 
 /// Creates a pass to perform optimizations relying on memref dataflow such as
 /// store to load forwarding, elimination of dead stores, and dead allocs.
 std::unique_ptr<OperationPass<FuncOp>> createMemRefDataFlowOptPass();
 
 /// Creates a pass to strip debug information from a function.
 std::unique_ptr<Pass> createStripDebugInfoPass();
 

mlir/lib/Transforms/ParallelLoopCollapsing.cpp

 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "parallel-loop-collapsing"
 
 using namespace mlir;
 
 namespace {
+
 struct ParallelLoopCollapsing
     : public ParallelLoopCollapsingBase<ParallelLoopCollapsing> {
+  ParallelLoopCollapsing() = default;
+  ParallelLoopCollapsing(const ParallelLoopCollapsing &){};
+  ParallelLoopCollapsing(ArrayRef<unsigned> collapsedIndices0,
+                         ArrayRef<unsigned> collapsedIndices1,
+                         ArrayRef<unsigned> collapsedIndices2) {
+    clCollapsedIndices0 = collapsedIndices0;
+    clCollapsedIndices1 = collapsedIndices1;
+    clCollapsedIndices2 = collapsedIndices2;
+  }
+
   void runOnOperation() override {
     Operation *module = getOperation();
 
     module->walk([&](loop::ParallelOp op) {
       // The common case for GPU dialect will be simplifying the ParallelOp to 3
       // arguments, so we do that here to simplify things.
       llvm::SmallVector<std::vector<unsigned>, 3> combinedLoops;
-      if (clCollapsedIndices0.size())
+      if (!clCollapsedIndices0.empty())
         combinedLoops.push_back(clCollapsedIndices0);
-      if (clCollapsedIndices1.size())
+      if (!clCollapsedIndices1.empty())
         combinedLoops.push_back(clCollapsedIndices1);
-      if (clCollapsedIndices2.size())
+      if (!clCollapsedIndices2.empty())
         combinedLoops.push_back(clCollapsedIndices2);
       collapseParallelLoops(op, combinedLoops);
     });
   }
 };
+
 } // namespace
 
-std::unique_ptr<Pass> mlir::createParallelLoopCollapsingPass() {
-  return std::make_unique<ParallelLoopCollapsing>();
+std::unique_ptr<Pass>
+mlir::createParallelLoopCollapsingPass(ArrayRef<unsigned> collapsedIndices0,
+                                       ArrayRef<unsigned> collapsedIndices1,
+                                       ArrayRef<unsigned> collapsedIndices2) {
+  return std::make_unique<ParallelLoopCollapsing>(
+      collapsedIndices0, collapsedIndices1, collapsedIndices2);
 }