[flang-commits] [flang] [Flang] Extracting internal constants from scalar literals (PR #73829)

[Mats Petersson via flang-commits]

================
@@ -0,0 +1,180 @@
+//===- ConstantArgumentGlobalisation.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
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace fir {
+#define GEN_PASS_DEF_CONSTANTARGUMENTGLOBALISATIONOPT
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+#define DEBUG_TYPE "flang-constang-argument-globalisation-opt"
+
+namespace {
+unsigned uniqueLitId = 1;
+
+class CallOpRewriter : public mlir::OpRewritePattern<fir::CallOp> {
+protected:
+  const mlir::DominanceInfo &di;
+
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  CallOpRewriter(mlir::MLIRContext *ctx, const mlir::DominanceInfo &_di)
+      : OpRewritePattern(ctx), di(_di) {}
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::CallOp callOp,
+                  mlir::PatternRewriter &rewriter) const override {
+    LLVM_DEBUG(llvm::dbgs() << "Processing call op: " << callOp << "\n");
+    auto module = callOp->getParentOfType<mlir::ModuleOp>();
+    bool needUpdate = false;
+    fir::FirOpBuilder builder(rewriter, module);
+    llvm::SmallVector<mlir::Value> newOperands;
+    llvm::SmallVector<mlir::Operation *> toErase;
+    for (const mlir::Value &a : callOp.getArgs()) {
+      auto alloca = mlir::dyn_cast_or_null<fir::AllocaOp>(a.getDefiningOp());
+      // We can convert arguments that are alloca, and that has
+      // the value by reference attribute. All else is just added
+      // to the argument list.
+      if (!alloca || !alloca->hasAttr(fir::getAdaptToByRefAttrName())) {
+        newOperands.push_back(a);
+        continue;
+      }
+
+      mlir::Type varTy = alloca.getInType();
+      assert(!fir::hasDynamicSize(varTy) &&
+             "only expect statically sized scalars to be by value");
+
+      // Find immediate store with const argument
+      mlir::Operation *store = nullptr;
+      for (mlir::Operation *s : alloca->getUsers()) {
+        if (mlir::isa<fir::StoreOp>(s) && di.dominates(s, callOp)) {
+          // We can only deal with ONE store - if already found one,
+          // set to nullptr and exit the loop.
+          if (store) {
+            store = nullptr;
+            break;
+          }
+          store = s;
+        }
+      }
+
+      // If we didn't find one signle store, add argument as is, and move on.
+      if (!store) {
+        newOperands.push_back(a);
+        continue;
+      }
+
+      LLVM_DEBUG(llvm::dbgs() << " found store " << *store << "\n");
+
+      mlir::Operation *constant_def = store->getOperand(0).getDefiningOp();
+      // Expect constant definition operation or force legalisation of the
+      // callOp and continue with its next argument
+      if (!mlir::isa<mlir::arith::ConstantOp>(constant_def)) {
+        // Unable to remove alloca arg
+        newOperands.push_back(a);
+        continue;
+      }
+
+      LLVM_DEBUG(llvm::dbgs() << " found define " << *constant_def << "\n");
+
+      std::string globalName =
+          "_global_const_." + std::to_string(uniqueLitId++);
+      assert(!builder.getNamedGlobal(globalName) &&
+             "We should have a unique name here");
+
+      unsigned count = 0;
+      for (mlir::Operation *s : alloca->getUsers())
+        if (di.dominates(store, s))
+          ++count;
+
+      // Delete if dominates itself and one more operation (which should
+      // be callOp)
+      if (count == 2)
+        toErase.push_back(store);
----------------
Leporacanthicus wrote:

Well, that's not so easy, because we don't create the new operation until outside the loop. I'll take another look into that tomorrow. I think we should cope with this case.

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