[Mlir-commits] [llvm] [openmp] [mlir] [OpenMP][OMPIRBuilder] Handle replace uses of ConstantExpr's inside of Target regions (PR #71891)

[Jan Leyonberg]

================
@@ -4817,11 +4817,40 @@ static Function *createOutlinedFunction(
     Builder.restoreIP(
         ArgAccessorFuncCB(Arg, Input, InputCopy, AllocaIP, Builder.saveIP()));
 
-    // Collect all the instructions
+    // 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
+    // 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).
+    auto ReplaceConstantUsedInFunction = [](ConstantExpr *ConstExpr,
+                                            Function *Func) {
+      for (User *User : make_early_inc_range(ConstExpr->users()))
+        if (auto *Instr = dyn_cast<Instruction>(User))
+          if (Instr->getFunction() == Func)
+            Instr->replaceUsesOfWith(ConstExpr,
+                                     ConstExpr->getAsInstruction(Instr));
+    };
+
     for (User *User : make_early_inc_range(Input->users()))
-      if (auto Instr = dyn_cast<Instruction>(User))
-        if (Instr->getFunction() == Func)
+      if (auto *Const = dyn_cast<Constant>(User))
----------------
jsjodin wrote:

Replacing the ConstExprs could be useful elsewhere so make it a separate function.

https://github.com/llvm/llvm-project/pull/71891