[flang-commits] [flang] [flang] Set LLVM specific attributes to fir.call's of Fortran runtime. (PR #128093)

[Slava Zakharin via flang-commits]

================
@@ -0,0 +1,252 @@
+//===- SetRuntimeCallAttributes.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
+//
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+/// \file
+/// SetRuntimeCallAttributesPass looks for fir.call operations
+/// that are calling into Fortran runtime, and tries to set different
+/// attributes on them to enable more optimizations in LLVM backend
+/// (granted that they are preserved all the way to LLVM IR).
+/// This pass is currently only attaching fir.call wide atttributes,
+/// such as ones corresponding to llvm.memory, nosync, nocallbac, etc.
+/// It is not designed to attach attributes to the arguments and the results
+/// of a call.
+//===----------------------------------------------------------------------===//
+#include "flang/Common/static-multimap-view.h"
+#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "flang/Optimizer/Support/InternalNames.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "flang/Runtime/io-api-consts.h"
+#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
+
+namespace fir {
+#define GEN_PASS_DEF_SETRUNTIMECALLATTRIBUTES
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+#define DEBUG_TYPE "set-runtime-call-attrs"
+
+using namespace Fortran::runtime;
+using namespace Fortran::runtime::io;
+
+#define mkIOKey(X) FirmkKey(IONAME(X))
+#define mkRTKey(X) FirmkKey(RTNAME(X))
+
+// Return LLVM dialect MemoryEffectsAttr for the given Fortran runtime call.
+// This function is computing a generic value of this attribute
+// by analyzing the arguments and their types.
+// It tries to figure out if an "indirect" memory access is possible
+// during this call. If it is not possible, then the memory effects
+// are:
+//   * other = NoModRef
+//   * argMem = ModRef
+//   * inaccessibleMem = ModRef
+//
+// Otherwise, it returns an empty attribute meaning ModRef for all kinds
+// of memory.
+//
+// The attribute deduction is conservative in a sense that it applies
+// to most of the runtime calls, but it may still be incorrect for some
+// runtime calls.
+static mlir::LLVM::MemoryEffectsAttr getGenericMemoryAttr(fir::CallOp callOp) {
+  bool maybeIndirectAccess = false;
+  for (auto arg : callOp.getArgOperands()) {
+    mlir::Type argType = arg.getType();
+    if (mlir::isa<fir::BaseBoxType>(argType)) {
+      // If it is a null box, then this particular call
+      // cannot access memory indirectly through the box's
+      // base_addr.
+      auto def = arg.getDefiningOp();
+      if (!def || !mlir::isa<fir::ZeroOp>(def)) {
+        maybeIndirectAccess = true;
+        break;
+      }
+    }
+    if (auto refType = mlir::dyn_cast<fir::ReferenceType>(argType)) {
+      if (!fir::isa_trivial(refType.getElementType())) {
+        maybeIndirectAccess = true;
+        break;
+      }
+    }
+    if (auto ptrType = mlir::dyn_cast<mlir::LLVM::LLVMPointerType>(argType)) {
+      maybeIndirectAccess = true;
+      break;
+    }
+  }
+  if (!maybeIndirectAccess) {
+    return mlir::LLVM::MemoryEffectsAttr::get(
+        callOp->getContext(),
+        {/*other=*/mlir::LLVM::ModRefInfo::NoModRef,
+         /*argMem=*/mlir::LLVM::ModRefInfo::ModRef,
+         /*inaccessibleMem=*/mlir::LLVM::ModRefInfo::ModRef});
+  }
+
+  return {};
+}
+
+namespace {
+class SetRuntimeCallAttributesPass
+    : public fir::impl::SetRuntimeCallAttributesBase<
+          SetRuntimeCallAttributesPass> {
+public:
+  void runOnOperation() override;
+};
+
+// A helper to match a type against a list of types.
+template <typename T, typename... Ts>
+constexpr bool IsAny = std::disjunction_v<std::is_same<T, Ts>...>;
+} // end anonymous namespace
+
+// MemoryAttrDesc type provides get() method for computing
+// mlir::LLVM::MemoryEffectsAttr for the given Fortran runtime call.
+// If needed, add specializations for particular runtime calls.
+namespace {
+// Default implementation just uses getGenericMemoryAttr().
+// Note that it may be incorrect for some runtime calls.
+template <typename KEY, typename Enable = void>
+struct MemoryAttrDesc {
+  static mlir::LLVM::MemoryEffectsAttr get(fir::CallOp callOp) {
+    return getGenericMemoryAttr(callOp);
+  }
+};
+} // end anonymous namespace
+
+// NosyncAttrDesc type provides get() method for computing
+// LLVM nosync attribute for the given call.
+namespace {
+// Default implementation always returns LLVM nosync.
+// This should be true for the majority of the Fortran runtime calls.
+template <typename KEY, typename Enable = void>
+struct NosyncAttrDesc {
+  static std::optional<mlir::NamedAttribute> get(fir::CallOp callOp) {
+    // TODO: replace llvm.nosync with an LLVM dialect callback.
+    return mlir::NamedAttribute("llvm.nosync",
+                                mlir::UnitAttr::get(callOp->getContext()));
+  }
+};
+} // end anonymous namespace
+
+// NocallbackAttrDesc type provides get() method for computing
+// LLVM nocallback attribute for the given call.
+namespace {
+// Default implementation always returns LLVM nocallback.
+// It must be specialized for Fortran runtime functions that may call
+// user functions during their execution (e.g. defined IO, assignment).
+template <typename KEY, typename Enable = void>
+struct NocallbackAttrDesc {
+  static std::optional<mlir::NamedAttribute> get(fir::CallOp callOp) {
+    // TODO: replace llvm.nocallback with an LLVM dialect callback.
+    return mlir::NamedAttribute("llvm.nocallback",
+                                mlir::UnitAttr::get(callOp->getContext()));
+  }
+};
+
+// Derived types IO may call back into a Fortran module.
----------------
vzakhari wrote:

You are right, Tom.  I will have to provide specializations for those non-I/O APIs.

Another point is that distinguishing these cases from trivial ones (when no callback may happen) might be not qutie reliable that late in the pipeline, so I think it may be worth introducing separate APIs for trivial and non-trivial cases of assignment/etc.  For the trivial cases, this will also help to avoid generating Fortran runtime calls that end up using call recursion, which is a problem for some offload devices.

If we split the APIs in such a way, then we will be able to set the attributes based on the APIs name rather than based on the recognition of the actual arguments.

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