[flang] Add AnnotateConstant pass

[schweitz]

Summary

Annotate constant arguments to FIR ops.

[clementval]

LGTM

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

 std::unique_ptr<mlir::Pass> createFirToCfgPass();
 std::unique_ptr<mlir::Pass> createCharacterConversionPass();
 std::unique_ptr<mlir::Pass> createExternalNameConversionPass();
 std::unique_ptr<mlir::Pass> createMemDataFlowOptPass();
 std::unique_ptr<mlir::Pass> createPromoteToAffinePass();
 std::unique_ptr<mlir::Pass> createMemoryAllocationPass();
 std::unique_ptr<mlir::Pass>
 createMemoryAllocationPass(bool dynOnHeap, std::size_t maxStackSize);
+std::unique_ptr<mlir::Pass> createAnnotateConstantOperandsPass();
 
 // declarative passes
 #define GEN_PASS_REGISTRATION
 #include "flang/Optimizer/Transforms/Passes.h.inc"
 
 } // namespace fir
 
 #endif // FORTRAN_OPTIMIZER_TRANSFORMS_PASSES_H

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

let description = [{
     `fir.load` and `fir.store` with `!fir.ref<?>` types is required.
   }];
   let constructor = "::fir::createAffineDemotionPass()";
   let dependentDialects = [
     "fir::FIROpsDialect", "mlir::func::FuncDialect", "mlir::AffineDialect"
   ];
 }
 
+def AnnotateConstantOperands : Pass<"annotate-constant"> {
+  let summary = "Annotate constant operands to all FIR operations";
+  let description = [{
+    The MLIR canonicalizer makes a distinction between constants based on how
+    they are packaged in the IR. A constant value is wrapped in an Attr and that
+    Attr can be attached to an Op. There is a distinguished Op, ConstantOp, that
+    merely has one of these Attr attached.
+
+    The MLIR canonicalizer treats constants referenced by an Op and constants
+    referenced through a ConstantOp as having distinct semantics. This pass
+    eliminates that distinction, so hashconsing of Ops, basic blocks, etc.
+    behaves as one would expect.
+  }];
+  let constructor = "::fir::createAnnotateConstantOperandsPass()";
+  let dependentDialects = [ "fir::FIROpsDialect" ];
+}
+
 def ArrayValueCopy : Pass<"array-value-copy", "::mlir::FuncOp"> {
   let summary = "Convert array value operations to memory operations.";
   let description = [{
     Transform the set of array value primitives to a memory-based array
     representation.
 
     The Ops `array_load`, `array_store`, `array_fetch`, and `array_update` are
     used to manage abstract aggregate array values. A simple analysis is done
     to determine if there are potential dependences between these operations.
     If not, these array operations can be lowered to work directly on the memory
     representation. If there is a potential conflict, a temporary is created
     along with appropriate copy-in/copy-out operations. Here, a more refined
     analysis might be deployed, such as using the affine framework.
 
     This pass is required before code gen to the LLVM IR dialect.
   }];
   let constructor = "::fir::createArrayValueCopyPass()";
+  let dependentDialects = [ "fir::FIROpsDialect" ];
 }
 
 def CharacterConversion : Pass<"character-conversion"> {
   let summary = "Convert CHARACTER entities with different KINDs";
   let description = [{
     Translates entities of one CHARACTER KIND to another.
 
     By default the translation is to naively zero-extend or truncate a code

flang/lib/Optimizer/Transforms/AnnotateConstant.cpp

+//===-- AnnotateConstant.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 "PassDetail.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+
+#define DEBUG_TYPE "flang-annotate-constant"
+
+using namespace fir;
+
+namespace {
+struct AnnotateConstantOperands
+    : AnnotateConstantOperandsBase<AnnotateConstantOperands> {
+  void runOnOperation() override {
+    auto *context = &getContext();
+    mlir::Dialect *firDialect = context->getLoadedDialect("fir");
+    getOperation()->walk([&](mlir::Operation *op) {
+      // We filter out other dialects even though they may undergo merging of
+      // non-equal constant values by the canonicalizer as well.
+      if (op->getDialect() == firDialect) {
+        llvm::SmallVector<mlir::Attribute> attrs;
+        bool hasOneOrMoreConstOpnd = false;
+        for (mlir::Value opnd : op->getOperands()) {
+          if (auto constOp = mlir::dyn_cast_or_null<mlir::arith::ConstantOp>(
+                  opnd.getDefiningOp())) {
+            attrs.push_back(constOp.getValue());
+            hasOneOrMoreConstOpnd = true;
+          } else if (auto addrOp = mlir::dyn_cast_or_null<fir::AddrOfOp>(
+                         opnd.getDefiningOp())) {
+            attrs.push_back(addrOp.getSymbol());
+            hasOneOrMoreConstOpnd = true;
+          } else {
+            attrs.push_back(mlir::UnitAttr::get(context));
+          }
+        }
+        if (hasOneOrMoreConstOpnd)
+          op->setAttr("canonicalize_constant_operands",
+                      mlir::ArrayAttr::get(context, attrs));
+      }
+    });
+  }
+};
+
+} // namespace
+
+std::unique_ptr<mlir::Pass> fir::createAnnotateConstantOperandsPass() {
+  return std::make_unique<AnnotateConstantOperands>();
+}

flang/lib/Optimizer/Transforms/CMakeLists.txt

 add_flang_library(FIRTransforms
   AbstractResult.cpp
   AffinePromotion.cpp
   AffineDemotion.cpp
+  AnnotateConstant.cpp
   CharacterConversion.cpp
   ArrayValueCopy.cpp
   ExternalNameConversion.cpp
   MemoryAllocation.cpp
   MemRefDataFlowOpt.cpp
   RewriteLoop.cpp
 
   DEPENDS

flang/test/Fir/annotate-constant.fir

+// RUN: fir-opt -annotate-constant %s | FileCheck %s
+
+// CHECK-LABEL: func @annotate_test() -> !fir.ref<!fir.array<?xi32>> {
+func @annotate_test() -> !fir.ref<!fir.array<?xi32>> {
+  %1 = arith.constant 5 : index
+  // CHECK: %[[a:.*]] = fir.alloca !fir.array<?xi32>, %{{.*}} {canonicalize_constant_operands = [5 : index]}
+  %2 = fir.alloca !fir.array<?xi32>, %1
+  return %2 : !fir.ref<!fir.array<?xi32>>
+}