Fix some parfor tests (#181)

Author: Hardcode84

azure-pipelines.yml

@@ -41,7 +41,7 @@ jobs:
   - script: |
       call "C:\Miniconda\Scripts\activate"
       call cd numba_dpcomp
-      pytest -n1 -vv --capture=tee-sys
+      pytest -n1 -vv --capture=tee-sys -rX
     displayName: 'Tests'
 
   - script: |
@@ -127,7 +127,7 @@ jobs:
       export SYCL_DEVICE_FILTER=opencl:cpu
       export NUMBA_DISABLE_PERFORMANCE_WARNINGS=1
       export DPCOMP_ENABLE_PARFOR_TESTS=1
-      pytest -n1 -vv --capture=tee-sys
+      pytest -n1 -vv --capture=tee-sys -rX
     displayName: 'Tests'
 
   - script: |
@@ -197,7 +197,7 @@ jobs:
       source /usr/local/miniconda/bin/activate
       cd numba_dpcomp
       conda activate test_env
-      pytest -n1 -vv --capture=tee-sys
+      pytest -n1 -vv --capture=tee-sys -rX
     displayName: 'Tests'
 
   - script: |

mlir/lib/dialect/plier_util/dialect.cpp

@@ -1416,6 +1416,95 @@ struct SignCastMemrefSubviewPropagate
   }
 };
 
