Skip to content

[MLIR][OpenMP] Simplify OpenMP device codegen #137201

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.

Sign up for GitHub

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

Open
wants to merge 2 commits into
base: users/skatrak/map-rework-03-rm-host-ops
Choose a base branch
from
Open

[MLIR][OpenMP] Simplify OpenMP device codegen #137201

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9b5e759
[MLIR][OpenMP] Simplify OpenMP device codegen
skatrak Apr 22, 2025
a5261d9
Make omp.target[_{enter,exit}]_data and omp.target_update host-only ops
skatrak Aug 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 87 additions & 0 deletions flang/test/Integration/OpenMP/target-nesting-in-host-ops.f90
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
!===----------------------------------------------------------------------===!
! This directory can be used to add Integration tests involving multiple
! stages of the compiler (for eg. from Fortran to LLVM IR). It should not
! contain executable tests. We should only add tests here sparingly and only
! if there is no other way to test. Repeat this message in each test that is
! added to this directory and sub-directories.
!===----------------------------------------------------------------------===!

!REQUIRES: amdgpu-registered-target
!RUN: %flang_fc1 -triple amdgcn-amd-amdhsa -emit-llvm -fopenmp -fopenmp-version=50 -fopenmp-is-target-device %s -o - | FileCheck %s

! CHECK-NOT: define void @nested_target_in_parallel
! CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_nested_target_in_parallel_{{.*}}(ptr %{{.*}}, ptr %{{.*}})
subroutine nested_target_in_parallel(v)
implicit none
integer, intent(inout) :: v(10)

!$omp parallel
!$omp target map(tofrom: v)
!$omp end target
!$omp end parallel
end subroutine

! CHECK-NOT: define void @nested_target_in_wsloop
! CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_nested_target_in_wsloop_{{.*}}(ptr %{{.*}}, ptr %{{.*}})
subroutine nested_target_in_wsloop(v)
implicit none
integer, intent(inout) :: v(10)
integer :: i

!$omp do
do i=1, 10
!$omp target map(tofrom: v)
!$omp end target
end do
end subroutine

! CHECK-NOT: define void @nested_target_in_parallel_with_private
! CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_nested_target_in_parallel_with_private_{{.*}}(ptr %{{.*}}, ptr %{{.*}}, ptr %{{.*}})
subroutine nested_target_in_parallel_with_private(v)
implicit none
integer, intent(inout) :: v(10)
integer :: x
x = 10

!$omp parallel firstprivate(x)
!$omp target map(tofrom: v(1:x))
!$omp end target
!$omp end parallel
end subroutine

! CHECK-NOT: define void @nested_target_in_task_with_private
! CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_nested_target_in_task_with_private_{{.*}}(ptr %{{.*}}, ptr %{{.*}}, ptr %{{.*}})
subroutine nested_target_in_task_with_private(v)
implicit none
integer, intent(inout) :: v(10)
integer :: x
x = 10

!$omp task firstprivate(x)
!$omp target map(tofrom: v(1:x))
!$omp end target
!$omp end task
end subroutine

! CHECK-NOT: define void @target_and_atomic_update
! CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_target_and_atomic_update_{{.*}}(ptr %{{.*}})
subroutine target_and_atomic_update(x, expr)
implicit none
integer, intent(inout) :: x, expr

!$omp target
!$omp end target

!$omp atomic update
x = x + expr
end subroutine

! CHECK-NOT: define void @nested_target_in_associate
! CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_nested_target_in_associate_{{.*}}(ptr %{{.*}}, ptr %{{.*}}, ptr %{{.*}})
subroutine nested_target_in_associate(x)
integer, pointer, contiguous :: x(:)
associate(y => x)
!$omp target map(tofrom: y)
!$omp end target
end associate
end subroutine
46 changes: 46 additions & 0 deletions flang/test/Integration/OpenMP/target-use-device-nested.f90
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
!===----------------------------------------------------------------------===!
! This directory can be used to add Integration tests involving multiple
! stages of the compiler (for eg. from Fortran to LLVM IR). It should not
! contain executable tests. We should only add tests here sparingly and only
! if there is no other way to test. Repeat this message in each test that is
! added to this directory and sub-directories.
!===----------------------------------------------------------------------===!

