[flang] add FIR to FIR pass to lower assumed-rank operations (llvm#93344)

jeanPerier · web-flow · commit f1d13bbd6662 · 2024-05-27T11:45:39.000+02:00
Add pass to lower assumed-rank operations. The current patch adds
codegen for fir.rebox_assumed_rank. It will be the pass lowering
fir.select_rank.
    
fir.rebox_assumed_rank is lowered to a call to CopyAndUpdateDescriptor
runtime API.
    
Note that the lowering ends-up allocating two new descriptors at the
LLVM level (one alloca created by the pass for the CopyAndUpdateDescriptor
result descriptor argument, the second one is created by the fir.load
of the result descriptor in codegen).
LLVM is currently unable to properly optimize and merge those allocas.
The "nocapture" attribute added to CopyAndUpdateDescriptor arguments
gives part of the information to LLVM, but the fir.load codegen of
descriptors must be updated to use llvm.memcpy instead of
llvm.load+store to allow LLVM to optimize it. This will be done in later patch.
diff --git a/flang/include/flang/Optimizer/Builder/FIRBuilder.h b/flang/include/flang/Optimizer/Builder/FIRBuilder.h
@@ -50,8 +50,10 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
                         mlir::SymbolTable *symbolTable = nullptr)
       : OpBuilder{op, /*listener=*/this}, kindMap{std::move(kindMap)},
         symbolTable{symbolTable} {}
-  explicit FirOpBuilder(mlir::OpBuilder &builder, fir::KindMapping kindMap)
-      : OpBuilder(builder), OpBuilder::Listener(), kindMap{std::move(kindMap)} {
+  explicit FirOpBuilder(mlir::OpBuilder &builder, fir::KindMapping kindMap,
+                        mlir::SymbolTable *symbolTable = nullptr)
+      : OpBuilder(builder), OpBuilder::Listener(), kindMap{std::move(kindMap)},
+        symbolTable{symbolTable} {
     setListener(this);
   }
   explicit FirOpBuilder(mlir::OpBuilder &builder, mlir::ModuleOp mod)
diff --git a/flang/include/flang/Optimizer/Builder/Runtime/RTBuilder.h b/flang/include/flang/Optimizer/Builder/Runtime/RTBuilder.h
@@ -130,6 +130,12 @@ constexpr TypeBuilderFunc getModel<signed char>() {
   };
 }
 template <>
