added parallel back

Signed-off-by: Alexandre Eichenberger <[email protected]>

Author: AlexandreEichenberger

src/Accelerators/NNPA/Compiler/NNPACompilerUtils.cpp

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
+#include "mlir/Conversion/Passes.h"
 #include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
@@ -215,9 +216,10 @@ void addPassesNNPA(mlir::OwningOpRef<mlir::ModuleOp> &module,
         addKrnlToAffinePasses(pm);
         // Optimizations at ZLow that needs affine map in MemRef.
         pm.addPass(zlow::createZLowRewritePass());
-        if (nnpaEnableCompilerStickUnstick) {
-          pm.addPass(zlow::createZLowStickExpansionPass());
-        }
+        // Late generation of code for stick/unstick, needed to be after a
+        // ZLowRewrite pass.
+        if (nnpaEnableCompilerStickUnstick)
+          pm.addPass(zlow::createZLowStickExpansionPass(enableParallel));
         pm.addPass(mlir::createCanonicalizerPass());
         // Normalize MemRefs.
         normalizeMemRefsPasses(pm);
@@ -226,6 +228,11 @@ void addPassesNNPA(mlir::OwningOpRef<mlir::ModuleOp> &module,
         addKrnlToAffinePasses(pm);
         // Optimizations at ZLow after normalizing MemRefs.
         pm.addPass(zlow::createZLowRewritePass());
+        // The createZLowStickExpansion pass may create parallel constructs,
+        // they need to be handled here.
+        if (nnpaEnableCompilerStickUnstick && enableParallel)
+          pm.addPass(mlir::createConvertSCFToOpenMPPass());
+
         pm.addPass(mlir::createCanonicalizerPass());
         // Constant folding for std.alloc.
         pm.addNestedPass<func::FuncOp>(onnx_mlir::createFoldStdAllocPass());

src/Accelerators/NNPA/Pass/NNPAPasses.hpp

@@ -61,7 +61,8 @@ namespace zlow {
 std::unique_ptr<mlir::Pass> createZLowRewritePass();
 
 /// Add pass for rewriting ZLow ops.
-std::unique_ptr<mlir::Pass> createZLowStickExpansionPass();
+std::unique_ptr<mlir::Pass> createZLowStickExpansionPass(
+    bool enableParallel = false);
 
 /// Add pass for rewriting ZLow ops.
 std::unique_ptr<mlir::Pass> createZLowDummyOpForMultiDerefPass();

src/Accelerators/NNPA/Transform/ZLow/ZLowStickExpansion.cpp

@@ -451,7 +451,11 @@ class ZLowStickExpansionPass
     : public PassWrapper<ZLowStickExpansionPass, OperationPass<func::FuncOp>> {
 
 public:
-  bool enableParallelism = false; // hi alex, fix this.
+  ZLowStickExpansionPass(bool enableParallel)
+      : PassWrapper<ZLowStickExpansionPass, OperationPass<func::FuncOp>>(),
+        enableParallel(enableParallel) {}
+
+  bool enableParallel;
 
   StringRef getArgument() const override { return "zlow-stick-expansion"; }
 
@@ -465,8 +469,8 @@ class ZLowStickExpansionPass
     llvm::SmallDenseSet<ZLowStickOp, 4> removableStickOps;
     ConversionTarget target(getContext());
     RewritePatternSet patterns(&getContext());
-    patterns.insert<StickExpansionPattern>(&getContext(), enableParallelism);
-    patterns.insert<UnstickExpansionPattern>(&getContext(), enableParallelism);
+    patterns.insert<StickExpansionPattern>(&getContext(), enableParallel);
+    patterns.insert<UnstickExpansionPattern>(&getContext(), enableParallel);
     // patterns.insert<UnstickExpansionPattern>(&getContext());
 
     fprintf(stderr, "hi alex, apply patterns in zlow stick expansion\n");
@@ -484,8 +488,8 @@ class ZLowStickExpansionPass
   }
 };
 
-std::unique_ptr<Pass> createZLowStickExpansionPass() {
-  return std::make_unique<ZLowStickExpansionPass>();
+std::unique_ptr<Pass> createZLowStickExpansionPass(bool enableParallel) {
+  return std::make_unique<ZLowStickExpansionPass>(enableParallel);
 }
 
 } // namespace zlow