[Affine] Create memory banks for memories used inside affine parallel loops

mlir/include/mlir/Dialect/Affine/Passes.h

@@ -25,6 +25,7 @@ class FuncOp;
 
 namespace affine {
 class AffineForOp;
+class AffineParallelOp;
 
 /// Fusion mode to attempt. The default mode `Greedy` does both
 /// producer-consumer and sibling fusion.
@@ -108,6 +109,13 @@ std::unique_ptr<OperationPass<func::FuncOp>> createLoopUnrollPass(
 std::unique_ptr<OperationPass<func::FuncOp>>
 createLoopUnrollAndJamPass(int unrollJamFactor = -1);
 
+/// Creates a memory banking pass to explicitly partition the memories used
+/// inside affine parallel operations
+std::unique_ptr<OperationPass<func::FuncOp>> createParallelBankingPass(
+    int unrollFactor = -1,
+    const std::function<unsigned(AffineParallelOp)> &getBankingFactor =
+        nullptr);
+
 /// Creates a pass to pipeline explicit movement of data across levels of the
 /// memory hierarchy.
 std::unique_ptr<OperationPass<func::FuncOp>> createPipelineDataTransferPass();

mlir/include/mlir/Dialect/Affine/Passes.td

@@ -381,6 +381,15 @@ def AffineParallelize : Pass<"affine-parallelize", "func::FuncOp"> {
   ];
 }
 
+def AffineParallelBanking : Pass<"affine-parallel-banking", "func::FuncOp"> {
+  let summary = "Partition the memories used in affine parallel loops into banks";
+  let constructor = "mlir::affine::createParallelBankingPass()";
+  let options = [
+    Option<"bankingFactor", "banking-factor", "unsigned", /*default=*/"1",
+           "Use this banking factor for all memories being partitioned">
+  ];
+}
+
 def AffineLoopNormalize : Pass<"affine-loop-normalize", "func::FuncOp"> {
   let summary = "Apply normalization transformations to affine loop-like ops";
   let constructor = "mlir::affine::createAffineLoopNormalizePass()";

mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt

@@ -11,6 +11,7 @@ add_mlir_dialect_library(MLIRAffineTransforms
   LoopTiling.cpp
   LoopUnroll.cpp
   LoopUnrollAndJam.cpp
+  ParallelBanking.cpp
   PipelineDataTransfer.cpp
   ReifyValueBounds.cpp
   SuperVectorize.cpp
@@ -39,5 +40,6 @@ add_mlir_dialect_library(MLIRAffineTransforms
   MLIRValueBoundsOpInterface
   MLIRVectorDialect
   MLIRVectorUtils
+  MLIRSCFDialect
   )