[Flang][AMD] Skip setting alloc idx for allocate in target (llvm#2550)

This patch skips setting _FortranAAMDAllocatableSetAllocIdx for allocate
within "omp target" or for functions declared as "declare target".

Author: skc7

flang/lib/Lower/Allocatable.cpp

@@ -37,6 +37,7 @@
 #include "flang/Runtime/pointer.h"
 #include "flang/Semantics/tools.h"
 #include "flang/Semantics/type.h"
+#include "mlir/Dialect/OpenMP/OpenMPInterfaces.h"
 #include "llvm/Support/CommandLine.h"
 
 /// By default fir memory operation fir::AllocMemOp/fir::FreeMemOp are used.
@@ -174,11 +175,34 @@ static void genRuntimeInitCharacter(fir::FirOpBuilder &builder,
   builder.create<fir::CallOp>(loc, callee, convertedArgs);
 }
 
+/// Check if region is nested in omp.target or
+/// region nested in function with declare target
+bool isRegionNestedInOmpTarget(mlir::Region &region) {
+  mlir::Operation *parentOp = region.getParentOp();
+  while (parentOp) {
+    if (auto declareTargetOp =
+            llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(parentOp)) {
+      if (declareTargetOp.isDeclareTarget())
+        return true;
+    }
+    if (llvm::isa<mlir::omp::TargetOp>(parentOp))
+      return true;
+    mlir::Region *parentRegion = parentOp->getParentRegion();
+    if (!parentRegion)
+      break;
+    parentOp = parentRegion->getParentOp();
+  }
+
+  return false;
+}
+
 /// Generate a runtime call to set allocator idx of descriptor for target amd.
 static void genAMDRuntimeDescriptorSetAllocIdx(fir::FirOpBuilder &builder,
                                                mlir::Location loc,
                                                const fir::MutableBoxValue &box,
                                                int allocatorId) {
+  if (isRegionNestedInOmpTarget(builder.getRegion()))
+    return;
   auto *context = builder.getContext();
   mlir::Type descriptorTy = box.getAddr().getType();
   mlir::IntegerType posTy = builder.getI32Type();

flang/test/Lower/AMDGPU/allocate_deallocate_omp_declare_target.f90

@@ -0,0 +1,24 @@
+! RUN: %flang -ffast-amd-memory-allocator -S -emit-llvm -fopenmp -fopenmp-targets=amdgcn-amd-amdhsa --offload-arch=gfx90a -o - %s | FileCheck %s --check-prefix=CHECK-OMP
+! RUN: %flang -ffast-amd-memory-allocator -S -emit-llvm -target amdgcn-- -o - %s | FileCheck %s --check-prefix=CHECK
+
+subroutine func_t_device()
+  !$omp declare target enter(func_t_device) device_type(nohost)
+    integer, ALLOCATABLE :: poly
+
+! CHECK-OMP-NOT: call void @_FortranAAMDAllocatableSetAllocIdx({{.*}}, i32 1)
+! CHECK: call void @_FortranAAMDAllocatableSetAllocIdx({{.*}}, i32 1)
+! CHECK-OMP: call i32 @_FortranAAllocatableAllocate
+! CHECK: call i32 @_FortranAAllocatableAllocate   
+    ALLOCATE(poly)
+
+! CHECK-OMP: call i32 @_FortranAAllocatableDeallocate
+! CHECK: call i32 @_FortranAAllocatableDeallocate
+    DEALLOCATE(poly)
+end subroutine func_t_device
+
+program main
+  implicit none
+  !$omp target
+    call func_t_device()
+  !$omp end target
+end program

flang/test/Lower/AMDGPU/allocate_deallocate_omp_declare_target_nested.f90

@@ -0,0 +1,25 @@
+! RUN: %flang -ffast-amd-memory-allocator -S -emit-llvm -fopenmp -fopenmp-targets=amdgcn-amd-amdhsa --offload-arch=gfx90a -o - %s | FileCheck %s --check-prefix=CHECK-OMP
+! RUN: %flang -ffast-amd-memory-allocator -S -emit-llvm -target amdgcn-- -o - %s | FileCheck %s --check-prefix=CHECK
+
+subroutine func_t_device()
+  !$omp declare target enter(func_t_device) device_type(nohost)
+    integer, ALLOCATABLE :: poly
+    do j=1,10
+! CHECK-OMP-NOT: call void @_FortranAAMDAllocatableSetAllocIdx({{.*}}, i32 1)
+! CHECK: call void @_FortranAAMDAllocatableSetAllocIdx({{.*}}, i32 1)
+! CHECK-OMP: call i32 @_FortranAAllocatableAllocate
+! CHECK: call i32 @_FortranAAllocatableAllocate
+      ALLOCATE(poly)
+
+! CHECK-OMP: call i32 @_FortranAAllocatableDeallocate
+! CHECK: call i32 @_FortranAAllocatableDeallocate
+      DEALLOCATE(poly)
+    end do
+end subroutine func_t_device
+
+program main
+  implicit none
+  !$omp target
+    call func_t_device()
+  !$omp end target
+end program

flang/test/Lower/AMDGPU/allocate_deallocate_omp_target.f90

@@ -0,0 +1,24 @@
+! RUN: %flang -ffast-amd-memory-allocator -S -emit-llvm -fopenmp -fopenmp-targets=amdgcn-amd-amdhsa --offload-arch=gfx90a -o - %s | FileCheck %s --check-prefix=CHECK-OMP
+! RUN: %flang -ffast-amd-memory-allocator -S -emit-llvm -target amdgcn-- -o - %s | FileCheck %s --check-prefix=CHECK
+
+program main
+   implicit none
+   !$omp requires unified_shared_memory
+   REAL, DIMENSION(:), ALLOCATABLE :: poly
+   integer,parameter :: n = 10
+   integer :: i,j
+     !$omp target teams distribute parallel do private(poly)
+     do j=1,n
+
+! CHECK-OMP-NOT: call void @_FortranAAMDAllocatableSetAllocIdx({{.*}}, i32 1)
+! CHECK: call void @_FortranAAMDAllocatableSetAllocIdx({{.*}}, i32 1)
+! CHECK-OMP: call i32 @_FortranAAllocatableAllocate
+! CHECK: call i32 @_FortranAAllocatableAllocate
+       ALLOCATE(poly(1:3))
+       poly = 2.0_8
+! CHECK-OMP: call i32 @_FortranAAllocatableDeallocate
+! CHECK: call i32 @_FortranAAllocatableDeallocate
+       DEALLOCATE(poly)
+     enddo
+     !$omp end target teams distribute parallel do
+end program