From cde7e805ad48e8766a07c95c2046693c6a6e236c Mon Sep 17 00:00:00 2001
From: Flakebi <flakebi@t-online.de>
Date: Mon, 10 Feb 2025 21:38:44 +0100
Subject: [PATCH] Cast allocas to default address space

Pointers for variables all need to be in the same address space for
correct compilation. Therefore ensure that even if an `alloca` is
created in a different address space, it is casted to the default
address space before its value is used.

This is necessary for the amdgpu target and others where the default
address space for `alloca`s is not 0.

For example the following code compiles incorrectly when not casting the
address space to the default one:

```rust
fn f(p: *const i8 /* addrspace(0) */) -> *const i8 /* addrspace(0) */ {
    let local = 0i8; /* addrspace(5) */
    let res = if cond { p } else { &raw const local };
    res
}
```

results in

```llvm
    %local = alloca addrspace(5) i8
    %res = alloca addrspace(5) ptr

if:
    ; Store 64-bit flat pointer
    store ptr %p, ptr addrspace(5) %res

else:
    ; Store 32-bit scratch pointer
    store ptr addrspace(5) %local, ptr addrspace(5) %res

ret:
    ; Load and return 64-bit flat pointer
    %res.load = load ptr, ptr addrspace(5) %res
    ret ptr %res.load
```

For amdgpu, `addrspace(0)` are 64-bit pointers, `addrspace(5)` are
32-bit pointers.
The above code may store a 32-bit pointer and read it back as a 64-bit
pointer, which is obviously wrong and cannot work. Instead, we need to
`addrspacecast %local to ptr addrspace(0)`, then we store and load the
correct type.
---
 compiler/rustc_codegen_llvm/src/builder.rs |  6 ++++--
 tests/codegen/amdgpu-addrspacecast.rs      | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)
 create mode 100644 tests/codegen/amdgpu-addrspacecast.rs

diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
index ab0e43e60a4f5..fbeedef314dfb 100644
--- a/compiler/rustc_codegen_llvm/src/builder.rs
+++ b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -421,7 +421,8 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         unsafe {
             let alloca = llvm::LLVMBuildAlloca(bx.llbuilder, ty, UNNAMED);
             llvm::LLVMSetAlignment(alloca, align.bytes() as c_uint);
-            alloca
+            // Cast to default addrspace if necessary
+            llvm::LLVMBuildPointerCast(bx.llbuilder, alloca, self.cx().type_ptr(), UNNAMED)
         }
     }
 
@@ -430,7 +431,8 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
             let alloca =
                 llvm::LLVMBuildArrayAlloca(self.llbuilder, self.cx().type_i8(), size, UNNAMED);
             llvm::LLVMSetAlignment(alloca, align.bytes() as c_uint);
-            alloca
+            // Cast to default addrspace if necessary
+            llvm::LLVMBuildPointerCast(self.llbuilder, alloca, self.cx().type_ptr(), UNNAMED)
         }
     }
 
diff --git a/tests/codegen/amdgpu-addrspacecast.rs b/tests/codegen/amdgpu-addrspacecast.rs
new file mode 100644
index 0000000000000..7fe630a7efa0e
--- /dev/null
+++ b/tests/codegen/amdgpu-addrspacecast.rs
@@ -0,0 +1,18 @@
+// Check that pointers are casted to addrspace(0) before they are used
+
+//@ compile-flags: --crate-type=rlib --target=amdgcn-amd-amdhsa -Ctarget-cpu=gfx900
+//@ needs-llvm-components: amdgpu
+//@ add-core-stubs
+#![feature(no_core)]
+#![no_core]
+
+extern crate minicore;
+
+// CHECK-LABEL: @ref_of_local
+// CHECK: [[alloca:%[0-9]]] = alloca
+// CHECK: %i = addrspacecast ptr addrspace(5) [[alloca]] to ptr
+#[no_mangle]
+pub fn ref_of_local(f: fn(&i32)) {
+    let i = 0;
+    f(&i);
+}