[clang] [CIR] Fix address space of dest cleanup slot (PR #222628)

[via cfe-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-clang

Author: Mariya Podchishchaeva (Fznamznon)

<details>
<summary>Changes</summary>

Make sure that flatten cfg creates new allocas for dest cleanup slot in correct address space. This is important for address-space aware targets like amdgpu.

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


3 Files Affected:

- (modified) clang/include/clang/CIR/Dialect/IR/CIRDataLayout.h (+9) 
- (modified) clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp (+2-1) 
- (added) clang/test/CIR/CodeGenHIP/cleanup-alloca-addrspace.hip (+39) 


``````````diff
diff --git a/clang/include/clang/CIR/Dialect/IR/CIRDataLayout.h b/clang/include/clang/CIR/Dialect/IR/CIRDataLayout.h
index 17b339417ad5a..71214c798afbf 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIRDataLayout.h
+++ b/clang/include/clang/CIR/Dialect/IR/CIRDataLayout.h
@@ -120,6 +120,15 @@ class CIRDataLayout {
   bool typeSizeEqualsStoreSize(mlir::Type ty) const {
     return getTypeSizeInBits(ty) == getTypeStoreSizeInBits(ty);
   }
+
+  mlir::ptr::MemorySpaceAttrInterface
+  getAllocaAddrSpace(mlir::MLIRContext *ctx) {
+    auto allocaASAttr = mlir::dyn_cast_if_present<mlir::IntegerAttr>(
+        layout.getAllocaMemorySpace());
+    if (!allocaASAttr)
+      return {};
+    return cir::TargetAddressSpaceAttr::get(ctx, allocaASAttr.getUInt());
+  }
 };
 
 } // namespace cir
diff --git a/clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp b/clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp
index eda0b59cc2367..a074aa70898e9 100644
--- a/clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp
+++ b/clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp
@@ -830,8 +830,9 @@ static cir::AllocaOp getOrCreateCleanupDestSlot(cir::FuncOp funcOp,
   rewriter.setInsertionPointToStart(&entryBlock);
   cir::IntType s32Type =
       cir::IntType::get(rewriter.getContext(), 32, /*isSigned=*/true);
-  cir::PointerType ptrToS32Type = cir::PointerType::get(s32Type);
   cir::CIRDataLayout dataLayout(funcOp->getParentOfType<mlir::ModuleOp>());
+  cir::PointerType ptrToS32Type = cir::PointerType::get(
+      s32Type, dataLayout.getAllocaAddrSpace(rewriter.getContext()));
   uint64_t alignment = dataLayout.getAlignment(s32Type, true).value();
   auto allocaOp = cir::AllocaOp::create(
       rewriter, loc, ptrToS32Type, "__cleanup_dest_slot",
diff --git a/clang/test/CIR/CodeGenHIP/cleanup-alloca-addrspace.hip b/clang/test/CIR/CodeGenHIP/cleanup-alloca-addrspace.hip
new file mode 100644
index 0000000000000..83b05227f0845
--- /dev/null
+++ b/clang/test/CIR/CodeGenHIP/cleanup-alloca-addrspace.hip
@@ -0,0 +1,39 @@
+// RUN: %clang_cc1 -triple amdgcn-amd-amdhsa -x hip -fclangir -fcuda-is-device -emit-cir %s -o %t.cir
+// RUN: cir-opt --cir-flatten-cfg %t.cir -o %t-flat.cir
+// RUN: FileCheck --input-file=%t-flat.cir %s --check-prefix=CIR-FLAT
+// RUN: %clang_cc1 -triple amdgcn-amd-amdhsa -x hip -fclangir -fcuda-is-device -emit-llvm %s -o - | FileCheck %s --check-prefix=LLVM
+
+// Check that alloca created for cleanup destination slot for a function has
+// correct address space after flatten cfg. This is important for address-space
+// aware targets like amdgpu.
+
+// CIR-FLAT-LABEL: cir.func {{.*}} @_Z1fv
+// CIR-FLAT:         cir.alloca "__cleanup_dest_slot" {{.*}} : !cir.ptr<!s32i, target_address_space(5)>
+// CIR-FLAT:         cir.alloca "__retval" {{.*}} : !cir.ptr<!rec_S, target_address_space(5)>
+// CIR-FLAT:         cir.alloca "nrvo" {{.*}} : !cir.ptr<!cir.bool, target_address_space(5)>
+
+// CIR-FLAT-LABEL: cir.func {{.*}} @_ZN1SD1Ev
+// CIR-FLAT:         cir.alloca "this" {{.*}} : !cir.ptr<!cir.ptr<!rec_S>, target_address_space(5)>
+
+// CIR-FLAT-LABEL: cir.func {{.*}} @_ZN1SD2Ev
+// CIR-FLAT:         cir.alloca "this" {{.*}} : !cir.ptr<!cir.ptr<!rec_S>, target_address_space(5)>
+
+// LLVM-LABEL: define {{.*}} @_Z1fv
+// LLVM:         alloca i32, align 4, addrspace(5)
+// LLVM:         alloca %struct.S, align 1, addrspace(5)
+// LLVM:         alloca i8, align 1, addrspace(5)
+
+// LLVM-LABEL: define {{.*}} @_ZN1SD1Ev
+// LLVM:         alloca ptr, align 8, addrspace(5)
+
+// LLVM-LABEL: define {{.*}} @_ZN1SD2Ev
+// LLVM:         alloca ptr, align 8, addrspace(5)
+
+struct S {
+  __attribute__((device)) ~S() {}
+};
+
+__attribute__((device)) S f() {
+  S s;
+  return s;
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/222628