[flang-commits] [flang] [mlir] [MLIR][OpenMP] Changes to function-filtering pass (PR #71850)

[Sergio Afonso via flang-commits]

================
@@ -56,17 +56,30 @@ class OMPFunctionFilteringPass
       if (declareTargetOp && declareTargetOp.isDeclareTarget())
         declareType = declareTargetOp.getDeclareTargetDeviceType();
 
-      // Filtering a function here means removing its body and explicitly
-      // setting its omp.declare_target attribute, so that following
-      // translation/lowering/transformation passes will skip processing its
-      // contents, but preventing the calls to undefined symbols that could
-      // result if the function were deleted. The second stage of function
-      // filtering, at the MLIR to LLVM IR translation level, will remove these
-      // from the IR thanks to the mismatch between the omp.declare_target
-      // attribute and the target device.
+      // Filtering a function here means deleting it if it doesn't containt a
+      // target region. Else we explicitly setting its omp.declare_target
+      // attribute. The second stage of function filtering at the MLIR to LLVM
+      // IR translation level will remove functions that contain the target
+      // region from the generated llvm IR.
       if (declareType == omp::DeclareTargetDeviceType::host) {
-        funcOp.eraseBody();
-        funcOp.setVisibility(SymbolTable::Visibility::Private);
+        auto funcUses = *funcOp.getSymbolUses(op);
+        for (auto use : funcUses) {
+          auto *callOp = use.getUser();
+          // If the callOp has users then replace them with Undef values.
+          if (!callOp->use_empty()) {
+            SmallVector<Value> undefResults;
+            for (auto res : callOp->getResults()) {
+              opBuilder.setInsertionPoint(callOp);
+              undefResults.emplace_back(
+                  opBuilder.create<fir::UndefOp>(res.getLoc(), res.getType()));
+            }
+            callOp->replaceAllUsesWith(undefResults);
+          }
+          // Remove the callOp
+          callOp->erase();
+        }
+        if (!hasTargetRegion)
+          funcOp.erase();
         if (declareTargetOp)
----------------
skatrak wrote:

I think this should be inside of an `else` block or moved above the `if (!hasTargetRegion)` block. Otherwise, we're adding attributes to an operation we just deleted (`declareTargetOp` and `funcOp` point to the same operation).

https://github.com/llvm/llvm-project/pull/71850