From 99c4ae46a610a62699a6a59f66d21723734b911f Mon Sep 17 00:00:00 2001
From: Gabriel Baraldi <baraldigabriel@gmail.com>
Date: Sat, 30 Sep 2023 12:27:10 -0300
Subject: [PATCH] Make allocopt respect the GC verifier rules with non usual
 address spaces (#51520)

On AMDGPU, this was generating a `addrspace(10)` pointer to an `alloca`
which is illegal and lead to other issues.

(cherry picked from commit af9a7af3b27c0ff22179a7dcd30bf6753d3d575f)
---
 src/llvm-alloc-opt.cpp                          |  4 +---
 test/llvmpasses/alloc-opt-gcframe-addrspaces.ll | 11 ++++++-----
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/src/llvm-alloc-opt.cpp b/src/llvm-alloc-opt.cpp
index b87a5a6799b0b..751cae6419dda 100644
--- a/src/llvm-alloc-opt.cpp
+++ b/src/llvm-alloc-opt.cpp
@@ -643,8 +643,6 @@ void Optimizer::moveToStack(CallInst *orig_inst, size_t sz, bool has_ref)
     }
     insertLifetime(ptr, ConstantInt::get(Type::getInt64Ty(prolog_builder.getContext()), sz), orig_inst);
     Instruction *new_inst = cast<Instruction>(prolog_builder.CreateBitCast(ptr, JuliaType::get_pjlvalue_ty(prolog_builder.getContext(), buff->getType()->getPointerAddressSpace())));
-    if (orig_inst->getModule()->getDataLayout().getAllocaAddrSpace() != 0)
-        new_inst = cast<Instruction>(prolog_builder.CreateAddrSpaceCast(new_inst, JuliaType::get_pjlvalue_ty(prolog_builder.getContext(), orig_inst->getType()->getPointerAddressSpace())));
     new_inst->takeName(orig_inst);
 
     auto simple_replace = [&] (Instruction *orig_i, Instruction *new_i) {
@@ -692,7 +690,7 @@ void Optimizer::moveToStack(CallInst *orig_inst, size_t sz, bool has_ref)
         else if (auto call = dyn_cast<CallInst>(user)) {
             auto callee = call->getCalledOperand();
             if (pass.pointer_from_objref_func == callee) {
-                call->replaceAllUsesWith(new_i);
+                call->replaceAllUsesWith(prolog_builder.CreateAddrSpaceCast(new_i, call->getCalledFunction()->getReturnType()));
                 call->eraseFromParent();
                 return;
             }
diff --git a/test/llvmpasses/alloc-opt-gcframe-addrspaces.ll b/test/llvmpasses/alloc-opt-gcframe-addrspaces.ll
index b33f2cdac7dd4..1eefa7e56922f 100644
--- a/test/llvmpasses/alloc-opt-gcframe-addrspaces.ll
+++ b/test/llvmpasses/alloc-opt-gcframe-addrspaces.ll
@@ -19,15 +19,16 @@ declare {}* @julia.pointer_from_objref({} addrspace(11)*)
 ; Test that non-0 addrspace allocas are properly emitted and handled
 
 ; CHECK-LABEL: @non_zero_addrspace
-; CHECK: %1 = alloca i32, align 8, addrspace(5)
+; TYPED: %1 = alloca i32, align 8, addrspace(5)
 
 ; TYPED: %2 = bitcast i32 addrspace(5)* %1 to i8 addrspace(5)*
-; TYPED: %3 = bitcast i8 addrspace(5)* %2 to {} addrspace(5)*
-; TYPED: %var1 = addrspacecast {} addrspace(5)* %3 to {} addrspace(10)*
+; TYPED: %var1 = bitcast i8 addrspace(5)* %2 to {} addrspace(5)*
+; TYPED: %3 = addrspacecast {} addrspace(5)* %var1 to {}*
 ; TYPED: call void @llvm.lifetime.start.p5i8(i64 4, i8 addrspace(5)* %2)
 
-; OPAQUE: %var1 = addrspacecast ptr addrspace(5) %1 to ptr addrspace(10)
-; OPAQUE: call void @llvm.lifetime.start.p5(i64 4, ptr addrspace(5) %1)
+; OPAQUE: %var1 = alloca i32, align 8, addrspace(5)
+; OPAQUE: %1 = addrspacecast ptr addrspace(5) %var1 to ptr
+; OPAQUE: call void @llvm.lifetime.start.p5(i64 4, ptr addrspace(5) %var1)
 
 ; CHECK: ret void
 define void @non_zero_addrspace() {