[NVPTX][InferAS] assume alloca instructions are in local AS

[AlexMaclean]

No description provided.

[llvmbot]

@llvm/pr-subscribers-llvm-selectiondag
@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-backend-nvptx

Author: Alex MacLean (AlexMaclean)

Changes

---

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

3 Files Affected:

• (modified) llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp (+9-1)
• (modified) llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h (+2)
• (added) llvm/test/Transforms/InferAddressSpaces/NVPTX/alloca.ll (+17)

diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp b/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp
index 4ec2ec100ab08d..3fe52fb33c5bd0 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp
@@ -20,6 +20,7 @@
 #include "llvm/IR/Value.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/NVPTXAddrSpace.h"
 #include "llvm/Transforms/InstCombine/InstCombiner.h"
 #include <optional>
 using namespace llvm;
@@ -562,4 +563,11 @@ Value *NVPTXTTIImpl::rewriteIntrinsicWithAddressSpace(IntrinsicInst *II,
   }
   }
   return nullptr;
-}
\ No newline at end of file
+}
+
+unsigned NVPTXTTIImpl::getAssumedAddrSpace(const Value *V) const {
+  if (isa<AllocaInst>(V))
+    return ADDRESS_SPACE_LOCAL;
+
+  return -1;
+}
diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h b/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h
index 0f4fb280b2d996..a378b8a37f7043 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h
@@ -129,6 +129,8 @@ class NVPTXTTIImpl : public BasicTTIImplBase<NVPTXTTIImpl> {
 
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const;
+
+  unsigned getAssumedAddrSpace(const Value *V) const;
 };
 
 } // end namespace llvm
diff --git a/llvm/test/Transforms/InferAddressSpaces/NVPTX/alloca.ll b/llvm/test/Transforms/InferAddressSpaces/NVPTX/alloca.ll
new file mode 100644
index 00000000000000..fa063cdf8d8054
--- /dev/null
+++ b/llvm/test/Transforms/InferAddressSpaces/NVPTX/alloca.ll
@@ -0,0 +1,17 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=infer-address-spaces %s | FileCheck %s
+
+target triple = "nvptx64-nvidia-cuda"
+
+
+define float @load_alloca() {
+; CHECK-LABEL: define float @load_alloca() {
+; CHECK-NEXT:    [[ADDR:%.*]] = alloca float, align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr [[ADDR]] to ptr addrspace(5)
+; CHECK-NEXT:    [[VAL:%.*]] = load float, ptr addrspace(5) [[TMP1]], align 4
+; CHECK-NEXT:    ret float [[VAL]]
+;
+  %addr = alloca float
+  %val = load float, ptr %addr
+  ret float %val
+}

[jdoerfert]

LG

[github-actions]

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

[arsenm]

This is worse than just changing the alloca addrspace

[AlexMaclean]

While I tend to agree that converting all allocas to local addrspace would be a good idea, it's also a deep can of worms. As @Artem-B points out here #106127 (comment), optimizations may not handle allocas in specific AS as well, and supporting local allocas would likely require changes in the backend and datalayout.

I do hope to pursue this in the future, but in the meantime this change is both simple and correct. Even if we were to address the above it still seems preferable to handle generic allocas here in case InferAS happens to encounter one in the IR.

[arsenm]

The only optimization using the non-0 addrspace really breaks folding compares to null out, which rarely matters if you can see the original alloca. We should also make the non-0 address space null handling datalayout configurable

[AlexMaclean]

@Artem-B were there other concerns you had with moving allocas to non-0 AS?

Even if we want to go this route I fear it will be pretty difficult and invasive. getAssumedAddrSpace is supposed to return the true AS of a given value in the AS 0. Handling this case there doesn't preclude moving allocas to non-0 AS in the future. I think it is good to add this check regardless to make InferAS make robust to different possible IRs.

[AlexMaclean]

Ping @Artem-B / @arsenm for any further thoughts on the feasibility of switching allocas to the local addrspace and whether this change is acceptable as a complimentary/intermediate-term solution.

[AlexMaclean]

@arsenm would it be alright if I proceeded with this change despite your point about switching to specific-AS allocas? I agree this is worth investigating but it will take a fair bit of work I think, and this solution is small, shows immediate improvements, and won't hamper any efforts to handle specific-AS allocas in the future.

[AlexMaclean]

@arsenm ping on this discussion.

[AlexMaclean]

Since I haven't heard back on this discussion for a while now, and this seems more like a possible future enhancement than a blocking issue, I'll plan to land this change at the end of the week unless I hear back otherwise. CC @arsenm

[Artem-B]

LGTM with a minor nit.

[Artem-B]

still LGTM for me.