! This tests check that target code nested inside a target data region which
! has only use_device_ptr mapping corectly generates code on the device pass.

!REQUIRES: amdgpu-registered-target
!RUN: %flang_fc1 -triple amdgcn-amd-amdhsa -emit-llvm -fopenmp -fopenmp-version=50 -fopenmp-is-target-device %s -o - | FileCheck %s

program main
use iso_c_binding
implicit none
type(c_ptr) :: a
!$omp target data use_device_ptr(a)
!$omp target map(tofrom: a)
call foo(a)
!$omp end target
!$omp end target data
end program

! CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading{{.*}}main_
! CHECK-NEXT: entry:
! CHECK-NEXT: %[[VAL_3:.*]] = alloca ptr, align 8, addrspace(5)
! CHECK-NEXT: %[[ASCAST:.*]] = addrspacecast ptr addrspace(5) %[[VAL_3]] to ptr
! CHECK-NEXT: store ptr %[[VAL_4:.*]], ptr %[[ASCAST]], align 8
! CHECK-NEXT: %[[VAL_5:.*]] = call i32 @__kmpc_target_init(ptr addrspacecast (ptr addrspace(1) @__omp_offloading_{{.*}}_kernel_environment to ptr), ptr %[[VAL_6:.*]])
! CHECK-NEXT: %[[VAL_7:.*]] = icmp eq i32 %[[VAL_5]], -1
! CHECK-NEXT: br i1 %[[VAL_7]], label %[[VAL_8:.*]], label %[[VAL_9:.*]]
! CHECK: user_code.entry: ; preds = %[[VAL_10:.*]]
! CHECK-NEXT: %[[VAL_11:.*]] = load ptr, ptr %[[ASCAST]], align 8
! CHECK-NEXT: br label %[[AFTER_ALLOC:.*]]

! CHECK: [[AFTER_ALLOC]]:
! CHECK-NEXT: br label %[[VAL_12:.*]]

! CHECK: [[VAL_12]]:
! CHECK-NEXT: br label %[[TARGET_REG_ENTRY:.*]]

! CHECK: [[TARGET_REG_ENTRY]]: ; preds = %[[VAL_12]]
! CHECK-NEXT: call void @{{.*}}foo{{.*}}(ptr %[[VAL_11]])
! CHECK-NEXT: br label
37 changes: 37 additions & 0 deletions flang/test/Integration/OpenMP/task-target-device.f90
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
!===----------------------------------------------------------------------===!
! This directory can be used to add Integration tests involving multiple
! stages of the compiler (for eg. from Fortran to LLVM IR). It should not
! contain executable tests. We should only add tests here sparingly and only
! if there is no other way to test. Repeat this message in each test that is
! added to this directory and sub-directories.
!===----------------------------------------------------------------------===!

!REQUIRES: amdgpu-registered-target
!RUN: %flang_fc1 -triple amdgcn-amd-amdhsa -emit-llvm -fopenmp -fopenmp-version=50 -fopenmp-is-target-device %s -o - | FileCheck %s

! This tests the fix for https://github.com/llvm/llvm-project/issues/84606
! We are only interested in ensuring that the -mlir-to-llmvir pass doesn't crash.

! CHECK: define weak_odr protected amdgpu_kernel void @{{.*}}QQmain{{.*}}({{.*}})
program main
implicit none
integer, parameter :: N = 5
integer, dimension(5) :: a
integer :: i
integer :: target_a = 0

!$omp task depend(out:a)
do i = 1, N
a(i) = i
end do
!$omp end task

!$omp target map(tofrom:target_a) map(tofrom:a)
do i = 1, N
target_a = target_a + i
a(i) = a(i) + i
end do
!$omp end target
print*, target_a
print*, a
end program main
40 changes: 40 additions & 0 deletions flang/test/Integration/OpenMP/threadprivate-target-device.f90
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
!===----------------------------------------------------------------------===!
! This directory can be used to add Integration tests involving multiple
! stages of the compiler (for eg. from Fortran to LLVM IR). It should not
! contain executable tests. We should only add tests here sparingly and only
! if there is no other way to test. Repeat this message in each test that is
! added to this directory and sub-directories.
!===----------------------------------------------------------------------===!

