From 192c0e5a7a879c814bd51702c710b1d3f2f59981 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault@amd.com>
Date: Wed, 4 Jan 2023 11:07:34 -0500
Subject: [PATCH] IROutliner: Fix assert with non-0 alloca addrspace

The arguments are passed as stored to new allocas so the address space
needs to match.
---
 llvm/lib/Transforms/IPO/IROutliner.cpp             |  7 ++--
 .../test/Transforms/IROutliner/alloca-addrspace.ll | 47 ++++++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Transforms/IROutliner/alloca-addrspace.ll

diff --git a/llvm/lib/Transforms/IPO/IROutliner.cpp b/llvm/lib/Transforms/IPO/IROutliner.cpp
index 35366d1..f5c52e5 100644
--- a/llvm/lib/Transforms/IPO/IROutliner.cpp
+++ b/llvm/lib/Transforms/IPO/IROutliner.cpp
@@ -1277,7 +1277,7 @@ static std::optional<unsigned> getGVNForPHINode(OutlinableRegion &Region,
 /// \param [in,out] Region - The region of code to be analyzed.
 /// \param [in] Outputs - The values found by the code extractor.
 static void
-findExtractedOutputToOverallOutputMapping(OutlinableRegion &Region,
+findExtractedOutputToOverallOutputMapping(Module &M, OutlinableRegion &Region,
                                           SetVector<Value *> &Outputs) {
   OutlinableGroup &Group = *Region.Parent;
   IRSimilarityCandidate &C = *Region.Candidate;
@@ -1350,7 +1350,8 @@ findExtractedOutputToOverallOutputMapping(OutlinableRegion &Region,
     // the output, so we add a pointer type to the argument types of the overall
     // function to handle this output and create a mapping to it.
     if (!TypeFound) {
-      Group.ArgumentTypes.push_back(PointerType::getUnqual(Output->getType()));
+      Group.ArgumentTypes.push_back(Output->getType()->getPointerTo(
+          M.getDataLayout().getAllocaAddrSpace()));
       // Mark the new pointer type as the last value in the aggregate argument
       // list.
       unsigned ArgTypeIdx = Group.ArgumentTypes.size() - 1;
@@ -1418,7 +1419,7 @@ void IROutliner::findAddInputsOutputs(Module &M, OutlinableRegion &Region,
 
   // Map the outputs found by the CodeExtractor to the arguments found for
   // the overall function.
-  findExtractedOutputToOverallOutputMapping(Region, Outputs);
+  findExtractedOutputToOverallOutputMapping(M, Region, Outputs);
 }
 
 /// Replace the extracted function in the Region with a call to the overall
diff --git a/llvm/test/Transforms/IROutliner/alloca-addrspace.ll b/llvm/test/Transforms/IROutliner/alloca-addrspace.ll
new file mode 100644
index 0000000..8f06d93
--- /dev/null
+++ b/llvm/test/Transforms/IROutliner/alloca-addrspace.ll
@@ -0,0 +1,47 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --include-generated-funcs
+; RUN: opt -S -passes=iroutliner -ir-outlining-no-cost < %s | FileCheck %s
+
+; Check alloca with non-0 address spaces works correctly.
+
+target datalayout = "A5"
+
+define i32 @outlineable() {
+bb:
+  %i = tail call i32 @func(i32 0, i32 1)
+  %i1 = or i32 0, %i
+  %i2 = tail call i32 @func(i32 %i1, i32 0)
+  %i3 = or i32 %i1, %i2
+  ret i32 0
+}
+
+declare i32 @func(i32, i32)
+; CHECK-LABEL: define {{[^@]+}}@outlineable() {
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[I1_LOC:%.*]] = alloca i32, align 4, addrspace(5)
+; CHECK-NEXT:    [[LT_CAST:%.*]] = addrspacecast ptr addrspace(5) [[I1_LOC]] to ptr
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 -1, ptr [[LT_CAST]])
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32 0, i32 1, ptr addrspace(5) [[I1_LOC]], i32 0)
+; CHECK-NEXT:    [[I1_RELOAD:%.*]] = load i32, ptr addrspace(5) [[I1_LOC]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 -1, ptr [[LT_CAST]])
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32 [[I1_RELOAD]], i32 0, ptr addrspace(5) null, i32 -1)
+; CHECK-NEXT:    ret i32 0
+;
+;
+; CHECK-LABEL: define {{[^@]+}}@outlined_ir_func_0
+; CHECK-SAME: (i32 [[TMP0:%.*]], i32 [[TMP1:%.*]], ptr addrspace(5) [[TMP2:%.*]], i32 [[TMP3:%.*]]) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    br label [[BB_TO_OUTLINE:%.*]]
+; CHECK:       bb_to_outline:
+; CHECK-NEXT:    [[I:%.*]] = tail call i32 @func(i32 [[TMP0]], i32 [[TMP1]])
+; CHECK-NEXT:    [[I1:%.*]] = or i32 [[TMP0]], [[I]]
+; CHECK-NEXT:    br label [[BB_AFTER_OUTLINE_EXITSTUB:%.*]]
+; CHECK:       bb_after_outline.exitStub:
+; CHECK-NEXT:    switch i32 [[TMP3]], label [[FINAL_BLOCK_0:%.*]] [
+; CHECK-NEXT:    i32 0, label [[OUTPUT_BLOCK_0_0:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       output_block_0_0:
+; CHECK-NEXT:    store i32 [[I1]], ptr addrspace(5) [[TMP2]], align 4
+; CHECK-NEXT:    br label [[FINAL_BLOCK_0]]
+; CHECK:       final_block_0:
+; CHECK-NEXT:    ret void
+;
-- 
2.7.4