[Mlir-commits] [llvm] [mlir] [OMPIRBuilder][OpenMP][LLVM] Modify and use ReplaceConstant utility in convertTarget (PR #94541)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 12 13:41:08 PDT 2024
https://github.com/agozillon updated https://github.com/llvm/llvm-project/pull/94541
>From 476c9ba4b08be2ac42394dad591cbd61cd870c52 Mon Sep 17 00:00:00 2001
From: agozillon <Andrew.Gozillon at amd.com>
Date: Wed, 5 Jun 2024 17:45:38 -0500
Subject: [PATCH 1/4] [OMPIRBuilder][OpenMP][LLVM] Modify and use
ReplaceConstant utility in convertTarget
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 leave dead constants
alone 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).
The one possible downside to this replacement is that the constant -> instruction
rewriting is no longer constrained to within the kernel, it will expand the available
uses of an input argument that is constant and has constant uses in the module.
This hasn't lowered the correctness of the examples I have tested with, however,
it may impact performance, a possibility in the future may be to optionally
constrain rewrites of uses of constants in convertUsersOfConstantsToInstructions
to a provided llvm::Function.
---
llvm/include/llvm/IR/ReplaceConstant.h | 3 +-
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 48 +++++++----------
llvm/lib/IR/ReplaceConstant.cpp | 8 +--
...target-fortran-allocatable-types-host.mlir | 6 ++-
.../dtype-array-constant-index-map.f90 | 51 +++++++++++++++++++
5 files changed, 79 insertions(+), 37 deletions(-)
create mode 100644 offload/test/offloading/fortran/dtype-array-constant-index-map.f90
diff --git a/llvm/include/llvm/IR/ReplaceConstant.h b/llvm/include/llvm/IR/ReplaceConstant.h
index 56027d7fba6f8..8f58548e3755a 100644
--- a/llvm/include/llvm/IR/ReplaceConstant.h
+++ b/llvm/include/llvm/IR/ReplaceConstant.h
@@ -21,7 +21,8 @@ class Constant;
/// Replace constant expressions users of the given constants with
/// instructions. Return whether anything was changed.
-bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts);
+bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts,
+ bool RemoveDeadConstants = true);
} // end namespace llvm
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 92213e19c9d9d..1d398426f51e0 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/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}, false);
// Collect all the instructions
for (User *User : make_early_inc_range(Input->users()))
diff --git a/llvm/lib/IR/ReplaceConstant.cpp b/llvm/lib/IR/ReplaceConstant.cpp
index 9b07bd8040492..e400be8884a74 100644
--- a/llvm/lib/IR/ReplaceConstant.cpp
+++ b/llvm/lib/IR/ReplaceConstant.cpp
@@ -49,7 +49,8 @@ static SmallVector<Instruction *, 4> expandUser(BasicBlock::iterator InsertPt,
return NewInsts;
}
-bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts) {
+bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts,
+ bool RemoveDeadConstants) {
// Find all expandable direct users of Consts.
SmallVector<Constant *> Stack;
for (Constant *C : Consts)
@@ -102,8 +103,9 @@ bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts) {
}
}
- for (Constant *C : Consts)
- C->removeDeadConstantUsers();
+ if (RemoveDeadConstants)
+ for (Constant *C : Consts)
+ C->removeDeadConstantUsers();
return Changed;
}
diff --git a/mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir b/mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir
index 9b46f84e5050f..7845012e1de6f 100644
--- a/mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir
+++ b/mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir
@@ -66,9 +66,11 @@ module attributes {omp.is_target_device = false} {
// CHECK: define void @_QQmain()
// CHECK: %[[SCALAR_ALLOCA:.*]] = alloca { ptr, i64, i32, i8, i8, i8, i8 }, i64 1, align 8
-// CHECK: %[[FULL_ARR_SIZE5:.*]] = load i64, ptr getelementptr ({ ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr @[[FULL_ARR_GLOB]], i32 0, i32 7, i64 0, i32 1), align 4
+// CHECK: %[[FULL_ARR_GEP:.*]] = getelementptr { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr @[[FULL_ARR_GLOB]], i32 0, i32 7, i64 0, i32 1
+// CHECK: %[[FULL_ARR_SIZE5:.*]] = load i64, ptr %[[FULL_ARR_GEP]], align 4
// CHECK: %[[FULL_ARR_SIZE4:.*]] = sub i64 %[[FULL_ARR_SIZE5]], 1
-// CHECK: %[[ARR_SECT_OFFSET3:.*]] = load i64, ptr getelementptr ({ ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr @[[ARR_SECT_GLOB]], i32 0, i32 7, i64 0, i32 0), align 4
+// CHECK: %[[ARR_SECT_GEP:.*]] = getelementptr { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr @[[ARR_SECT_GLOB]], i32 0, i32 7, i64 0, i32 0
+// CHECK: %[[ARR_SECT_OFFSET3:.*]] = load i64, ptr %[[ARR_SECT_GEP]], align 4
// CHECK: %[[ARR_SECT_OFFSET2:.*]] = sub i64 2, %[[ARR_SECT_OFFSET3]]
// CHECK: %[[ARR_SECT_SIZE4:.*]] = sub i64 5, %[[ARR_SECT_OFFSET3]]
// CHECK: %[[SCALAR_BASE:.*]] = getelementptr { ptr, i64, i32, i8, i8, i8, i8 }, ptr %[[SCALAR_ALLOCA]], i32 0, i32 0
diff --git a/offload/test/offloading/fortran/dtype-array-constant-index-map.f90 b/offload/test/offloading/fortran/dtype-array-constant-index-map.f90
new file mode 100644
index 0000000000000..580eef42e51a8
--- /dev/null
+++ b/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
\ No newline at end of file
>From f3ce7c0f397a3f1910e456ce8e2765c15aa0a306 Mon Sep 17 00:00:00 2001
From: agozillon <Andrew.Gozillon at amd.com>
Date: Wed, 12 Jun 2024 14:38:53 -0500
Subject: [PATCH 2/4] restrict constant -> instruction upgrade to kernel/a
specified function
---
llvm/include/llvm/IR/ReplaceConstant.h | 2 ++
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 2 +-
llvm/lib/IR/ReplaceConstant.cpp | 4 +++-
3 files changed, 6 insertions(+), 2 deletions(-)
diff --git a/llvm/include/llvm/IR/ReplaceConstant.h b/llvm/include/llvm/IR/ReplaceConstant.h
index 8f58548e3755a..5826fefcafe35 100644
--- a/llvm/include/llvm/IR/ReplaceConstant.h
+++ b/llvm/include/llvm/IR/ReplaceConstant.h
@@ -18,10 +18,12 @@ 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,
+ Function *RestrictToFunc = nullptr,
bool RemoveDeadConstants = true);
} // end namespace llvm
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 1d398426f51e0..3e7966736b1de 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -5187,7 +5187,7 @@ static Function *createOutlinedFunction(
// 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}, false);
+ convertUsersOfConstantsToInstructions({Const}, Func, false);
// Collect all the instructions
for (User *User : make_early_inc_range(Input->users()))
diff --git a/llvm/lib/IR/ReplaceConstant.cpp b/llvm/lib/IR/ReplaceConstant.cpp
index e400be8884a74..67b6fe6fda3b9 100644
--- a/llvm/lib/IR/ReplaceConstant.cpp
+++ b/llvm/lib/IR/ReplaceConstant.cpp
@@ -50,6 +50,7 @@ static SmallVector<Instruction *, 4> expandUser(BasicBlock::iterator InsertPt,
}
bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts,
+ Function *RestrictToFunc,
bool RemoveDeadConstants) {
// Find all expandable direct users of Consts.
SmallVector<Constant *> Stack;
@@ -75,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;
>From e7807839f75c6d0297d838bb02b9d3f07282c5d1 Mon Sep 17 00:00:00 2001
From: agozillon <Andrew.Gozillon at amd.com>
Date: Wed, 12 Jun 2024 15:27:36 -0500
Subject: [PATCH 3/4] Remove unnecessary brace initialization
---
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 3e7966736b1de..ccec66fcb7bac 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -5187,7 +5187,7 @@ static Function *createOutlinedFunction(
// 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);
+ convertUsersOfConstantsToInstructions(Const, Func, false);
// Collect all the instructions
for (User *User : make_early_inc_range(Input->users()))
>From 996b763900c6512eece2732aeda02a4a8de885b0 Mon Sep 17 00:00:00 2001
From: agozillon <Andrew.Gozillon at amd.com>
Date: Wed, 12 Jun 2024 15:40:49 -0500
Subject: [PATCH 4/4] Revert no longer necessary alterations to MLIR test now
that constant -> instruction replacement has been restricted
---
.../LLVMIR/omptarget-fortran-allocatable-types-host.mlir | 6 ++----
1 file changed, 2 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir b/mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir
index 7845012e1de6f..9b46f84e5050f 100644
--- a/mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir
+++ b/mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir
@@ -66,11 +66,9 @@ module attributes {omp.is_target_device = false} {
// CHECK: define void @_QQmain()
// CHECK: %[[SCALAR_ALLOCA:.*]] = alloca { ptr, i64, i32, i8, i8, i8, i8 }, i64 1, align 8
-// CHECK: %[[FULL_ARR_GEP:.*]] = getelementptr { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr @[[FULL_ARR_GLOB]], i32 0, i32 7, i64 0, i32 1
-// CHECK: %[[FULL_ARR_SIZE5:.*]] = load i64, ptr %[[FULL_ARR_GEP]], align 4
+// CHECK: %[[FULL_ARR_SIZE5:.*]] = load i64, ptr getelementptr ({ ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr @[[FULL_ARR_GLOB]], i32 0, i32 7, i64 0, i32 1), align 4
// CHECK: %[[FULL_ARR_SIZE4:.*]] = sub i64 %[[FULL_ARR_SIZE5]], 1
-// CHECK: %[[ARR_SECT_GEP:.*]] = getelementptr { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr @[[ARR_SECT_GLOB]], i32 0, i32 7, i64 0, i32 0
-// CHECK: %[[ARR_SECT_OFFSET3:.*]] = load i64, ptr %[[ARR_SECT_GEP]], align 4
+// CHECK: %[[ARR_SECT_OFFSET3:.*]] = load i64, ptr getelementptr ({ ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr @[[ARR_SECT_GLOB]], i32 0, i32 7, i64 0, i32 0), align 4
// CHECK: %[[ARR_SECT_OFFSET2:.*]] = sub i64 2, %[[ARR_SECT_OFFSET3]]
// CHECK: %[[ARR_SECT_SIZE4:.*]] = sub i64 5, %[[ARR_SECT_OFFSET3]]
// CHECK: %[[SCALAR_BASE:.*]] = getelementptr { ptr, i64, i32, i8, i8, i8, i8 }, ptr %[[SCALAR_ALLOCA]], i32 0, i32 0
More information about the Mlir-commits
mailing list