Inline: Fix handling of byval using non-alloca addrspace

[arsenm]

Use the address space of the original pointer argument instead
of querying the datalayout. This avoids producing a verifier error
since this was changing the address space for the user instructions.

Fixes #97086

[arsenm]

• Inline: Fix handling of byval using non-alloca addrspace #97306 👈
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

Join @arsenm and the rest of your teammates on Graphite

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Matt Arsenault (arsenm)

Changes

Use the address space of the original pointer argument instead
of querying the datalayout. This avoids producing a verifier error
since this was changing the address space for the user instructions.

Fixes #97086

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/InlineFunction.cpp (+3-2)
• (added) llvm/test/Transforms/Inline/byval-with-non-alloca-addrspace.ll (+38)

diff --git a/llvm/lib/Transforms/Utils/InlineFunction.cpp b/llvm/lib/Transforms/Utils/InlineFunction.cpp
index 0725addfbb90a..46605492490ea 100644
--- a/llvm/lib/Transforms/Utils/InlineFunction.cpp
+++ b/llvm/lib/Transforms/Utils/InlineFunction.cpp
@@ -1674,8 +1674,9 @@ static Value *HandleByValArgument(Type *ByValType, Value *Arg,
   if (ByValAlignment)
     Alignment = std::max(Alignment, *ByValAlignment);
 
-  AllocaInst *NewAlloca = new AllocaInst(ByValType, DL.getAllocaAddrSpace(),
-                                         nullptr, Alignment, Arg->getName());
+  AllocaInst *NewAlloca = new AllocaInst(
+      ByValType, cast<PointerType>(Arg->getType())->getAddressSpace(), nullptr,
+      Alignment, Arg->getName());
   NewAlloca->insertBefore(Caller->begin()->begin());
   IFI.StaticAllocas.push_back(NewAlloca);
 
diff --git a/llvm/test/Transforms/Inline/byval-with-non-alloca-addrspace.ll b/llvm/test/Transforms/Inline/byval-with-non-alloca-addrspace.ll
new file mode 100644
index 0000000000000..42ec0f2bf5699
--- /dev/null
+++ b/llvm/test/Transforms/Inline/byval-with-non-alloca-addrspace.ll
@@ -0,0 +1,38 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=always-inline %s | FileCheck %s
+
+; The byval parameters use a different address space from the alloca
+; address space. Make sure this is gracefully handled by using the
+; original byval pointer type's address space instead of what the
+; datalayout says to use.
+
+target datalayout = "A5"
+
+%struct = type { i64, i64 }
+
+define i64 @bar(ptr byval(%struct) %a) alwaysinline {
+; CHECK-LABEL: define i64 @bar(
+; CHECK-SAME: ptr byval([[STRUCT:%.*]]) [[A:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [[STRUCT]], ptr [[A]], i64 0, i32 1
+; CHECK-NEXT:    [[TMP2:%.*]] = load i64, ptr [[TMP1]], align 4
+; CHECK-NEXT:    ret i64 0
+;
+  %1 = getelementptr %struct, ptr %a, i64 0, i32 1
+  %2 = load i64, ptr %1, align 4
+  ret i64 0
+}
+
+define i64 @foo(ptr %arg) {
+; CHECK-LABEL: define i64 @foo(
+; CHECK-SAME: ptr [[ARG:%.*]]) {
+; CHECK-NEXT:    [[ARG1:%.*]] = alloca [[STRUCT:%.*]], align 8
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 16, ptr [[ARG1]])
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr align 1 [[ARG1]], ptr align 1 [[ARG]], i64 16, i1 false)
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [[STRUCT]], ptr [[ARG1]], i64 0, i32 1
+; CHECK-NEXT:    [[TMP2:%.*]] = load i64, ptr [[TMP1]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 16, ptr [[ARG1]])
+; CHECK-NEXT:    ret i64 0
+;
+  %1 = call i64 @bar(ptr byval(%struct) align 8 %arg)
+  ret i64 0
+}

[nikic]

LGTM

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[MrSidims]

@arsenm
Digging a bit into git log the patch that added DL.getAllocaAddrSpace() was the following:

commit 3c1fc768ed6cf2b01463df036ae6ae3b1f0de632
Author: Matt Arsenault <[email protected]>
Date:   Mon Apr 10 22:27:50 2017 +0000

    Allow DataLayout to specify addrspace for allocas.

and reasoning behind that patch:
" These are problems for AMDGPU because alloca is used to
implement the private address space, which uses a 32-bit
index as the pointer value. Other pointers are 64-bit
and behave more like LLVM's notion of generic address
space. By changing the address space used for allocas,
we can change our generic pointer type to be LLVM's generic
pointer type which does have similar properties."

makes sense to me especially in OpenCL context, where default address space is 4 (generic). Also IMHO LLVM passes must respect target's datalayout when picking possible types for the created instructions.

Current fix brings some complications though (from OpenCL point of view, but there might be issues for other languages and targets outside of my expertise), if byval pointer of a function was in default aka genetic address space, then the alloca would be in generic. Meanwhile per SPIR-V (to which LLVM IR is being lowered starting from OpenCL 2.0) specification OpVariable must have non-generic storage class. It can probably be workarounded in https://github.com/KhronosGroup/SPIRV-LLVM-Translator by translating alloca addrspace(4) to OpVariable with Private storage class bringing extra address space casts, or by modifying clang, introducing an exception for byval pointers making them not be in generic address space, but in private, but I wonder if #97086 can be solved differently - instead here if Arg->getType()->getPointerAddressSpace() != DL.getAllocaAddrSpace() - we still create an alloca in default datalayout address space and add an AS cast, WDYT?

[davidchisnall]

instead here if Arg->getType()->getPointerAddressSpace() != DL.getAllocaAddrSpace() - we still create an alloca in default datalayout address space and add an AS cast, WDYT?

I'm not sure I fully understand the proposal, but for the CHERI use case (the AMD version was based on the CHERI fork, which we aim to start upstreaming soon), allocas must be created in the alloca address space, not some other address space and cast.

[MrSidims]

I mean: leave alloca in DL.getAllocaAddrSpace() AS and if it doesn't match with byval pointer's address space - create a cast from alloca to byval's address space and update users of the alloca to use the cast. But digging more into it might be not as trivial as I though initially - at least it seems it shouldn't be done in HandleByValArgument, but probably later in the stack.

[davidchisnall]

That would not break anything for CHERI, if it works for other use cases.

[arsenm]

No, this is the correct fix. We should follow what the existing IR says. The rule is not all allocas must have the alloca address space, it's the alloca address space is the default for contextless new values.

As far as OpenCL is concerned such IR would have been malformed in the first place. If we wanted to handle this as a generic pointer with a private backed alloca, it should be handled by lowering the alloca with the wrong address space to include the cast (which is essentially what NVPTX does, it's just more work for the backend and less optimizable IR)