nvptx-lower-alloca introduces pointer with incorrect address space: "GEP address space doesn't match type"

[ingomueller-net]

I am running into an issue related to address spaces in the following snippets:

; RUN: llc -mtriple nvptx64-nvidia-cuda %s
; ModuleID = 'LLVMDialectModule'
source_filename = "LLVMDialectModule"
target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"

; Function Attrs: nounwind
define void @main_kernel() {
.preheader11:
  %0 = alloca [2048 x float], align 4, addrspace(3)
  %1 = getelementptr float, ptr addrspace(3) %0, i64 undef
  call void @llvm.masked.store.v4f32.p3(<4 x float> undef, ptr addrspace(3) %1, i32 4, <4 x i1> zeroinitializer)
  ret void
}

; Function Attrs: mustprogress nocallback nofree nosync nounwind willreturn memory(argmem: write)
declare void @llvm.masked.store.v4f32.p3(<4 x float> %0, ptr addrspace(3) nocapture %1, i32 immarg %2, <4 x i1> %3) #4

Running the command from the first line crashes with the following output (on the latest main):

GEP address space doesn't match type
  %3 = getelementptr float, ptr %2, i64 undef
in function main_kernel
LLVM ERROR: Broken function found, compilation aborted!
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.      Program arguments: /data//git/llvm-project/build-Debug/bin/llc -mtriple nvptx64-nvidia-cuda /data//git/llvm-project/llvm/test/module.ll
1.      Running pass 'Function Pass Manager' on module '/data//git/llvm-project/llvm/test/module.ll'.
2.      Running pass 'Module Verifier' on function '@main_kernel'
 #0 0x00007f98a462195d llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) /data//git/llvm-project/llvm/lib/Support/Unix/Signals.inc:723:11
 #1 0x00007f98a4621e4b PrintStackTraceSignalHandler(void*) /data//git/llvm-project/llvm/lib/Support/Unix/Signals.inc:798:1
 #2 0x00007f98a461fe76 llvm::sys::RunSignalHandlers() /data//git/llvm-project/llvm/lib/Support/Signals.cpp:105:5
 #3 0x00007f98a4622665 SignalHandler(int) /data//git/llvm-project/llvm/lib/Support/Unix/Signals.inc:413:1
 #4 0x00007f98a965a540 (/lib/x86_64-linux-gnu/libc.so.6+0x3c540)
 #5 0x00007f98a96a812c __pthread_kill_implementation ./nptl/pthread_kill.c:44:76
 #6 0x00007f98a965a4a2 raise ./signal/../sysdeps/posix/raise.c:27:6
 #7 0x00007f98a96444b2 abort ./stdlib/abort.c:81:7
 #8 0x00007f98a44efdb4 llvm::report_fatal_error(llvm::Twine const&, bool) /data//git/llvm-project/llvm/lib/Support/ErrorHandling.cpp:125:5
 #9 0x00007f98a44efc22 /data//git/llvm-project/llvm/lib/Support/ErrorHandling.cpp:83:3
#10 0x00007f98a522b46b (anonymous namespace)::VerifierLegacyPass::runOnFunction(llvm::Function&) /data//git/llvm-project/llvm/lib/IR/Verifier.cpp:6619:5
#11 0x00007f98a50f06aa llvm::FPPassManager::runOnFunction(llvm::Function&) /data//git/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:1435:23
#12 0x00007f98a50f54d2 llvm::FPPassManager::runOnModule(llvm::Module&) /data//git/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:1481:16
#13 0x00007f98a50f0f5b (anonymous namespace)::MPPassManager::runOnModule(llvm::Module&) /data//git/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:1550:23
#14 0x00007f98a50f0add llvm::legacy::PassManagerImpl::run(llvm::Module&) /data//git/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:535:16
#15 0x00007f98a50f57b1 llvm::legacy::PassManager::run(llvm::Module&) /data//git/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:1677:3
#16 0x0000556b364c6e17 compileModule(char**, llvm::LLVMContext&) /data//git/llvm-project/llvm/tools/llc/llc.cpp:754:41
#17 0x0000556b364c515d main /data//git/llvm-project/llvm/tools/llc/llc.cpp:416:13
#18 0x00007f98a96456ca __libc_start_call_main ./csu/../sysdeps/nptl/libc_start_call_main.h:74:3
#19 0x00007f98a9645785 call_init ./csu/../csu/libc-start.c:128:20
#20 0x00007f98a9645785 __libc_start_main ./csu/../csu/libc-start.c:347:5
#21 0x0000556b364c4931 _start (/data//git/llvm-project/build-Debug/bin/llc+0x2e931)

The error message is caused by this check:

if (auto *PTy = dyn_cast<PointerType>(GEP.getType())) {
  Check(GEP.getAddressSpace() == PTy->getAddressSpace(),
        "GEP address space doesn't match type", &GEP);
}

If I print PTy and GEP.getPointerOperandType(), I get ptr addrspace(3) and ptr, respectively, which don't match. Note, however, that the %0 operand of the GEP instruction in the input above is actually annotated with addrspace(3), so that attribute seems to get lost somehow.

Unfortunately, I am not familiar enough with LLVM to debug this issue further.

I have run into this issue in a lowering pipeline from MLIR, which seems to be using this part of LLVM in its SerializeToCubin pass. The module generated by MLIR was several thousands lines; above snippet is an (approximatly) minimal example that produces the same error. Note that I inserted the undefs in order to reduce the involved registers.

[fhahn]

Looks like this is caused by -nvptx-lower-alloca: https://godbolt.org/z/sa85Y4o31

[ingomueller-net]

I have tried out all llvmorg-*-init branches back until 15 (i.e., around the time of LLVM 14) and they all have this problem. LLVM 14 only accepts the input with -opaque-pointers and LLVM 13 (where I tested a release from my distro) does not accept the input at all with the following error:

error: '%0' defined with type '[2048 x float] addrspace(3)*' but expected 'ptr addrspace(3)'

[Artem-B]

%0 = alloca [2048 x float], align 4, addrspace(3)

This looks like the source of the problem. Address space 3 is shared memory. Alloca only works for local memory. Shared memory is always allocated statically.

[ingomueller-net]

(Apart from not crashing) what's the way to fix this? I.e., what IR should MLIR generate?

[Artem-B]

It should not generate alloca [2048 x float], align 4, addrspace(3). I have no idea what the code is trying to do, but alloca is not going to work w/ shared memory. You either drop addrspace (3) from the alloca, or allocate shared memory externally and either use a static address or pass the shared memory pointer as an argument.

[ingomueller-net]

OK, thanks a lot for the pointers!