!REQUIRES: amdgpu-registered-target
!RUN: %flang_fc1 -triple amdgcn-amd-amdhsa -emit-llvm -fopenmp -fopenmp-version=50 -fopenmp-is-target-device %s -o - | FileCheck %s

! The aim of this test is to verify host threadprivate directives do not cause
! crashes during OpenMP target device codegen when used in conjunction with
! target code in the same function.

! CHECK: define weak_odr protected amdgpu_kernel void @{{.*}}(ptr %{{.*}}, ptr %[[ARG1:.*]], ptr %[[ARG2:.*]]) #{{[0-9]+}} {
! CHECK: %[[ALLOCA_X:.*]] = alloca ptr, align 8, addrspace(5)
! CHECK: %[[ASCAST_X:.*]] = addrspacecast ptr addrspace(5) %[[ALLOCA_X]] to ptr
! CHECK: store ptr %[[ARG1]], ptr %[[ASCAST_X]], align 8

! CHECK: %[[ALLOCA_N:.*]] = alloca ptr, align 8, addrspace(5)
! CHECK: %[[ASCAST_N:.*]] = addrspacecast ptr addrspace(5) %[[ALLOCA_N]] to ptr
! CHECK: store ptr %[[ARG2]], ptr %[[ASCAST_N]], align 8

! CHECK: %[[LOAD_X:.*]] = load ptr, ptr %[[ASCAST_X]], align 8
! CHECK: call void @bar_(ptr %[[LOAD_X]], ptr %[[ASCAST_N]])

module test
implicit none
integer :: n
!$omp threadprivate(n)

contains
subroutine foo(x)
integer, intent(inout) :: x(10)
!$omp target map(tofrom: x(1:n))
call bar(x, n)
!$omp end target
end subroutine
end module
11 changes: 0 additions & 11 deletions llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7080,17 +7080,6 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTargetData(
return InsertPointTy();

Builder.restoreIP(CodeGenIP);
// Disable TargetData CodeGen on Device pass.
if (Config.IsTargetDevice.value_or(false)) {
if (BodyGenCB) {
InsertPointOrErrorTy AfterIP =
BodyGenCB(Builder.saveIP(), BodyGenTy::NoPriv);
if (!AfterIP)
return AfterIP.takeError();
Builder.restoreIP(*AfterIP);
}
return Builder.saveIP();
}

bool IsStandAlone = !BodyGenCB;
MapInfosTy *MapInfo;
Expand Down
21 changes: 0 additions & 21 deletions llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6281,27 +6281,6 @@ TEST_F(OpenMPIRBuilderTest, TargetDataRegion) {
EXPECT_TRUE(TargetDataCall->getOperand(2)->getType()->isIntegerTy(32));
EXPECT_TRUE(TargetDataCall->getOperand(8)->getType()->isPointerTy());

// Check that BodyGenCB is still made when IsTargetDevice is set to true.
OMPBuilder.Config.setIsTargetDevice(true);
bool CheckDevicePassBodyGen = false;
auto BodyTargetCB = [&](InsertPointTy CodeGenIP, BodyGenTy BodyGenType) {
CheckDevicePassBodyGen = true;
Builder.restoreIP(CodeGenIP);
CallInst *TargetDataCall =
dyn_cast<CallInst>(BB->back().getPrevNode()->getPrevNode());
// Make sure no begin_mapper call is present for device pass.
EXPECT_EQ(TargetDataCall, nullptr);
return Builder.saveIP();
};
ASSERT_EXPECTED_INIT(
OpenMPIRBuilder::InsertPointTy, TargetDataIP2,
OMPBuilder.createTargetData(Loc, AllocaIP, Builder.saveIP(),
Builder.getInt64(DeviceID),
/* IfCond= */ nullptr, Info, GenMapInfoCB,
CustomMapperCB, nullptr, BodyTargetCB));
Builder.restoreIP(TargetDataIP2);
EXPECT_TRUE(CheckDevicePassBodyGen);

Builder.CreateRetVoid();
EXPECT_FALSE(verifyModule(*M, &errs()));
}
Expand Down
Loading
Loading