[OMPIRBuilder][OpenMP][LLVM] Modify and use ReplaceConstant utility in convertTarget (#94541)

This PR seeks to expand/replace the Constant -> Instruction conversion
that needs to occur inside of the OpenMP Target kernel generation to
allow kernel argument replacement of uses within the kernel (cannot
replace constant uses within constant expressions with non-constants).
It does so by making use of the new-ish utility
convertUsersOfConstantsToInstructions which is a much more expansive
version of what the smaller "version" of the function I wrote does,
effectively expanding uses of the input argument that are constant
expressions into instructions so that we can replace with the
appropriate kernel argument.

Also alters convertUsersOfConstantsToInstructions to optionally
restrict the replacement to a function and optionally leave
dead constants alone, the latter is necessary when lowering from 
MLIR as we cannot be sure we can remove the constants at this 
stage, even if rewritten to instructions the ModuleTranslation may 
maintain links to the original constants and utilise them in further 
lowering steps (as when we're lowering the kernel, the module is 
still in the process of being lowered). This can result in unusual 
ICEs later. These dead constants can be tidied up later (and 
appear to be in subsequent lowering from checking with 
emit-llvm).

Author: agozillon

llvm/include/llvm/IR/ReplaceConstant.h

@@ -18,10 +18,21 @@ namespace llvm {
 
 template <typename T> class ArrayRef;
 class Constant;
+class Function;
 
 /// Replace constant expressions users of the given constants with
 /// instructions. Return whether anything was changed.
-bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts);
+///
+/// Passing RestrictToFunc will restrict the constant replacement
+/// to the passed in functions scope, as opposed to the replacements
+/// occurring at module scope.
+///
+/// RemoveDeadConstants by default will remove all dead constants as
+/// the final step of the function after replacement, when passed
+/// false it will skip this final step.
+bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts,
+                                           Function *RestrictToFunc = nullptr,
+                                           bool RemoveDeadConstants = true);
 
 } // end namespace llvm
 

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -40,6 +40,7 @@
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/ReplaceConstant.h"
 #include "llvm/IR/Value.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/CommandLine.h"
@@ -5092,27 +5093,6 @@ FunctionCallee OpenMPIRBuilder::createDispatchFiniFunction(unsigned IVSize,
   return getOrCreateRuntimeFunction(M, Name);
 }
 
