[Flang][MLIR][OpenMP] Update OMPEarlyOutlining to support Bounds, MapEntry and declare target globals

[agozillon]

Summary

This patch is a required change for the device side IR to
maintain apporpiate links for declare target variables to
their global variables for later lowering.

It is also a requirement to clone over map bounds and
entry operations to maintain the correct information for
later lowering of the IR.

It simply tries to clone over the relevant information
maintaining the appropriate links they would have
maintained prior to the pass, rather than redirecting
them to new function arguments which causes a
loss of information in the case of Declare Target
and map information.

Depends on D158734

[agozillon]

Patch 3/4 of the map changes from the patch series made by: https://reviews.llvm.org/D158732, https://reviews.llvm.org/D158735, https://reviews.llvm.org/D158737 and https://reviews.llvm.org/D158734 that aims to expand the current map support of the OpenMP dialect and lower from Fortran -> MLIR -> LLVM IR, future support on expanding the lowering from the OpenMP dialect to LLVM IR for map on TargetOp (omp target) will be forthcoming for declare target and explicit map variables hopefully in the near future (implicit likely a little while after that).

This patch should not pass CI on its own, patch 4: https://reviews.llvm.org/D158737, should pass provided the dependencies are setup correctly.

[agozillon]

A small ping for some attention on this patch if a reviewer can spare some time, thank you very much!

[razvanlupusoru]

This change seems OK to me - but I have limited experience with the omp device pipeline.

