[llvm] 597b9b7 - CodeExtractor: Fix assertion with non-0 alloca address spaces

[Matt Arsenault via llvm-commits]

Author: Matt Arsenault
Date: 2022-10-23T15:16:55-07:00
New Revision: 597b9b7e8eaae1fbbab4f2603f1fe99c02f0ff5e

URL: https://github.com/llvm/llvm-project/commit/597b9b7e8eaae1fbbab4f2603f1fe99c02f0ff5e
DIFF: https://github.com/llvm/llvm-project/commit/597b9b7e8eaae1fbbab4f2603f1fe99c02f0ff5e.diff

LOG: CodeExtractor: Fix assertion with non-0 alloca address spaces

emitCallAndSwitchStatement creates placeholder allocas to pass
to these, so the types need to match.

Added: 
    llvm/test/tools/llvm-extract/address-space-assert.ll

Modified: 
    llvm/lib/Transforms/Utils/CodeExtractor.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index 5c8160e6c3939..23b943c69386d 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -831,6 +831,7 @@ Function *CodeExtractor::constructFunction(const ValueSet &inputs,
   std::vector<Type *> ParamTy;
   std::vector<Type *> AggParamTy;
   ValueSet StructValues;
+  const DataLayout &DL = M->getDataLayout();
 
   // Add the types of the input values to the function's argument list
   for (Value *value : inputs) {
@@ -849,7 +850,8 @@ Function *CodeExtractor::constructFunction(const ValueSet &inputs,
       AggParamTy.push_back(output->getType());
       StructValues.insert(output);
     } else
-      ParamTy.push_back(PointerType::getUnqual(output->getType()));
+      ParamTy.push_back(
+          PointerType::get(output->getType(), DL.getAllocaAddrSpace()));
   }
 
   assert(
@@ -864,7 +866,7 @@ Function *CodeExtractor::constructFunction(const ValueSet &inputs,
   StructType *StructTy = nullptr;
   if (AggregateArgs && !AggParamTy.empty()) {
     StructTy = StructType::get(M->getContext(), AggParamTy);
-    ParamTy.push_back(PointerType::getUnqual(StructTy));
+    ParamTy.push_back(PointerType::get(StructTy, DL.getAllocaAddrSpace()));
   }
 
   LLVM_DEBUG({

diff  --git a/llvm/test/tools/llvm-extract/address-space-assert.ll b/llvm/test/tools/llvm-extract/address-space-assert.ll
new file mode 100644
index 0000000000000..94f3fdd72598d
--- /dev/null
+++ b/llvm/test/tools/llvm-extract/address-space-assert.ll
@@ -0,0 +1,32 @@
+; RUN: llvm-extract -S --bb=func:bb4 -aggregate-extracted-args=0 < %s | FileCheck %s
+
+; FIXME: aggregate-extracted-args doesn't work for other reasons
+; XUN: llvm-extract -S --bb=func:bb4 -aggregate-extracted-args=1 < %s | FileCheck %s
+
+target datalayout = "A5-G1-ni:7"
+
+; Check that there's no assert from incorrect pointer types used in the new arguments.
+
+; CHECK-LABEL: define dso_local void @func.bb4(i32 %orig.arg.0, ptr addrspace(5) %tmp1.out, ptr addrspace(5) %add.out) {
+; CHECK: bb4:
+; CHECK-NEXT: %tmp0 = add i32 0, 0
+; CHECK-NEXT: %tmp1 = add i32 1, 1
+; CHECK-NEXT: store i32 %tmp1, ptr addrspace(5) %tmp1.out, align 4
+; CHECK-NEXT: %add = add i32 %tmp0, %orig.arg.0
+; CHECK-NEXT: store i32 %add, ptr addrspace(5) %add.out, align 4
+; CHECK-NEXT: br label %bb5.exitStub
+define void @func(i32 %orig.arg.0, ptr addrspace(1) %orig.arg.1) {
+bb:
+  br label %bb4
+
+bb4:                                              ; preds = %bb
+  %tmp0 = add i32 0, 0
+  %tmp1 = add i32 1, 1
+  %add = add i32 %tmp0, %orig.arg.0
+  br label %bb5
+
+bb5:                                              ; preds = %bb5, %bb4
+  %tmp6 = phi i32 [ %add, %bb4 ], [ 0, %bb5 ]
+  %tmp7 = phi i32 [ %tmp1, %bb4 ], [ 2, %bb5 ]
+  br label %bb5
+}