-static void replaceConstatExprUsesInFuncWithInstr(ConstantExpr *ConstExpr,
-                                                  Function *Func) {
-  for (User *User : make_early_inc_range(ConstExpr->users())) {
-    if (auto *Instr = dyn_cast<Instruction>(User)) {
-      if (Instr->getFunction() == Func) {
-        Instruction *ConstInst = ConstExpr->getAsInstruction();
-        ConstInst->insertBefore(*Instr->getParent(), Instr->getIterator());
-        Instr->replaceUsesOfWith(ConstExpr, ConstInst);
-      }
-    }
-  }
-}
-
-static void replaceConstantValueUsesInFuncWithInstr(llvm::Value *Input,
-                                                    Function *Func) {
-  for (User *User : make_early_inc_range(Input->users()))
-    if (auto *Const = dyn_cast<Constant>(User))
-      if (auto *ConstExpr = dyn_cast<ConstantExpr>(Const))
-        replaceConstatExprUsesInFuncWithInstr(ConstExpr, Func);
-}
-
 static Function *createOutlinedFunction(
     OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder, StringRef FuncName,
     SmallVectorImpl<Value *> &Inputs,
@@ -5191,17 +5171,23 @@ static Function *createOutlinedFunction(
 
     // Things like GEP's can come in the form of Constants. Constants and
     // ConstantExpr's do not have access to the knowledge of what they're
-    // contained in, so we must dig a little to find an instruction so we can
-    // tell if they're used inside of the function we're outlining. We also
-    // replace the original constant expression with a new instruction
+    // contained in, so we must dig a little to find an instruction so we
+    // can tell if they're used inside of the function we're outlining. We
+    // also replace the original constant expression with a new instruction
     // equivalent; an instruction as it allows easy modification in the
-    // following loop, as we can now know the constant (instruction) is owned by
-    // our target function and replaceUsesOfWith can now be invoked on it
-    // (cannot do this with constants it seems). A brand new one also allows us
-    // to be cautious as it is perhaps possible the old expression was used
-    // inside of the function but exists and is used externally (unlikely by the
-    // nature of a Constant, but still).
-    replaceConstantValueUsesInFuncWithInstr(Input, Func);
+    // following loop, as we can now know the constant (instruction) is
+    // owned by our target function and replaceUsesOfWith can now be invoked
+    // on it (cannot do this with constants it seems). A brand new one also
+    // allows us to be cautious as it is perhaps possible the old expression
+    // was used inside of the function but exists and is used externally
+    // (unlikely by the nature of a Constant, but still).
+    // NOTE: We cannot remove dead constants that have been rewritten to
+    // instructions at this stage, we run the risk of breaking later lowering
+    // by doing so as we could still be in the process of lowering the module
+    // from MLIR to LLVM-IR and the MLIR lowering may still require the original
+    // constants we have created rewritten versions of.
+    if (auto *Const = dyn_cast<Constant>(Input))
+      convertUsersOfConstantsToInstructions(Const, Func, false);
 
     // Collect all the instructions
     for (User *User : make_early_inc_range(Input->users()))

llvm/lib/IR/ReplaceConstant.cpp

@@ -49,7 +49,9 @@ static SmallVector<Instruction *, 4> expandUser(BasicBlock::iterator InsertPt,
   return NewInsts;
 }
 
-bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts) {
+bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts,
+                                           Function *RestrictToFunc,
+                                           bool RemoveDeadConstants) {
   // Find all expandable direct users of Consts.
   SmallVector<Constant *> Stack;
   for (Constant *C : Consts)
@@ -74,7 +76,8 @@ bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts) {
   for (Constant *C : ExpandableUsers)
     for (User *U : C->users())
       if (auto *I = dyn_cast<Instruction>(U))
-        InstructionWorklist.insert(I);
+        if (!RestrictToFunc || I->getFunction() == RestrictToFunc)
+          InstructionWorklist.insert(I);
 
   // Replace those expandable operands with instructions
   bool Changed = false;
@@ -102,8 +105,9 @@ bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts) {
     }
   }
 
-  for (Constant *C : Consts)
-    C->removeDeadConstantUsers();
+  if (RemoveDeadConstants)
+    for (Constant *C : Consts)
+      C->removeDeadConstantUsers();
 
   return Changed;
 }

offload/test/offloading/fortran/dtype-array-constant-index-map.f90

@@ -0,0 +1,51 @@
+! Offloading test which maps a specific element of a
+! derived type to the device and then accesses the
+! element alongside an individual element of an array
+! that the derived type contains. In particular, this
+! test helps to check that we can replace the constants
+! within the kernel with instructions and then replace
+! these instructions with the kernel parameters.
+! REQUIRES: flang
+! UNSUPPORTED: nvptx64-nvidia-cuda
+! UNSUPPORTED: nvptx64-nvidia-cuda-LTO
+! UNSUPPORTED: aarch64-unknown-linux-gnu
+! UNSUPPORTED: aarch64-unknown-linux-gnu-LTO
+! UNSUPPORTED: x86_64-pc-linux-gnu
+! UNSUPPORTED: x86_64-pc-linux-gnu-LTO
+
+! RUN: %libomptarget-compile-fortran-run-and-check-generic
+module test_0
+    type dtype
+      integer elements(20)
+      integer value
+    end type dtype
+
+    type (dtype) array_dtype(5)
+  contains
+
+  subroutine assign()
+    implicit none
+!$omp target map(tofrom: array_dtype(5))
+    array_dtype(5)%elements(5) = 500
+!$omp end target
+  end subroutine
+
+  subroutine add()
+    implicit none
+
+!$omp target map(tofrom: array_dtype(5))
+    array_dtype(5)%elements(5) = array_dtype(5)%elements(5) + 500
+!$omp end target
+  end subroutine
+end module test_0
+
+program main
+   use test_0
+
+  call assign()
+  call add()
+
+  print *, array_dtype(5)%elements(5)
+end program
+
+! CHECK: 1000