[llvm] [AMDGPU] Re-enable closed-world assumption as an opt-out feature (PR #115371)

[Shilei Tian via llvm-commits]

https://github.com/shiltian created https://github.com/llvm/llvm-project/pull/115371

Although the ABI (if any exists) doesn’t explicitly prohibit cross-device-image
function calls, especially since our loader can handle them, for all officially
supported programming models, this is not actually allowed. Given this, assuming
a closed-world model at link time is safe. However, there are certain cases,
such as the GPU libc project, that use non-standard approaches which could break
this assumption. This PR introduces an option to disable this assumption when
needed.

>From 0ef2313eecc010403c0fd293dd2b931856a8d4e8 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Thu, 7 Nov 2024 16:09:12 -0500
Subject: [PATCH] [AMDGPU] Re-enable closed-world assumption as an opt-out
 feature
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Although the ABI (if any exists) doesn’t explicitly prohibit cross-device-image
function calls, especially since our loader can handle them, for all officially
supported programming models, this is not actually allowed. Given this, assuming
a closed-world model at link time is safe. However, there are certain cases,
such as the GPU libc project, that use non-standard approaches which could break
this assumption. This PR introduces an option to disable this assumption when
needed.
---
 llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp     |  4 ++++
 llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp  | 13 +++++++++++--
 llvm/test/LTO/AMDGPU/closed-world-assumption.ll | 12 ++++++++++++
 3 files changed, 27 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/LTO/AMDGPU/closed-world-assumption.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
index 2ae34636005eac..0a33ff7072be08 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
@@ -1068,6 +1068,10 @@ static bool runImpl(Module &M, AnalysisGetter &AG, TargetMachine &TM,
 
   Attributor A(Functions, InfoCache, AC);
 
+  LLVM_DEBUG(dbgs() << "[AMDGPUAttributor] Module " << M.getName() << " is "
+                    << (AC.IsClosedWorldModule ? "" : "not ")
+                    << "assumed to be a closed world.\n");
+
   for (auto *F : Functions) {
     A.getOrCreateAAFor<AAAMDAttributes>(IRPosition::function(*F));
     A.getOrCreateAAFor<AAUniformWorkGroupSize>(IRPosition::function(*F));
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index 786baa6820e860..6b93a659debb7b 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -449,6 +449,11 @@ static cl::opt<bool>
                            cl::desc("Enable AMDGPUAttributorPass"),
                            cl::init(true), cl::Hidden);
 
+static cl::opt<bool> HasClosedWorldAssumption(
+    "amdgpu-link-time-closed-world",
+    cl::desc("Whether has closed-world assumption at link time"),
+    cl::init(true), cl::Hidden);
+
 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAMDGPUTarget() {
   // Register the target
   RegisterTargetMachine<R600TargetMachine> X(getTheR600Target());
@@ -836,8 +841,12 @@ void AMDGPUTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
             PM.addPass(InternalizePass(mustPreserveGV));
             PM.addPass(GlobalDCEPass());
           }
-          if (EnableAMDGPUAttributor)
-            PM.addPass(AMDGPUAttributorPass(*this));
+          if (EnableAMDGPUAttributor) {
+            AMDGPUAttributorOptions Opt;
+            if (HasClosedWorldAssumption)
+              Opt.IsClosedWorld = true;
+            PM.addPass(AMDGPUAttributorPass(*this, Opt));
+          }
         }
       });
 
diff --git a/llvm/test/LTO/AMDGPU/closed-world-assumption.ll b/llvm/test/LTO/AMDGPU/closed-world-assumption.ll
new file mode 100644
index 00000000000000..cfd3b0db74ccb0
--- /dev/null
+++ b/llvm/test/LTO/AMDGPU/closed-world-assumption.ll
@@ -0,0 +1,12 @@
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -O3 -debug-only=amdgpu-attributor -o - %s 2>&1 | FileCheck %s --check-prefix=NO-CW
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes="lto<O3>" -debug-only=amdgpu-attributor -o - %s 2>&1 | FileCheck %s --check-prefix=CW
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes="lto<O3>" -debug-only=amdgpu-attributor -amdgpu-link-time-closed-world=0 -o - %s 2>&1 | FileCheck %s --check-prefix=NO-CW
+
+; REQUIRES: amdgpu-registered-target
+; REQUIRES: asserts
+
+; NO-CW: Module {{.*}} is not assumed to be a closed world.
+; CW: Module {{.*}} is assumed to be a closed world.
+define hidden noundef i32 @_Z3foov() {
+  ret i32 1
+}