+struct SignCastForPropagate : public mlir::OpRewritePattern<mlir::scf::ForOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::scf::ForOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto &body = op.getLoopBody().front();
+    auto term = mlir::cast<mlir::scf::YieldOp>(body.getTerminator());
+    auto termResults = term.getResults();
+    auto initArgs = op.getInitArgs();
+    auto count = static_cast<unsigned>(initArgs.size());
+
+    assert(termResults.size() == count);
+    llvm::SmallVector<mlir::Value> newInitArgs(count);
+    bool needUpdate = false;
+    for (auto i : llvm::seq(0u, count)) {
+      auto initArg = initArgs[i];
+      auto yieldArg = termResults[i];
+      assert(initArg.getType() == yieldArg.getType());
+      auto initCast = initArg.getDefiningOp<plier::SignCastOp>();
+      auto yieldCast = yieldArg.getDefiningOp<plier::SignCastOp>();
+      if (initCast && yieldCast &&
+          initCast.value().getType() == yieldCast.value().getType()) {
+        newInitArgs[i] = initCast.value();
+        needUpdate = true;
+      } else {
+        newInitArgs[i] = initArg;
+      }
+    }
+
+    if (!needUpdate)
+      return mlir::failure();
+
+    auto bodyBuilder = [&](mlir::OpBuilder &builder, mlir::Location loc,
+                           mlir::Value iter, mlir::ValueRange iterVals) {
+      assert(iterVals.size() == count);
+      mlir::BlockAndValueMapping mapping;
+      mapping.map(body.getArguments()[0], iter);
+      auto oldIterVals = body.getArguments().drop_front(1);
+      for (auto i : llvm::seq(0u, count)) {
+        auto iterVal = iterVals[i];
+        auto oldIterVal = oldIterVals[i];
+        auto oldType = oldIterVal.getType();
+        if (iterVal.getType() != oldType) {
+          auto newIterVal =
+              builder.create<plier::SignCastOp>(loc, oldType, iterVal);
+          mapping.map(oldIterVal, newIterVal.getResult());
+        } else {
+          mapping.map(oldIterVal, iterVal);
+        }
+      }
+
+      for (auto &bodyOp : body.without_terminator())
+        builder.clone(bodyOp, mapping);
+
+      llvm::SmallVector<mlir::Value> newYieldArgs(count);
+      for (auto i : llvm::seq(0u, count)) {
+        auto val = mapping.lookupOrDefault(termResults[i]);
+        auto newType = newInitArgs[i].getType();
+        if (val.getType() != newType)
+          val = builder.create<plier::SignCastOp>(loc, newType, val);
+
+        newYieldArgs[i] = val;
+      }
+      builder.create<mlir::scf::YieldOp>(loc, newYieldArgs);
+    };
+
+    auto loc = op->getLoc();
+    auto newResults = rewriter
+                          .create<mlir::scf::ForOp>(
+                              loc, op.getLowerBound(), op.getUpperBound(),
+                              op.getStep(), newInitArgs, bodyBuilder)
+                          ->getResults();
+
+    for (auto i : llvm::seq(0u, count)) {
+      auto oldRersultType = initArgs[i].getType();
+      mlir::Value newResult = newResults[i];
+      if (newResult.getType() != oldRersultType)
+        newResult =
+            rewriter.create<plier::SignCastOp>(loc, oldRersultType, newResult);
+
+      newInitArgs[i] = newResult;
+    }
+
+    rewriter.replaceOp(op, newInitArgs);
+    return mlir::success();
+  }
+};
+
 } // namespace
 
 void SignCastOp::getCanonicalizationPatterns(::mlir::RewritePatternSet &results,
@@ -1428,7 +1517,7 @@ void SignCastOp::getCanonicalizationPatterns(::mlir::RewritePatternSet &results,
       SignCastAllocPropagate<mlir::memref::AllocaOp>,
       SignCastTensorFromElementsPropagate, SignCastTensorCollapseShapePropagate,
       SignCastTensorToMemrefPropagate, SignCastMemrefToTensorPropagate,
-      SignCastMemrefSubviewPropagate>(context);
+      SignCastMemrefSubviewPropagate, SignCastForPropagate>(context);
 }
 
 void ExtractMemrefMetadataOp::build(::mlir::OpBuilder &odsBuilder,

numba_dpcomp/numba_dpcomp/mlir/tests/test_numba_parfor.py

@@ -52,33 +52,20 @@ def _gen_tests():
     ]
 
     xfail_tests = {
-        "test_prange26",
-        "test_prange03mul",
         "test_prange09",
         "test_prange03sub",
-        "test_prange10",
-        "test_prange03",
         "test_prange03div",
         "test_prange07",
         "test_prange06",
-        "test_prange16",
         "test_prange12",
-        "test_prange04",
-        "test_prange13",
         "test_prange25",
-        "test_prange21",
-        "test_prange14",
         "test_prange18",
         "test_prange_nested_reduction1",
         "test_list_setitem_hoisting",
-        "test_prange23",
-        "test_prange24",
         "test_list_comprehension_prange",
-        "test_prange22",
         "test_prange_raises_invalid_step_size",
         "test_issue7501",
         "test_parfor_race_1",
-        "test_check_alias_analysis",
         "test_nested_parfor_push_call_vars",
         "test_record_array_setitem_yield_array",
         "test_record_array_setitem",
@@ -127,7 +114,6 @@ def _gen_tests():
         "test_namedtuple2",
         "test_simple19",
         "test_no_hoisting_with_member_function_call",
-        "test_reduction_var_reuse",
         "test_parfor_dtype_type",
         "test_tuple3",
         "test_parfor_array_access3",
@@ -179,7 +165,6 @@ def _gen_tests():
         "test_parfor_generate_fuse",
         "test_parfor_slice7",
         "test_parfor_bitmask2",
-        "test_parfor_alias2",
         "test_one_d_array_reduction",
     }
 
@@ -217,19 +202,31 @@ def wrapper(*args, **kwargs):
                 return wrapper
 
             def _gen_parallel_fastmath(self, func):
+                ops = (
+                    "fadd",
+                    "fsub",
+                    "fmul",
+                    "fdiv",
+                    "frem",
+                    "fcmp",
+                )
+
                 def wrapper(*args, **kwargs):
                     with print_pass_ir([], ["PostLLVMLowering"]):
                         res = njit(parallel=True, fastmath=True)(func)(*args, **kwargs)
                         ir = get_print_buffer()
                         # Check some fastmath llvm flags were generated
-                        count = 0
+                        opCount = 0
+                        fastCount = 0
                         for line in ir.splitlines():
-                            for op in ("fadd", "fsub", "fmul", "fdiv", "frem", "fcmp"):
-                                if line.count("llvm." + op) and line.count(
-                                    "llvm.fastmath<fast>"
-                                ):
-                                    count += 1
-                        assert count > 0, ir
+                            for op in ops:
+                                if line.count("llvm." + op) > 0:
+                                    opCount += 1
+                                    if line.count("llvm.fastmath<fast>") > 0:
+                                        fastCount += 1
+                                    break
+                        if opCount > 0:
+                            assert fastCount > 0, it
                     return res
 
                 return wrapper