[Flang] Bail out if lower-workdistribute didn't patternmatch.

Author: skc7

flang-rt/lib/runtime/assign_omp.cpp

@@ -68,7 +68,7 @@ RT_EXT_API_GROUP_BEGIN
 void RTDEF(Assign_omp)(Descriptor &to, const Descriptor &from,
     const char *sourceFile, int sourceLine, omp::OMPDeviceTy omp_device) {
   Terminator terminator{sourceFile, sourceLine};
-  omp::Assign(to, from, terminator,
+  Fortran::runtime::omp::Assign(to, from, terminator,
       MaybeReallocate | NeedFinalization | ComponentCanBeDefinedAssignment,
       omp_device);
 }

flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp

@@ -125,6 +125,7 @@ static T getPerfectlyNested(Operation *op) {
 /// E()
 
 struct FissionWorkdistribute : public OpRewritePattern<omp::WorkdistributeOp> {
+  static bool fissionWorkdistributePatternMatched;
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(omp::WorkdistributeOp workdistribute,
                                 PatternRewriter &rewriter) const override {
@@ -210,9 +211,12 @@ struct FissionWorkdistribute : public OpRewritePattern<omp::WorkdistributeOp> {
         changed = true;
       }
     }
+    if (changed)
+      fissionWorkdistributePatternMatched = true;
     return success(changed);
   }
 };
+bool FissionWorkdistribute::fissionWorkdistributePatternMatched = false;
 
 /// If fir.do_loop is present inside teams workdistribute
 ///
@@ -296,6 +300,7 @@ static void genWsLoopOp(mlir::PatternRewriter &rewriter, fir::DoLoopOp doLoop,
 }
 
 struct WorkdistributeDoLower : public OpRewritePattern<omp::WorkdistributeOp> {
+  static bool workdistributeDoLowerPatternMatched;
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(omp::WorkdistributeOp workdistribute,
                                 PatternRewriter &rewriter) const override {
@@ -309,12 +314,15 @@ struct WorkdistributeDoLower : public OpRewritePattern<omp::WorkdistributeOp> {
       genLoopNestClauseOps(rewriter, doLoop, loopNestClauseOps);
       genWsLoopOp(rewriter, doLoop, loopNestClauseOps, true);
       rewriter.eraseOp(workdistribute);
+      workdistributeDoLowerPatternMatched = true;
       return success();
     }
     return failure();
   }
 };
 
+bool WorkdistributeDoLower::workdistributeDoLowerPatternMatched = false;
+
 /// If A() and B () are present inside teams workdistribute
 ///
 /// omp.teams {
@@ -331,6 +339,7 @@ struct WorkdistributeDoLower : public OpRewritePattern<omp::WorkdistributeOp> {
 ///
 
 struct TeamsWorkdistributeToSingle : public OpRewritePattern<omp::TeamsOp> {
+  static bool teamsWorkdistributeToSinglePatternMatched;
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(omp::TeamsOp teamsOp,
                                 PatternRewriter &rewriter) const override {
@@ -343,9 +352,12 @@ struct TeamsWorkdistributeToSingle : public OpRewritePattern<omp::TeamsOp> {
     rewriter.eraseOp(workdistributeBlock->getTerminator());
     rewriter.inlineBlockBefore(workdistributeBlock, teamsOp);
     rewriter.eraseOp(workdistributeOp);
+    teamsWorkdistributeToSinglePatternMatched = true;
     return success();
   }
 };
+bool TeamsWorkdistributeToSingle::teamsWorkdistributeToSinglePatternMatched =
+    false;
 
 struct SplitTargetResult {
   omp::TargetOp targetOp;
@@ -517,28 +529,6 @@ static bool usedOutsideSplit(Value v, Operation *split) {
   return false;
 };
 
-static bool isOpToBeCached(Operation *op) {
-  if (auto loadOp = dyn_cast<fir::LoadOp>(op)) {  
-    Value memref = loadOp.getMemref();  
-    if (auto blockArg = dyn_cast<BlockArgument>(memref)) {  
-      // 'op' is an operation within the targetOp that 'splitBefore' is also in.
-      Operation *parentOpOfLoadBlock = op->getBlock()->getParentOp();  
-      // Ensure the blockArg belongs to the entry block of this parent omp.TargetOp.  
-      // This implies the load is from a variable directly mapped into the target region.  
-      if (isa<omp::TargetOp>(parentOpOfLoadBlock) &&  
-          !parentOpOfLoadBlock->getRegions().empty()) {  
-        Block *targetOpEntryBlock = &parentOpOfLoadBlock->getRegions().front().front();  
-        if (blockArg.getOwner() == targetOpEntryBlock) {  
-          // This load is from a direct argument of the target op.  
-          // It's safe to recompute.
-          return false;  
-        }  
-      }  
-    }  
-  }
-  return true;
-}
-
 static bool isRecomputableAfterFission(Operation *op, Operation *splitBefore) {
   if (isa<fir::DeclareOp>(op))
     return true;
@@ -892,23 +882,31 @@ class LowerWorkdistributePass
     config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled);
 
     Operation *op = getOperation();
+    bool anyPatternChanged = false;
     {
       RewritePatternSet patterns(&context);
       patterns.insert<FissionWorkdistribute, WorkdistributeDoLower>(&context);
       if (failed(applyPatternsGreedily(op, std::move(patterns), config))) {
         emitError(op->getLoc(), DEBUG_TYPE " pass failed\n");
         signalPassFailure();
       }
+      anyPatternChanged |=
+          FissionWorkdistribute::fissionWorkdistributePatternMatched;
+      anyPatternChanged |=
+          WorkdistributeDoLower::workdistributeDoLowerPatternMatched;
     }
     {
       RewritePatternSet patterns(&context);
-      patterns.insert<TeamsWorkdistributeToSingle>(&context);
+      patterns.insert<WorkdistributeDoLower, TeamsWorkdistributeToSingle>(
+          &context);
       if (failed(applyPatternsGreedily(op, std::move(patterns), config))) {
         emitError(op->getLoc(), DEBUG_TYPE " pass failed\n");
         signalPassFailure();
       }
+      anyPatternChanged |= TeamsWorkdistributeToSingle::
+          teamsWorkdistributeToSinglePatternMatched;
     }
-    {
+    if (anyPatternChanged) {
       SmallVector<omp::TargetOp> targetOps;
       op->walk([&](omp::TargetOp targetOp) { targetOps.push_back(targetOp); });
       IRRewriter rewriter(&context);
@@ -917,7 +915,6 @@ class LowerWorkdistributePass
         if (res) fissionTarget(res->targetOp, rewriter);
       }
     }
-
   }
 };
 } // namespace

flang/test/Transforms/OpenMP/lower-workdistribute-target.mlir

@@ -10,8 +10,7 @@
 // CHECK:           %[[VAL_7:.*]] = fir.coordinate_of %[[VAL_6]], r : (!fir.ref<!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTnested{i:i32,r:f32}>>) -> !fir.ref<f32>
 // CHECK:           %[[VAL_8:.*]] = omp.map.info var_ptr(%[[VAL_7]] : !fir.ref<f32>, f32) map_clauses(tofrom) capture(ByRef) -> !fir.ref<f32> {name = "sa%[[VAL_4]]%[[VAL_9:.*]]"}
 // CHECK:           %[[VAL_10:.*]] = omp.map.info var_ptr(%[[VAL_0]] : !fir.ref<!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTscalar_and_array{r:f32,n:!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTnested{i:i32,r:f32}>}>>, !fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTscalar_and_array{r:f32,n:!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTnested{i:i32,r:f32}>}>) map_clauses(tofrom) capture(ByRef) members(%[[VAL_3]], %[[VAL_8]] : [1, 0], [1, 1] : !fir.ref<i32>, !fir.ref<f32>) -> !fir.ref<!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTscalar_and_array{r:f32,n:!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTnested{i:i32,r:f32}>}>> {name = "sa", partial_map = true}
-// CHECK:           %[[VAL_11:.*]] = omp.map.info var_ptr(%[[VAL_0]] : !fir.ref<!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTscalar_and_array{r:f32,n:!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTnested{i:i32,r:f32}>}>>, !fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTscalar_and_array{r:f32,n:!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTnested{i:i32,r:f32}>}>) map_clauses(exit_release_or_enter_alloc) capture(ByRef) members(%[[VAL_3]], %[[VAL_8]] : [1, 0], [1, 1] : !fir.ref<i32>, !fir.ref<f32>) -> !fir.ref<!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTscalar_and_array{r:f32,n:!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTnested{i:i32,r:f32}>}>> {name = "sa", partial_map = true}
-// CHECK:           omp.target_data map_entries(%[[VAL_10]] : !fir.ref<!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTscalar_and_array{r:f32,n:!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTnested{i:i32,r:f32}>}>>) {
+// CHECK:           omp.target map_entries(%[[VAL_10]] -> %[[VAL_11:.*]] : !fir.ref<!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTscalar_and_array{r:f32,n:!fir.type<_QFmaptype_derived_nested_explicit_multiple_membersTnested{i:i32,r:f32}>}>>) {
 // CHECK:             omp.terminator
 // CHECK:           }
 // CHECK:           return