Just one comment - in another change (https://reviews.llvm.org/D158734) I noticed an explicit mention about " Add To+USM mode case"
I am guessing the reason declare variables are handled specially here is because in some cases there is a device version which can be resolved at link time and thus doesn't need passed into the kernel. I guess a safer alternate implementation is to always pass the pointer in - this would work regardless of the memory mode.

flang/lib/Optimizer/Transforms/OMPEarlyOutlining.cpp

228

Commented out code should be removed.

[agozillon]

In D158735#4632349, @razvanlupusoru wrote:

This change seems OK to me - but I have limited experience with the omp device pipeline.

Just one comment - in another change (https://reviews.llvm.org/D158734) I noticed an explicit mention about " Add To+USM mode case"

Yes, a lot of the flow in that segment (and the lowering I'll create patches for on Github eventually) try to mimic segments of Clang's current code generation process. And I think that specific check is something related to a specific case the OpenMP specification requests, but I'd have to look up the specifics as I don't have them on hand at the moment.

I am guessing the reason declare variables are handled specially here is because in some cases there is a device version which can be resolved at link time and thus doesn't need passed into the kernel. I guess a safer alternate implementation is to always pass the pointer in - this would work regardless of the memory mode.

I believe declare target "to/enter" clause variables are statically allocated on the devices data environment, whereas I think declare target "link" variables mapping/lifetime is handled by the runtime using OpenMP's data mapping control rules (emitted to IR as a nullptr and the runtime magics the value into it for the device to use), unsure on "indirect" as it's not something I've looked into handling at the moment. Someone with more OpenMP knowledge can likely fact check me on this, as it's been a little while since I've taken a look at declare targets semantics and I imagine I've gotten something wrong (and learnings always good!). Declare target variables are device side globals for the most part I believe that persist for the length of there data environment (longer perhaps).

However, It's quite possible that it would be a safer implementation, but the lowering, for the moment, is trying to stick somewhat close to the IR emitted from Clang's OpenMP implementation in part to try and achieve a shared implementation but also to make it easier to complete features/reach the standard Clang currently has, the old if it ain't broken don't fix it adage. But I imagine as we progress further we'll deviate, especially as we try to become more performant rather than just conformant.

Hopefully this makes some sense, it's a little late at night here!

flang/lib/Optimizer/Transforms/OMPEarlyOutlining.cpp

228

nice catch, thank you! that one snuck by me it seems.

[agozillon]

In D158735#4632349, @razvanlupusoru wrote:

This change seems OK to me - but I have limited experience with the omp device pipeline.

Just one comment - in another change (https://reviews.llvm.org/D158734) I noticed an explicit mention about " Add To+USM mode case"
I am guessing the reason declare variables are handled specially here is because in some cases there is a device version which can be resolved at link time and thus doesn't need passed into the kernel. I guess a safer alternate implementation is to always pass the pointer in - this would work regardless of the memory mode.

Sorry, I managed to conflate two different issues with the last message (that teaches me for trying to answer one last message when I'm tired I suppose)! The more appropriate answer is that this pass simply outlines the target region into a separate function, isolating it from other instructions (and optimisations to the region which can cause difficulties for the current lowering, which the pass is here to try and avoid). Most map clauses are mapped into arguments to this function then remapped to the target and the getUsedValuesDefinedAbove is a placeholder for implicit map captures that don't get added to map clauses upstream currently (but I have something that's wip downstream that will raise implicit captures into clauses). The issue with this is it loses information in certain cases like declare target as we'd be pointing at a block argument and the declare target interface only applies to globals and we'd be missing symbol information also required for the lowering (but perhaps worked around by renaming the block argument). In other cases we need to clone over specific operations (bounds, map entry and in the future likely fir box).

[agozillon]

remove left over comment

[razvanlupusoru]

Based on last explanation on why cloning is needed, this change seems reasonable to me.

[agozillon]

In D158735#4640008, @razvanlupusoru wrote:

Based on last explanation on why cloning is needed, this change seems reasonable to me.

Thank you very much for the review! If anyone else has anything they wish to add please feel free to do so.

[agozillon]

rebase

[agozillon]

@razvanlupusoru thank you very much for your review!

I will leave the patch series open until Thursday to give time for further comments/feedback from other reviewers (in particular @kiranchandramohan and @jsjodin) at that point if no further comments/change requests are received I will push it upstream!

[agozillon]

Update based on name changes to operation, just altering the pass slightly with the new type name for the op

[agozillon]

Updating early outliner pass to clone box types as it does downstream (also some slight extensions beyond it to handle certain edge cases better), tidying it up and rebasing as well.

flang/lib/Optimizer/Transforms/OMPEarlyOutlining.cpp

 namespace {
 class OMPEarlyOutliningPass
     : public fir::impl::OMPEarlyOutliningPassBase<OMPEarlyOutliningPass> {
 
   std::string getOutlinedFnName(llvm::StringRef parentName, unsigned count) {
     return std::string(parentName) + "_omp_outline_" + std::to_string(count);
   }
 
+  // Given a value this function will iterate over an operators results
+  // and return the relevant index for the result the value corresponds to.
+  // There may be a simpler way to do this however.
+  unsigned getResultIndex(mlir::Value value, mlir::Operation *op) {
+    for (unsigned i = 0; i < op->getNumResults(); ++i) {
+      if (op->getResult(i) == value)
+        return i;
+    }
+    return 0;
+  }
+
+  bool isDeclareTargetOp(mlir::Operation *op) {
+    if (fir::AddrOfOp addressOfOp = mlir::dyn_cast<fir::AddrOfOp>(op))
+      if (fir::GlobalOp gOp = mlir::dyn_cast<fir::GlobalOp>(
+              addressOfOp->getParentOfType<mlir::ModuleOp>().lookupSymbol(
+                  addressOfOp.getSymbol())))
+        if (auto declareTargetGlobal =
+                llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
+                    gOp.getOperation()))
+          if (declareTargetGlobal.isDeclareTarget())
+            return true;
+    return false;
+  }
+
+  // Currently used for cloning arguments that are nested. Should be
+  // extendable where required, perhaps via operation
+  // specialisation/overloading, if something needs specialised handling.
+  // NOTE: Results in duplication of some values that would otherwise be
+  // a single SSA value shared between operations, this is tidied up on
+  // lowering to some extent.
+  mlir::Operation *
+  cloneArgAndChildren(mlir::OpBuilder &builder, mlir::Operation *op,
+                      llvm::SetVector<mlir::Value> &inputs,
+                      mlir::Block::BlockArgListType &newInputs) {
+    mlir::IRMapping valueMap;
+    for (auto opValue : op->getOperands()) {
+      if (opValue.getDefiningOp()) {
+        auto resIdx = getResultIndex(opValue, opValue.getDefiningOp());
+        valueMap.map(opValue,
+                     cloneArgAndChildren(builder, opValue.getDefiningOp(),
+                                         inputs, newInputs)
+                         ->getResult(resIdx));
+      } else {
+        for (auto inArg : llvm::zip(inputs, newInputs)) {
+          if (opValue == std::get<0>(inArg))
+            valueMap.map(opValue, std::get<1>(inArg));
+        }
+      }
+    }
+
+    return builder.clone(*op, valueMap);
+  }
+
+  void cloneMapOpVariables(mlir::OpBuilder &builder, mlir::IRMapping &valueMap,
+                           mlir::IRMapping &mapInfoMap,
+                           llvm::SetVector<mlir::Value> &inputs,
+                           mlir::Block::BlockArgListType &newInputs,
+                           mlir::Value varPtr) {
+    if (fir::BoxAddrOp boxAddrOp =
+            mlir::dyn_cast_if_present<fir::BoxAddrOp>(varPtr.getDefiningOp())) {
+      mlir::Value newV =
+          cloneArgAndChildren(builder, boxAddrOp, inputs, newInputs)
+              ->getResult(0);
+      mapInfoMap.map(varPtr, newV);
+      valueMap.map(boxAddrOp, newV);
+      return;
+    }
+
+    if (varPtr.getDefiningOp() && isDeclareTargetOp(varPtr.getDefiningOp())) {
+      fir::AddrOfOp addrOp =
+          mlir::dyn_cast<fir::AddrOfOp>(varPtr.getDefiningOp());
+      mlir::Value newV = builder.clone(*addrOp)->getResult(0);
+      mapInfoMap.map(varPtr, newV);
+      valueMap.map(addrOp, newV);
+      return;
+    }
+
+    for (auto inArg : llvm::zip(inputs, newInputs)) {
+      if (varPtr == std::get<0>(inArg))
+        mapInfoMap.map(varPtr, std::get<1>(inArg));
+    }
+  }
+
   mlir::func::FuncOp outlineTargetOp(mlir::OpBuilder &builder,
                                      mlir::omp::TargetOp &targetOp,
                                      mlir::func::FuncOp &parentFunc,
                                      unsigned count) {
+    // NOTE: once implicit captures are handled appropriately in the initial
+    // PFT lowering if it is possible, we can remove the usage of
+    // getUsedValuesDefinedAbove and instead just iterate over the target op's
+    // operands (or just the map arguments) and perhaps refactor this function
+    // a little.
     // Collect inputs
     llvm::SetVector<mlir::Value> inputs;
-    for (auto operand : targetOp.getOperation()->getOperands())
-      inputs.insert(operand);
-
     mlir::Region &targetRegion = targetOp.getRegion();
     mlir::getUsedValuesDefinedAbove(targetRegion, inputs);
 
+    // filter out declareTarget and map entries which are specially handled
+    // at the moment, so we do not wish these to end up as function arguments
+    // which would just be more noise in the IR.
+    for (auto value : inputs)
+      if (value.getDefiningOp())
+        if (mlir::isa<mlir::omp::MapInfoOp>(value.getDefiningOp()) ||
+            isDeclareTargetOp(value.getDefiningOp()))
+          inputs.remove(value);
+
     // Create new function and initialize
     mlir::FunctionType funcType = builder.getFunctionType(
         mlir::TypeRange(inputs.getArrayRef()), mlir::TypeRange());
     std::string parentName(parentFunc.getName());
     std::string funcName = getOutlinedFnName(parentName, count);
     auto loc = targetOp.getLoc();
     mlir::func::FuncOp newFunc =
         mlir::func::FuncOp::create(loc, funcName, funcType);
     mlir::Block *entryBlock = newFunc.addEntryBlock();
     builder.setInsertionPointToStart(entryBlock);
-    mlir::ValueRange newInputs = entryBlock->getArguments();
+    mlir::Block::BlockArgListType newInputs = entryBlock->getArguments();
 
     // Set the declare target information, the outlined function
     // is always a host function.
     if (auto parentDTOp = llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
             parentFunc.getOperation()))
       if (auto newDTOp = llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
               newFunc.getOperation()))
         newDTOp.setDeclareTarget(mlir::omp::DeclareTargetDeviceType::host,
                                  parentDTOp.getDeclareTargetCaptureClause());
 
     // Set the early outlining interface parent name
     if (auto earlyOutlineOp =
             llvm::dyn_cast<mlir::omp::EarlyOutliningInterface>(
                 newFunc.getOperation()))
       earlyOutlineOp.setParentName(parentName);
 
-    // Create input map from inputs to function parameters.
+    // The value map for the newly generated Target Operation, we must
+    // remap most of the input.
     mlir::IRMapping valueMap;
-    for (auto InArg : llvm::zip(inputs, newInputs))
-      valueMap.map(std::get<0>(InArg), std::get<1>(InArg));
+
+    // Special handling for map, declare target and regular map variables
+    // are handled slightly differently for the moment, declare target has
+    // its addressOfOp cloned over, whereas we skip it for the regular map
+    // variables. We need knowledge of which global is linked to the map
+    // operation for declare target, whereas we aren't bothered for the
+    // regular map variables for the moment. We could treat both the same,
+    // however, cloning across the minimum for the moment to avoid
+    // optimisations breaking segments of the lowering seems prudent as this
+    // was the original intent of the pass.
+    for (auto oper : targetOp.getOperation()->getOperands()) {
+      if (auto mapEntry =
+              mlir::dyn_cast<mlir::omp::MapInfoOp>(oper.getDefiningOp())) {
+        mlir::IRMapping mapInfoMap;
+        for (auto bound : mapEntry.getBounds()) {
+          if (auto mapEntryBound = mlir::dyn_cast<mlir::omp::DataBoundsOp>(
+                  bound.getDefiningOp())) {
+            mapInfoMap.map(bound, cloneArgAndChildren(builder, mapEntryBound,
+                                                      inputs, newInputs)
+                                      ->getResult(0));
+          }
+        }
+
+        cloneMapOpVariables(builder, valueMap, mapInfoMap, inputs, newInputs,
+                            mapEntry.getVarPtr());
+
+        if (mapEntry.getVarPtrPtr())
+          cloneMapOpVariables(builder, valueMap, mapInfoMap, inputs, newInputs,
+                              mapEntry.getVarPtrPtr());
+
+        valueMap.map(
+            mapEntry,
+            builder.clone(*mapEntry.getOperation(), mapInfoMap)->getResult(0));
+      }
+    }
+
+    for (auto inArg : llvm::zip(inputs, newInputs))
+      valueMap.map(std::get<0>(inArg), std::get<1>(inArg));
 
     // Clone the target op into the new function
     builder.clone(*(targetOp.getOperation()), valueMap);
 
     // Create return op
     builder.create<mlir::func::ReturnOp>(loc);
 
     return newFunc;
   }
 
   // Returns true if a target region was found int the function.
   bool outlineTargetOps(mlir::OpBuilder &builder,
                         mlir::func::FuncOp &functionOp,
                         mlir::ModuleOp &moduleOp,
                         llvm::SmallVectorImpl<mlir::func::FuncOp> &newFuncs) {
     unsigned count = 0;
     for (auto TargetOp : functionOp.getOps<mlir::omp::TargetOp>()) {
       mlir::func::FuncOp outlinedFunc =
           outlineTargetOp(builder, TargetOp, functionOp, count);
       newFuncs.push_back(outlinedFunc);
       count++;
     }
     return count > 0;
   }
 
   void runOnOperation() override {
     mlir::ModuleOp moduleOp = getOperation();
     mlir::MLIRContext *context = &getContext();
     mlir::OpBuilder builder(context);
     llvm::SmallVector<mlir::func::FuncOp> newFuncs;