[OMPIRBuilder][OpenMP][LLVM] Modify and use ReplaceConstant utility in convertTarget

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