+constexpr TypeBuilderFunc getModel<unsigned char>() {
+  return [](mlir::MLIRContext *context) -> mlir::Type {
+    return mlir::IntegerType::get(context, 8 * sizeof(unsigned char));
+  };
+}
+template <>
 constexpr TypeBuilderFunc getModel<void *>() {
   return [](mlir::MLIRContext *context) -> mlir::Type {
     return fir::LLVMPointerType::get(context,
diff --git a/flang/include/flang/Optimizer/Builder/Runtime/Support.h b/flang/include/flang/Optimizer/Builder/Runtime/Support.h
@@ -0,0 +1,31 @@
+//===-- Support.h - generate support runtime API calls ----------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_OPTIMIZER_BUILDER_RUNTIME_SUPPORT_H
+#define FORTRAN_OPTIMIZER_BUILDER_RUNTIME_SUPPORT_H
+
+namespace mlir {
+class Value;
+class Location;
+} // namespace mlir
+
+namespace fir {
+class FirOpBuilder;
+}
+
+namespace fir::runtime {
+
+/// Generate call to `CopyAndUpdateDescriptor` runtime routine.
+void genCopyAndUpdateDescriptor(fir::FirOpBuilder &builder, mlir::Location loc,
+                                mlir::Value to, mlir::Value from,
+                                mlir::Value newDynamicType,
+                                mlir::Value newAttribute,
+                                mlir::Value newLowerBounds);
+
+} // namespace fir::runtime
+#endif // FORTRAN_OPTIMIZER_BUILDER_RUNTIME_SUPPORT_H
diff --git a/flang/include/flang/Optimizer/Dialect/FIRType.h b/flang/include/flang/Optimizer/Dialect/FIRType.h
@@ -53,6 +53,7 @@ class BaseBoxType : public mlir::Type {
   /// Return the same type, except for the shape, that is taken the shape
   /// of shapeMold.
   BaseBoxType getBoxTypeWithNewShape(mlir::Type shapeMold) const;
+  BaseBoxType getBoxTypeWithNewShape(int rank) const;
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(mlir::Type type);
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.h b/flang/include/flang/Optimizer/Transforms/Passes.h
@@ -36,6 +36,7 @@ namespace fir {
 #define GEN_PASS_DECL_AFFINEDIALECTDEMOTION
 #define GEN_PASS_DECL_ANNOTATECONSTANTOPERANDS
 #define GEN_PASS_DECL_ARRAYVALUECOPY
+#define GEN_PASS_DECL_ASSUMEDRANKOPCONVERSION
 #define GEN_PASS_DECL_CHARACTERCONVERSION
 #define GEN_PASS_DECL_CFGCONVERSION
 #define GEN_PASS_DECL_EXTERNALNAMECONVERSION
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -402,4 +402,16 @@ def FunctionAttr : Pass<"function-attr", "mlir::func::FuncOp"> {
   let constructor = "::fir::createFunctionAttrPass()";
 }
 
+def AssumedRankOpConversion : Pass<"fir-assumed-rank-op", "mlir::ModuleOp"> {
+  let summary =
+    "Simplify operations on assumed-rank types";
+  let description = [{
+    This pass breaks up the lowering of operations on assumed-rank types by 
+    introducing an intermediate FIR level that simplifies code generation. 
+  }];
+  let dependentDialects = [
+    "fir::FIROpsDialect", "mlir::func::FuncDialect"
+  ];
+}
+
 #endif // FLANG_OPTIMIZER_TRANSFORMS_PASSES
diff --git a/flang/include/flang/Tools/CLOptions.inc b/flang/include/flang/Tools/CLOptions.inc
@@ -292,6 +292,7 @@ inline void createDefaultFIROptimizerPassPipeline(
 
   // Polymorphic types
   pm.addPass(fir::createPolymorphicOpConversion());
+  pm.addPass(fir::createAssumedRankOpConversion());
 
   if (pc.AliasAnalysis && !disableFirAliasTags && !useOldAliasTags)
     pm.addPass(fir::createAddAliasTags());
diff --git a/flang/lib/Optimizer/Builder/CMakeLists.txt b/flang/lib/Optimizer/Builder/CMakeLists.txt
@@ -29,6 +29,7 @@ add_flang_library(FIRBuilder
   Runtime/Ragged.cpp
   Runtime/Reduction.cpp
   Runtime/Stop.cpp
+  Runtime/Support.cpp
   Runtime/TemporaryStack.cpp
   Runtime/Transformational.cpp
   TemporaryStorage.cpp
diff --git a/flang/lib/Optimizer/Builder/Runtime/Support.cpp b/flang/lib/Optimizer/Builder/Runtime/Support.cpp
@@ -0,0 +1,46 @@
+//===-- Support.cpp - generate support runtime API calls --------*- C++ -*-===//
+//
+// 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/Runtime/Support.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"
+#include "flang/Runtime/support.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+
+using namespace Fortran::runtime;
+
+template <>
+constexpr fir::runtime::TypeBuilderFunc
+fir::runtime::getModel<Fortran::runtime::LowerBoundModifier>() {
+  return [](mlir::MLIRContext *context) -> mlir::Type {
+    return mlir::IntegerType::get(
+        context, sizeof(Fortran::runtime::LowerBoundModifier) * 8);
+  };
+}
+
+void fir::runtime::genCopyAndUpdateDescriptor(fir::FirOpBuilder &builder,
+                                              mlir::Location loc,
+                                              mlir::Value to, mlir::Value from,
+                                              mlir::Value newDynamicType,
+                                              mlir::Value newAttribute,
+                                              mlir::Value newLowerBounds) {
+  mlir::func::FuncOp func =
+      fir::runtime::getRuntimeFunc<mkRTKey(CopyAndUpdateDescriptor)>(loc,
+                                                                     builder);
+  auto fTy = func.getFunctionType();
+  auto args =
+      fir::runtime::createArguments(builder, loc, fTy, to, from, newDynamicType,
+                                    newAttribute, newLowerBounds);
+  llvm::StringRef noCapture = mlir::LLVM::LLVMDialect::getNoCaptureAttrName();
+  if (!func.getArgAttr(0, noCapture)) {
+    mlir::UnitAttr unitAttr = mlir::UnitAttr::get(func.getContext());
+    func.setArgAttr(0, noCapture, unitAttr);
+    func.setArgAttr(1, noCapture, unitAttr);
+  }
+  builder.create<fir::CallOp>(loc, func, args);
+}
diff --git a/flang/lib/Optimizer/Dialect/FIRType.cpp b/flang/lib/Optimizer/Dialect/FIRType.cpp
@@ -1324,6 +1324,17 @@ fir::BaseBoxType::getBoxTypeWithNewShape(mlir::Type shapeMold) const {
   return mlir::cast<fir::BaseBoxType>(changeTypeShape(*this, newShape));
 }
 
+fir::BaseBoxType fir::BaseBoxType::getBoxTypeWithNewShape(int rank) const {
+  std::optional<fir::SequenceType::ShapeRef> newShape;
+  fir::SequenceType::Shape shapeVector;
+  if (rank > 0) {
+    shapeVector =
+        fir::SequenceType::Shape(rank, fir::SequenceType::getUnknownExtent());
+    newShape = shapeVector;
+  }
+  return mlir::cast<fir::BaseBoxType>(changeTypeShape(*this, newShape));
+}
+
 bool fir::BaseBoxType::isAssumedRank() const {
   if (auto seqTy =
           mlir::dyn_cast<fir::SequenceType>(fir::unwrapRefType(getEleTy())))
diff --git a/flang/lib/Optimizer/Transforms/AssumedRankOpConversion.cpp b/flang/lib/Optimizer/Transforms/AssumedRankOpConversion.cpp
@@ -0,0 +1,131 @@
+//===-- AssumedRankOpConversion.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/Common/Fortran.h"
+#include "flang/Lower/BuiltinModules.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/Runtime/Support.h"
+#include "flang/Optimizer/Builder/Todo.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Support/TypeCode.h"
+#include "flang/Optimizer/Support/Utils.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "flang/Runtime/support.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace fir {
+#define GEN_PASS_DEF_ASSUMEDRANKOPCONVERSION
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+using namespace fir;
+using namespace mlir;
+
+namespace {
+
+static int getCFIAttribute(mlir::Type boxType) {
+  if (fir::isAllocatableType(boxType))
+    return CFI_attribute_allocatable;
+  if (fir::isPointerType(boxType))
+    return CFI_attribute_pointer;
+  return CFI_attribute_other;
+}
+
+static Fortran::runtime::LowerBoundModifier
+getLowerBoundModifier(fir::LowerBoundModifierAttribute modifier) {
+  switch (modifier) {
+  case fir::LowerBoundModifierAttribute::Preserve:
+    return Fortran::runtime::LowerBoundModifier::Preserve;
+  case fir::LowerBoundModifierAttribute::SetToOnes:
+    return Fortran::runtime::LowerBoundModifier::SetToOnes;
+  case fir::LowerBoundModifierAttribute::SetToZeroes:
+    return Fortran::runtime::LowerBoundModifier::SetToZeroes;
+  }
+  llvm_unreachable("bad modifier code");
+}
+
+class ReboxAssumedRankConv
+    : public mlir::OpRewritePattern<fir::ReboxAssumedRankOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  ReboxAssumedRankConv(mlir::MLIRContext *context,
+                       mlir::SymbolTable *symbolTable, fir::KindMapping kindMap)
+      : mlir::OpRewritePattern<fir::ReboxAssumedRankOp>(context),
+        symbolTable{symbolTable}, kindMap{kindMap} {};
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::ReboxAssumedRankOp rebox,
+                  mlir::PatternRewriter &rewriter) const override {
+    fir::FirOpBuilder builder{rewriter, kindMap, symbolTable};
+    mlir::Location loc = rebox.getLoc();
+    auto newBoxType = mlir::cast<fir::BaseBoxType>(rebox.getType());
+    mlir::Type newMaxRankBoxType =
+        newBoxType.getBoxTypeWithNewShape(Fortran::common::maxRank);
+    // CopyAndUpdateDescriptor FIR interface requires loading
+    // !fir.ref<fir.box> input which is expensive with assumed-rank. It could
+    // be best to add an entry point that takes a non "const" from to cover
+    // this case, but it would be good to indicate to LLVM that from does not
+    // get modified.
+    if (fir::isBoxAddress(rebox.getBox().getType()))
+      TODO(loc, "fir.rebox_assumed_rank codegen with fir.ref<fir.box<>> input");
+    mlir::Value tempDesc = builder.createTemporary(loc, newMaxRankBoxType);
+    mlir::Value newDtype;
+    mlir::Type newEleType = newBoxType.unwrapInnerType();
+    auto oldBoxType = mlir::cast<fir::BaseBoxType>(
+        fir::unwrapRefType(rebox.getBox().getType()));
+    auto newDerivedType = mlir::dyn_cast<fir::RecordType>(newEleType);
+    if (newDerivedType && (newEleType != oldBoxType.unwrapInnerType()) &&
+        !fir::isPolymorphicType(newBoxType)) {
+      newDtype = builder.create<fir::TypeDescOp>(
+          loc, mlir::TypeAttr::get(newDerivedType));
+    } else {
+      newDtype = builder.createNullConstant(loc);
+    }
+    mlir::Value newAttribute = builder.createIntegerConstant(
+        loc, builder.getIntegerType(8), getCFIAttribute(newBoxType));
+    int lbsModifierCode =
+        static_cast<int>(getLowerBoundModifier(rebox.getLbsModifier()));
+    mlir::Value lowerBoundModifier = builder.createIntegerConstant(
+        loc, builder.getIntegerType(32), lbsModifierCode);
+    fir::runtime::genCopyAndUpdateDescriptor(builder, loc, tempDesc,
+                                             rebox.getBox(), newDtype,
+                                             newAttribute, lowerBoundModifier);
+
+    mlir::Value descValue = builder.create<fir::LoadOp>(loc, tempDesc);
+    mlir::Value castDesc = builder.createConvert(loc, newBoxType, descValue);
+    rewriter.replaceOp(rebox, castDesc);
+    return mlir::success();
+  }
+
+private:
+  mlir::SymbolTable *symbolTable = nullptr;
+  fir::KindMapping kindMap;
+};
+
+/// Convert FIR structured control flow ops to CFG ops.
+class AssumedRankOpConversion
+    : public fir::impl::AssumedRankOpConversionBase<AssumedRankOpConversion> {
+public:
+  void runOnOperation() override {
+    auto *context = &getContext();
+    mlir::ModuleOp mod = getOperation();
+    mlir::SymbolTable symbolTable(mod);
+    fir::KindMapping kindMap = fir::getKindMapping(mod);
+    mlir::RewritePatternSet patterns(context);
+    patterns.insert<ReboxAssumedRankConv>(context, &symbolTable, kindMap);
+    mlir::GreedyRewriteConfig config;
+    config.enableRegionSimplification = false;
+    (void)applyPatternsAndFoldGreedily(mod, std::move(patterns), config);
+  }
+};
+} // namespace
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -4,6 +4,7 @@ add_flang_library(FIRTransforms
   AffinePromotion.cpp
   AffineDemotion.cpp
   AnnotateConstant.cpp
+  AssumedRankOpConversion.cpp
   CharacterConversion.cpp
   ControlFlowConverter.cpp
   ArrayValueCopy.cpp
diff --git a/flang/test/Driver/bbc-mlir-pass-pipeline.f90 b/flang/test/Driver/bbc-mlir-pass-pipeline.f90
@@ -46,6 +46,7 @@
 ! CHECK-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 
 ! CHECK-NEXT: PolymorphicOpConversion
+! CHECK-NEXT: AssumedRankOpConversion
 
 ! CHECK-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
 ! CHECK-NEXT: 'fir.global' Pipeline
diff --git a/flang/test/Driver/mlir-debug-pass-pipeline.f90 b/flang/test/Driver/mlir-debug-pass-pipeline.f90
@@ -73,6 +73,7 @@
 ! ALL-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 
 ! ALL-NEXT: PolymorphicOpConversion
+! ALL-NEXT: AssumedRankOpConversion
 
 ! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
 ! ALL-NEXT:   'fir.global' Pipeline
diff --git a/flang/test/Driver/mlir-pass-pipeline.f90 b/flang/test/Driver/mlir-pass-pipeline.f90
@@ -80,6 +80,7 @@
 ! ALL-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 
 ! ALL-NEXT: PolymorphicOpConversion
+! ALL-NEXT: AssumedRankOpConversion
 ! O2-NEXT:  AddAliasTags
 
 ! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
diff --git a/flang/test/Fir/basic-program.fir b/flang/test/Fir/basic-program.fir
@@ -80,6 +80,7 @@ func.func @_QQmain() {
 // PASSES-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 
 // PASSES-NEXT: PolymorphicOpConversion
+// PASSES-NEXT: AssumedRankOpConversion
 // PASSES-NEXT: AddAliasTags
 
 // PASSES-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
diff --git a/flang/test/Fir/rebox_assumed_rank_codegen.fir b/flang/test/Fir/rebox_assumed_rank_codegen.fir
