Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 5 commits into from
Jun 13, 2024

Conversation

agozillon
Copy link
Contributor

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 (so far at least, it has only increased correctness e.g. test added in this PR), 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.

@llvmbot
Copy link
Collaborator

llvmbot commented Jun 5, 2024

@llvm/pr-subscribers-llvm-ir

@llvm/pr-subscribers-flang-openmp

Author: None (agozillon)

Changes

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 (so far at least, it has only increased correctness e.g. test added in this PR), 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.


Full diff: https://github.com/llvm/llvm-project/pull/94541.diff

5 Files Affected:

  • (modified) llvm/include/llvm/IR/ReplaceConstant.h (+2-1)
  • (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+17-31)
  • (modified) llvm/lib/IR/ReplaceConstant.cpp (+5-3)
  • (modified) mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir (+4-2)
  • (added) offload/test/offloading/fortran/dtype-array-constant-index-map.f90 (+51)
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 2c4b45255d059..78b5badad217c 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

@llvmbot
Copy link
Collaborator

llvmbot commented Jun 5, 2024

@llvm/pr-subscribers-mlir

Author: None (agozillon)

Changes

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 (so far at least, it has only increased correctness e.g. test added in this PR), 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.


Full diff: https://github.com/llvm/llvm-project/pull/94541.diff

5 Files Affected:

  • (modified) llvm/include/llvm/IR/ReplaceConstant.h (+2-1)
  • (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+17-31)
  • (modified) llvm/lib/IR/ReplaceConstant.cpp (+5-3)
  • (modified) mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir (+4-2)
  • (added) offload/test/offloading/fortran/dtype-array-constant-index-map.f90 (+51)
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 2c4b45255d059..78b5badad217c 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

@llvmbot
Copy link
Collaborator

llvmbot commented Jun 5, 2024

@llvm/pr-subscribers-offload

Author: None (agozillon)

Changes

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 (so far at least, it has only increased correctness e.g. test added in this PR), 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.


Full diff: https://github.com/llvm/llvm-project/pull/94541.diff

5 Files Affected:

  • (modified) llvm/include/llvm/IR/ReplaceConstant.h (+2-1)
  • (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+17-31)
  • (modified) llvm/lib/IR/ReplaceConstant.cpp (+5-3)
  • (modified) mlir/test/Target/LLVMIR/omptarget-fortran-allocatable-types-host.mlir (+4-2)
  • (added) offload/test/offloading/fortran/dtype-array-constant-index-map.f90 (+51)
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 2c4b45255d059..78b5badad217c 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

@agozillon
Copy link
Contributor Author

As a slight aside, this PR will have some affect on the following PR: #91829 as the common block PR makes use of the function I am replacing here (with the much better and more advanced/comprehensive utility function that now exists) to achieve it's results. So if/when this lands, I'll replicate these changes to the relevant locations in the other PR which should simplify the changeset as well.

@agozillon
Copy link
Contributor Author

Small ping for a reviewer or two for this PR, it'd be greatly appreciated!

@nikic
Copy link
Contributor

nikic commented Jun 12, 2024

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 (so far at least, it has only increased correctness e.g. test added in this PR), 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.

I'd feel more comfortable with this if it did not change other functions. It seems like supporting this in convertUsersOfConstantsToInstructions would just be one extra check when creating the InstructionWorklist, so maybe just go ahead and do it now?

@agozillon
Copy link
Contributor Author

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 (so far at least, it has only increased correctness e.g. test added in this PR), 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.

I'd feel more comfortable with this if it did not change other functions. It seems like supporting this in convertUsersOfConstantsToInstructions would just be one extra check when creating the InstructionWorklist, so maybe just go ahead and do it now?

Can do so! I was a bit hesitant to do it without some input as I wasn't sure exactly the best place to put the check to constrain it. But it should only take one extra check and an extra argument for the current function/kernel the conversion should be restricted to (at least from my simplistic understanding of the function)

…n 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.
@agozillon agozillon force-pushed the constant-to-instruction-expansion branch from c137271 to f3ce7c0 Compare June 12, 2024 20:07
Copy link
Contributor

@nikic nikic left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

// 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);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
convertUsersOfConstantsToInstructions({Const}, Func, false);
convertUsersOfConstantsToInstructions(Const, Func, false);

ArrayRef has a single-value overload.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, thank you, I wasn't aware! So used to brace initializing most things. I'll make this alteration and update in a little bit :-)

@agozillon
Copy link
Contributor Author

LGTM

Thank you very much! I'll pester one other reviewer for a review/sign-off and then hopefully it'll be good to land.

Copy link
Contributor

@jsjodin jsjodin left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just a nit, otherwise it looks good to me.


/// 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,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add descriptions of the new parameters to the function in the comment above.

@agozillon
Copy link
Contributor Author

Updated the comment in the last commit to provide details on the arguments! Thank you very much @jsjodin and @nikic I will merge the changes now.

@agozillon agozillon merged commit 0aeaa2d into llvm:main Jun 13, 2024
4 of 6 checks passed
nikic added a commit that referenced this pull request Jun 13, 2024
Reapplying without changes. The flang+openmp buildbot failure
should be addressed by #94541.

-----

This is a followup to #93823
and drops the DataLayout-unaware GEP of GEP fold entirely. All cases are
now left to the DataLayout-aware constant folder, which will fold
everything to a single i8 GEP.

We didn't have any test coverage for this fold in LLVM, but some Clang
tests change.
EthanLuisMcDonough pushed a commit to EthanLuisMcDonough/llvm-project that referenced this pull request Aug 13, 2024
…n convertTarget (llvm#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).
EthanLuisMcDonough pushed a commit to EthanLuisMcDonough/llvm-project that referenced this pull request Aug 13, 2024
Reapplying without changes. The flang+openmp buildbot failure
should be addressed by llvm#94541.

-----

This is a followup to llvm#93823
and drops the DataLayout-unaware GEP of GEP fold entirely. All cases are
now left to the DataLayout-aware constant folder, which will fold
everything to a single i8 GEP.

We didn't have any test coverage for this fold in LLVM, but some Clang
tests change.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants