[flang-commits] [flang] [Flang][Transform] Modify stack reclaim pass to use allocation address space when generating intrinsics (PR #96836)

[via flang-commits]

https://github.com/agozillon created https://github.com/llvm/llvm-project/pull/96836

This PR aims to factor in the allocation address space provided by an architectures data layout when generating the intrinsic instructions, this allows them to be lowered later with the address spaces in tow. This aligns the intrinsic creation with the LLVM IRBuilder's https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/IRBuilder.h#L1053

This is also necessary for the below example to compile for OpenMP AMD GPU and not ICE the compiler in ISEL as AMD's stackrestore and stacksave are expected to have the appropriate allocation address space for AMD GPU.

program main
    integer(4), allocatable :: test
    allocate(test)

!$omp target map(tofrom:test)
    do i = 1, 10
      test = test + 50
    end do
!$omp end target

  deallocate(test)
end program

The PR also fixes the issue I opened a while ago which hits the same error when compiling for AMDGPU: https://github.com/llvm/llvm-project/issues/82368

Although, you have to have the appropriate GPU LIBC and Fortran offload runtime (both compiled for AMDGPU) added to the linker for the command or it will reach another ISEL error and ICE weirdly. But with the pre-requisites it works fine with this PR.

>From b28a90e0314faab704d0ff8417adcab6f542d78f Mon Sep 17 00:00:00 2001
From: agozillon <Andrew.Gozillon at amd.com>
Date: Wed, 26 Jun 2024 20:09:45 -0500
Subject: [PATCH] [Flang][Transform] Modify stack reclaim pass to use
 allocation address space when generating intrinsics

This PR aims to factor in the allocation address space provided by an architectures data layout when generating the intrinsic instructions, this allows them to be lowered later with the address spaces in tow. This aligns the intrinsic creation with the LLVM IRBuilder's https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/IRBuilder.h#L1053

This is also necessary for the below example to compile for OpenMP AMD GPU and not ICE the compiler in ISEL as AMD's stackrestore and stacksave are expected to have the appropriate allocation address space for AMD GPU.

program main
    integer(4), allocatable :: test
    allocate(test)

!$omp target map(tofrom:test)
    do i = 1, 10
      test = test + 50
    end do
!$omp end target

  deallocate(test)
end program

The PR also fixes the issue I opened a while ago which hits the same error when compiling for AMDGPU: https://github.com/llvm/llvm-project/issues/82368

Although, you have to have the appropriate GPU LIBC and Fortran offload runtime (both compiled for AMDGPU) added to the linker for the command or it will reach another ISEL error and ICE weirdly. But with the pre-requisites it works fine with this PR.
---
 flang/lib/Optimizer/Transforms/StackReclaim.cpp | 14 +++++++++++++-
 flang/test/Transforms/stack-reclaime.fir        | 17 +++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/flang/lib/Optimizer/Transforms/StackReclaim.cpp b/flang/lib/Optimizer/Transforms/StackReclaim.cpp
index e5e0e4eab8298..8a60a9e64f704 100644
--- a/flang/lib/Optimizer/Transforms/StackReclaim.cpp
+++ b/flang/lib/Optimizer/Transforms/StackReclaim.cpp
@@ -31,11 +31,23 @@ class StackReclaimPass : public fir::impl::StackReclaimBase<StackReclaimPass> {
 };
 } // namespace
 
+uint64_t getAllocaAddressSpace(Operation *op) {
+  mlir::ModuleOp module = mlir::dyn_cast_or_null<mlir::ModuleOp>(op);
+  if (!module)
+    module = op->getParentOfType<mlir::ModuleOp>();
+
+  if (mlir::Attribute addrSpace =
+          mlir::DataLayout(module).getAllocaMemorySpace())
+    return llvm::cast<mlir::IntegerAttr>(addrSpace).getUInt();
+  return 0;
+}
+
 void StackReclaimPass::runOnOperation() {
   auto *op = getOperation();
   auto *context = &getContext();
   mlir::OpBuilder builder(context);
-  mlir::Type voidPtr = mlir::LLVM::LLVMPointerType::get(context);
+  mlir::Type voidPtr =
+      mlir::LLVM::LLVMPointerType::get(context, getAllocaAddressSpace(op));
 
   op->walk([&](fir::DoLoopOp loopOp) {
     mlir::Location loc = loopOp.getLoc();
diff --git a/flang/test/Transforms/stack-reclaime.fir b/flang/test/Transforms/stack-reclaime.fir
index b53cc96035751..96416df8b2013 100644
--- a/flang/test/Transforms/stack-reclaime.fir
+++ b/flang/test/Transforms/stack-reclaime.fir
@@ -12,3 +12,20 @@ func.func @alloca_in_loop(%lb : index, %ub : index, %step : index, %b : i1, %add
 // CHECK: %[[STACKPTR:.*]] = llvm.intr.stacksave : !llvm.ptr
 // CHECK: %{{.*}} = fir.alloca !fir.box<!fir.heap<!fir.char<1,?>>>
 // CHECK: llvm.intr.stackrestore %0 : !llvm.ptr
+
+// -----
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui64>>} {
+  func.func @stack_restore_save_alloca_address(%lb : index, %ub : index, %step : index, %b : i1, %addr : !fir.ref<index>) {
+    fir.do_loop %iv = %lb to %ub step %step unordered {
+      %0 = fir.alloca !fir.box<!fir.heap<!fir.char<1,?>>>
+    }
+    return
+  }
+}
+
+// CHECK-LABEL: func.func @stack_restore_save_alloca_address
+// CHECK: fir.do_loop
+// CHECK: %[[STACKPTR:.*]] = llvm.intr.stacksave : !llvm.ptr<5>
+// CHECK: %{{.*}} = fir.alloca !fir.box<!fir.heap<!fir.char<1,?>>>
+// CHECK: llvm.intr.stackrestore %0 : !llvm.ptr<5>