[MLIR][Flang][OpenMP] Remove unused global variables in offload target module (llvm#97)

Author: DominikAdamski

flang/include/flang/Optimizer/Transforms/Passes.h

@@ -76,6 +76,7 @@ std::unique_ptr<mlir::Pass> createAlgebraicSimplificationPass();
 std::unique_ptr<mlir::Pass>
 createAlgebraicSimplificationPass(const mlir::GreedyRewriteConfig &config);
 
+std::unique_ptr<mlir::Pass> createOMPGlobalFilteringPass();
 std::unique_ptr<mlir::Pass> createVScaleAttrPass();
 std::unique_ptr<mlir::Pass>
 createVScaleAttrPass(std::pair<unsigned, unsigned> vscaleAttr);

flang/include/flang/Optimizer/Transforms/Passes.td

@@ -358,6 +358,16 @@ def OMPFunctionFiltering : Pass<"omp-function-filtering"> {
   ];
 }
 
+def OMPGlobalFiltering : Pass<"omp-global-filtering"> {
+  let summary = "Filters out globals intended for the host when compiling "
+                "for the target device.";
+  let constructor = "::fir::createOMPGlobalFilteringPass()";
+  let dependentDialects = [
+    "mlir::func::FuncDialect",
+    "fir::FIROpsDialect"
+  ];
+}
+
 def VScaleAttr : Pass<"vscale-attr", "mlir::func::FuncOp"> {
   let summary = "Add vscale_range attribute to functions";
   let description = [{

flang/include/flang/Tools/CLOptions.inc

@@ -355,8 +355,10 @@ inline void createOpenMPFIRPassPipeline(mlir::PassManager &pm,
       pm, fir::createOMPMapInfoFinalizationPass);
 
   pm.addPass(fir::createOMPMarkDeclareTargetPass());
-  if (isTargetDevice)
+  if (isTargetDevice) {
     pm.addPass(fir::createOMPFunctionFiltering());
+    pm.addPass(fir::createOMPGlobalFilteringPass());
+  }
 }
 
 #if !defined(FLANG_EXCLUDE_CODEGEN)

flang/lib/Optimizer/Transforms/CMakeLists.txt

@@ -19,6 +19,7 @@ add_flang_library(FIRTransforms
   PolymorphicOpConversion.cpp
   LoopVersioning.cpp
   OMPFunctionFiltering.cpp
+  OMPGlobalFiltering.cpp
   OMPMapInfoFinalization.cpp
   OMPMarkDeclareTarget.cpp
   StackReclaim.cpp

flang/lib/Optimizer/Transforms/OMPGlobalFiltering.cpp

@@ -0,0 +1,64 @@
+//===- OMPFunctionFiltering.cpp -------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements transforms to filter out functions intended for the host
+// when compiling for the device and vice versa.
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/Dialect/OpenMP/OpenMPInterfaces.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace fir {
+#define GEN_PASS_DEF_OMPGLOBALFILTERING
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+using namespace mlir;
+
+namespace {
+class OMPGlobalFilteringPass
+    : public fir::impl::OMPGlobalFilteringBase<OMPGlobalFilteringPass> {
+public:
+  OMPGlobalFilteringPass() = default;
+
+  void runOnOperation() override {
+    auto op = dyn_cast<omp::OffloadModuleInterface>(getOperation());
+    if (!op || !op.getIsTargetDevice())
+      return;
+
+    op->walk<WalkOrder::PreOrder>([&](fir::GlobalOp globalOp) {
+      bool symbolUnused = true;
+      SymbolTable::UseRange globalUses = *globalOp.getSymbolUses(op);
+      for (SymbolTable::SymbolUse use : globalUses) {
+        if (use.getUser() == globalOp)
+          continue;
+        symbolUnused = false;
+        break;
+      }
+
+      // Remove unused host symbols with external linkage
+      // TODO: Add support for declare target global variables
+      if (symbolUnused && !globalOp.getLinkName())
+        globalOp.erase();
+      return WalkResult::advance();
+    });
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> fir::createOMPGlobalFilteringPass() {
+  return std::make_unique<OMPGlobalFilteringPass>();
+}