Fix parfor tests (#180)

Hardcode84 · web-flow · commit 210d8aed3280 · 2022-03-19T18:39:04.000+03:00
diff --git a/dpcomp_runtime/lib/tbb_parallel.cpp b/dpcomp_runtime/lib/tbb_parallel.cpp
@@ -126,20 +126,23 @@ static void run_parallel_for(const InputRange *input_ranges, size_t depth,
   auto runFunc = [&](Dim *current) {
     auto thread_index =
         static_cast<size_t>(tbb::this_task_arena::current_thread_index());
+    assert(thread_index >= 0);
     std::array<Range, 8> static_ranges;
     std::unique_ptr<Range[]> dyn_ranges;
     auto *range_ptr = [&]() -> Range * {
-      if (num_loops <= static_ranges.size()) {
+      if (num_loops <= static_ranges.size())
         return static_ranges.data();
-      }
+
       dyn_ranges.reset(new Range[num_loops]);
       return dyn_ranges.get();
     }();
 
     for (size_t i = 0; i < num_loops; ++i) {
+      assert(current);
       range_ptr[num_loops - i - 1] = current->val;
       current = current->prev;
     }
+
     if (DEBUG) {
       std::lock_guard<std::mutex> lock(getDebugMutext());
       fprintf(stderr, "parallel_for func: thread_index=%d",
@@ -174,7 +177,7 @@ static void run_parallel_for(const InputRange *input_ranges, size_t depth,
       runFunc(prev);
     } else {
       auto next = depth + N;
-      parallel_for_nested(input_ranges, next, num_threads, num_loops, prev_dim,
+      parallel_for_nested(input_ranges, next, num_threads, num_loops, prev,
                           func, ctx);
     }
   };
diff --git a/mlir/lib/dialect/plier_util/dialect.cpp b/mlir/lib/dialect/plier_util/dialect.cpp
@@ -20,6 +20,7 @@
 #include <mlir/IR/BuiltinTypes.h>
 #include <mlir/IR/DialectImplementation.h>
 #include <mlir/IR/Dominance.h>
+#include <mlir/IR/Matchers.h>
 #include <mlir/IR/PatternMatch.h>
 #include <mlir/Transforms/InliningUtils.h>
 
@@ -269,14 +270,91 @@ struct GenGlobalId : public mlir::OpRewritePattern<mlir::arith::AddIOp> {
     return mlir::failure();
   }
 };
+
+struct InvertCmpi : public mlir::OpRewritePattern<mlir::arith::CmpIOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::arith::CmpIOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+
+    if (!mlir::matchPattern(op.getLhs(), mlir::m_Constant()) ||
+        mlir::matchPattern(op.getRhs(), mlir::m_Constant()))
+      return mlir::failure();
+
+    using Pred = mlir::arith::CmpIPredicate;
+    const std::pair<Pred, Pred> inv[] = {
+        // clang-format off
+        {Pred::slt, Pred::sgt},
+        {Pred::sle, Pred::sge},
+        {Pred::ult, Pred::ugt},
+        {Pred::ule, Pred::uge},
+        {Pred::eq, Pred::eq},
+        {Pred::ne, Pred::ne},
+        // clang-format on
+    };
+
+    auto newPred = [&]() -> Pred {
+      auto oldPred = op.getPredicate();
+      for (auto it : inv) {
+        if (it.first == oldPred)
+          return it.second;
+        if (it.second == oldPred)
+          return it.first;
+      }
+
+      llvm_unreachable("Unknown predicate");
+    }();
+
+    rewriter.replaceOpWithNewOp<mlir::arith::CmpIOp>(op, newPred, op.getRhs(),
+                                                     op.getLhs());
+    ;
+    return mlir::success();
+  }
+};
+
+struct ReshapeAlloca : public mlir::OpRewritePattern<mlir::memref::ReshapeOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::memref::ReshapeOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto shapeOp = op.shape().getDefiningOp<mlir::memref::AllocOp>();
+    if (!shapeOp)
+      return mlir::failure();
+
+    for (auto user : shapeOp->getUsers())
+      if (!mlir::isa<mlir::memref::StoreOp, mlir::memref::ReshapeOp>(user))
+        return mlir::failure();
+
+    if (!shapeOp.dynamicSizes().empty() || !shapeOp.symbolOperands().empty())
+      return mlir::failure();
+
+    auto func = op->getParentOfType<mlir::FuncOp>();
+    if (!func)
+      return mlir::failure();
+
+    if (shapeOp->getParentOp() != func) {
+      rewriter.setInsertionPointToStart(&func.getBody().front());
+    } else {
+      rewriter.setInsertionPoint(shapeOp);
+    }
+
+    auto type = shapeOp.getType().cast<mlir::MemRefType>();
+    auto alignment = shapeOp.alignmentAttr().cast<mlir::IntegerAttr>();
+    rewriter.replaceOpWithNewOp<mlir::memref::AllocaOp>(shapeOp, type,
+                                                        alignment);
+    return mlir::success();
+  }
+};
 } // namespace
 
 void PlierUtilDialect::getCanonicalizationPatterns(
     mlir::RewritePatternSet &results) const {
   results.add<DimExpandShape<mlir::tensor::DimOp, mlir::tensor::ExpandShapeOp>,
               DimExpandShape<mlir::memref::DimOp, mlir::memref::ExpandShapeOp>,
-              DimInsertSlice, FillExtractSlice, SpirvInputCSE, GenGlobalId>(
-      getContext());
+              DimInsertSlice, FillExtractSlice, SpirvInputCSE, GenGlobalId,
+              InvertCmpi, ReshapeAlloca>(getContext());
 }
 
 OpaqueType OpaqueType::get(mlir::MLIRContext *context) {
@@ -330,35 +408,35 @@ EnforceShapeOp::fold(llvm::ArrayRef<mlir::Attribute> operands) {
 }
 
 namespace {
-struct EnforceShapeDim : public mlir::OpRewritePattern<mlir::memref::DimOp> {
-  using mlir::OpRewritePattern<mlir::memref::DimOp>::OpRewritePattern;
+template <typename DimOp>
+struct EnforceShapeDim : public mlir::OpRewritePattern<DimOp> {
+  using mlir::OpRewritePattern<DimOp>::OpRewritePattern;
 
   mlir::LogicalResult
-  matchAndRewrite(mlir::memref::DimOp op,
-                  mlir::PatternRewriter &rewriter) const override {
-    auto enforce_op = mlir::dyn_cast_or_null<plier::EnforceShapeOp>(
-        op.source().getDefiningOp());
-    if (!enforce_op) {
+  matchAndRewrite(DimOp op, mlir::PatternRewriter &rewriter) const override {
+    auto enforceOp =
+        op.source().template getDefiningOp<plier::EnforceShapeOp>();
+    if (!enforceOp)
       return mlir::failure();
-    }
-    auto const_ind = plier::getConstVal<mlir::IntegerAttr>(op.index());
-    if (!const_ind) {
+
+    auto constInd = mlir::getConstantIntValue(op.index());
+    if (!constInd)
       return mlir::failure();
-    }
-    auto index = const_ind.getInt();
-    if (index < 0 || index >= static_cast<int64_t>(enforce_op.sizes().size())) {
+
+    auto index = *constInd;
+    if (index < 0 || index >= static_cast<int64_t>(enforceOp.sizes().size()))
       return mlir::failure();
-    }
 
-    rewriter.replaceOp(op, enforce_op.sizes()[static_cast<unsigned>(index)]);
+    rewriter.replaceOp(op, enforceOp.sizes()[static_cast<unsigned>(index)]);
     return mlir::success();
   }
 };
 } // namespace
 
 void EnforceShapeOp::getCanonicalizationPatterns(
     ::mlir::RewritePatternSet &results, ::mlir::MLIRContext *context) {
-  results.insert<EnforceShapeDim>(context);
+  results.insert<EnforceShapeDim<mlir::tensor::DimOp>,
+                 EnforceShapeDim<mlir::memref::DimOp>>(context);
 }
 
 mlir::LogicalResult
diff --git a/numba_dpcomp/numba_dpcomp/mlir/tests/test_numba_parfor.py b/numba_dpcomp/numba_dpcomp/mlir/tests/test_numba_parfor.py
@@ -52,6 +52,7 @@ def _gen_tests():
     ]
 
     xfail_tests = {
+        "test_prange26",
         "test_prange03mul",
         "test_prange09",
         "test_prange03sub",
@@ -182,12 +183,7 @@ def _gen_tests():
         "test_one_d_array_reduction",
     }
 
-    skip_tests = {
-        "test_copy_global_for_parfor",
-        "test_high_dimension1",
-        "test_three_d_array_reduction",
-        "test_prange26",
-    }
+    skip_tests = {}
 
     def countParfors(test_func, args, **kws):
         pytest.xfail()
diff --git a/numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/plier_to_linalg.cpp b/numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/plier_to_linalg.cpp
@@ -2851,8 +2851,6 @@ static void populatePlierToLinalgOptPipeline(mlir::OpPassManager &pm) {
       mlir::bufferization::createBufferHoistingPass());
   pm.addNestedPass<mlir::FuncOp>(
       mlir::bufferization::createBufferLoopHoistingPass());
-  pm.addNestedPass<mlir::FuncOp>(
-      mlir::bufferization::createPromoteBuffersToStackPass());
 
   pm.addNestedPass<mlir::FuncOp>(std::make_unique<CloneArgsPass>());
   pm.addPass(std::make_unique<MakeStridedLayoutPass>());
@@ -2863,6 +2861,8 @@ static void populatePlierToLinalgOptPipeline(mlir::OpPassManager &pm) {
   pm.addPass(mlir::createCanonicalizerPass());
 
   pm.addNestedPass<mlir::FuncOp>(std::make_unique<LowerCloneOpsPass>());
+  pm.addNestedPass<mlir::FuncOp>(
+      mlir::bufferization::createPromoteBuffersToStackPass());
 
   pm.addPass(std::make_unique<LowerLinalgPass>());
   pm.addPass(plier::createForceInlinePass());
diff --git a/numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/plier_to_scf.cpp b/numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/plier_to_scf.cpp
@@ -432,17 +432,17 @@ struct ScfWhileRewrite : public mlir::OpRewritePattern<mlir::cf::BranchOp> {
         }
       }
 
-      llvm::transform(
-          beforeBlock->getArguments(), yieldVars.begin(),
-          [&](mlir::Value val) { return mapper.lookupOrDefault(val); });
+      llvm::transform(beforeBlockArgs, yieldVars.begin(), [&](mlir::Value val) {
+        return mapper.lookupOrDefault(val);
+      });
 
-      auto term =
-          mlir::cast<mlir::cf::CondBranchOp>(beforeBlock->getTerminator());
-      for (auto arg : term.getFalseDestOperands()) {
+      auto falseArgs = beforeTerm.getFalseDestOperands();
+      for (auto arg : falseArgs) {
         origVars.emplace_back(arg);
         yieldVars.emplace_back(mapper.lookupOrDefault(arg));
       }
-      auto cond = mapper.lookupOrDefault(term.getCondition());
+
+      auto cond = mapper.lookupOrDefault(beforeTerm.getCondition());
       builder.create<mlir::scf::ConditionOp>(loc, cond, yieldVars);
     };
     auto afterBody = [&](mlir::OpBuilder &builder, mlir::Location loc,
@@ -465,11 +465,22 @@ struct ScfWhileRewrite : public mlir::OpRewritePattern<mlir::cf::BranchOp> {
                                beforeBody, afterBody);
 
     assert(origVars.size() == whileOp.getNumResults());
-    for (auto arg : llvm::zip(origVars, whileOp.getResults()))
-      std::get<0>(arg).replaceAllUsesWith(std::get<1>(arg));
+    for (auto arg : llvm::zip(origVars, whileOp.getResults())) {
+      auto origVal = std::get<0>(arg);
+      for (auto &use : llvm::make_early_inc_range(origVal.getUses())) {
+        auto *owner = use.getOwner();
+        auto *block = owner->getBlock();
+        if (block != &whileOp.getBefore().front() &&
+            block != &whileOp.getAfter().front()) {
+          auto newVal = std::get<1>(arg);
+          rewriter.updateRootInPlace(owner, [&]() { use.set(newVal); });
+        }
+      }
+    }
 
-    rewriter.replaceOpWithNewOp<mlir::cf::BranchOp>(
-        op, postBlock, beforeTerm.getFalseDestOperands());
+    auto results =
+        whileOp.getResults().take_back(beforeTerm.getNumFalseOperands());
+    rewriter.replaceOpWithNewOp<mlir::cf::BranchOp>(op, postBlock, results);
     return mlir::success();
   }
 };
