[HLSL] Change default linkage of HLSL functions to internal (llvm#95331)

hekota · web-flow · commit db279c72f2fe · 2024-08-16T10:32:10.000-07:00
An HLSL function has internal linkage by default unless it is: 1. shader entry point function 2. marked with the `export` keyword (llvm#92812) 3. patch constant function (not implemented yet) This PR adds a link-time pass `DXILFinalizeLinkage` that updates the linkage of functions to make sure only shader entry points and exported functions are visible from the module (have _program linkage_). All other functions will be updated to have internal linkage. Related spec update: microsoft/hlsl-specs#295 Fixes #llvm#92071
diff --git a/llvm/lib/Target/DirectX/CMakeLists.txt b/llvm/lib/Target/DirectX/CMakeLists.txt
@@ -19,6 +19,7 @@ add_llvm_target(DirectXCodeGen
   DirectXSubtarget.cpp
   DirectXTargetMachine.cpp
   DXContainerGlobals.cpp
+  DXILFinalizeLinkage.cpp
   DXILIntrinsicExpansion.cpp
   DXILMetadata.cpp
   DXILOpBuilder.cpp
diff --git a/llvm/lib/Target/DirectX/DXILFinalizeLinkage.cpp b/llvm/lib/Target/DirectX/DXILFinalizeLinkage.cpp
@@ -0,0 +1,60 @@
+//===- DXILFinalizeLinkage.cpp - Finalize linkage of functions ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "DXILFinalizeLinkage.h"
+#include "DirectX.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalValue.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+
+#define DEBUG_TYPE "dxil-finalize-linkage"
+
+using namespace llvm;
+
+static bool finalizeLinkage(Module &M) {
+  SmallPtrSet<Function *, 8> EntriesAndExports;
+
+  // Find all entry points and export functions
+  for (Function &EF : M.functions()) {
+    if (!EF.hasFnAttribute("hlsl.shader") && !EF.hasFnAttribute("hlsl.export"))
+      continue;
+    EntriesAndExports.insert(&EF);
+  }
+
+  for (Function &F : M.functions()) {
+    if (F.getLinkage() == GlobalValue::ExternalLinkage &&
+        !EntriesAndExports.contains(&F)) {
+      F.setLinkage(GlobalValue::InternalLinkage);
+    }
+  }
+
+  return false;
+}
+
+PreservedAnalyses DXILFinalizeLinkage::run(Module &M,
+                                           ModuleAnalysisManager &AM) {
+  if (finalizeLinkage(M))
+    return PreservedAnalyses::none();
+  return PreservedAnalyses::all();
+}
+
+bool DXILFinalizeLinkageLegacy::runOnModule(Module &M) {
+  return finalizeLinkage(M);
+}
+
+char DXILFinalizeLinkageLegacy::ID = 0;
+
+INITIALIZE_PASS_BEGIN(DXILFinalizeLinkageLegacy, DEBUG_TYPE,
+                      "DXIL Finalize Linkage", false, false)
+INITIALIZE_PASS_END(DXILFinalizeLinkageLegacy, DEBUG_TYPE,
+                    "DXIL Finalize Linkage", false, false)
+
+ModulePass *llvm::createDXILFinalizeLinkageLegacyPass() {
+  return new DXILFinalizeLinkageLegacy();
+}
diff --git a/llvm/lib/Target/DirectX/DXILFinalizeLinkage.h b/llvm/lib/Target/DirectX/DXILFinalizeLinkage.h
@@ -0,0 +1,39 @@
+//===- DXILFinalizeLinkage.h - Finalize linkage of functions --------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// DXILFinalizeLinkage pass updates the linkage of functions to make sure only
+/// shader entry points and exported functions are visible from the module (have
+/// program linkage). All other functions will be updated to have internal
+/// linkage.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TARGET_DIRECTX_DXILFINALIZELINKAGE_H
+#define LLVM_TARGET_DIRECTX_DXILFINALIZELINKAGE_H
+
+#include "llvm/IR/PassManager.h"
+#include "llvm/Pass.h"
+
+namespace llvm {
+
+class DXILFinalizeLinkage : public PassInfoMixin<DXILFinalizeLinkage> {
+public:
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &);
+  static bool isRequired() { return true; }
+};
+
+class DXILFinalizeLinkageLegacy : public ModulePass {
+public:
+  DXILFinalizeLinkageLegacy() : ModulePass(ID) {}
+  bool runOnModule(Module &M) override;
+
+  static char ID; // Pass identification.
+};
+} // namespace llvm
+
+#endif // LLVM_TARGET_DIRECTX_DXILFINALIZELINKAGE_H
diff --git a/llvm/lib/Target/DirectX/DirectX.h b/llvm/lib/Target/DirectX/DirectX.h
@@ -63,6 +63,13 @@ void initializeDXContainerGlobalsPass(PassRegistry &);
 
 /// Pass for generating DXContainer part globals.
 ModulePass *createDXContainerGlobalsPass();
+
+/// Initializer for DXILFinalizeLinkage pass.
+void initializeDXILFinalizeLinkageLegacyPass(PassRegistry &);
+
+/// Pass to finalize linkage of functions.
+ModulePass *createDXILFinalizeLinkageLegacyPass();
+
 } // namespace llvm
 
 #endif // LLVM_LIB_TARGET_DIRECTX_DIRECTX_H
diff --git a/llvm/lib/Target/DirectX/DirectXTargetMachine.cpp b/llvm/lib/Target/DirectX/DirectXTargetMachine.cpp
@@ -48,6 +48,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeDirectXTarget() {
   initializeDXILTranslateMetadataPass(*PR);
   initializeDXILResourceMDWrapperPass(*PR);
   initializeShaderFlagsAnalysisWrapperPass(*PR);
+  initializeDXILFinalizeLinkageLegacyPass(*PR);
 }
 
 class DXILTargetObjectFile : public TargetLoweringObjectFile {
@@ -79,6 +80,7 @@ class DirectXPassConfig : public TargetPassConfig {
   void addCodeGenPrepare() override {
     addPass(createDXILIntrinsicExpansionLegacyPass());
     addPass(createDXILOpLoweringLegacyPass());
+    addPass(createDXILFinalizeLinkageLegacyPass());
     addPass(createDXILTranslateMetadataPass());
     addPass(createDXILPrepareModulePass());
   }
diff --git a/llvm/test/CodeGen/DirectX/conflicting-bitcast-insert.ll b/llvm/test/CodeGen/DirectX/conflicting-bitcast-insert.ll
@@ -7,7 +7,7 @@ define i64 @test(ptr %p) {
   ret i64 %v
 }
 
-; CHECK: define i64 @test(ptr %p) {
+; CHECK: define internal i64 @test(ptr %p) {
 ; CHECK-NEXT: %1 = bitcast ptr %p to ptr
 ; CHECK-NEXT: store i32 0, ptr %1, align 4
 ; CHECK-NEXT: %2 = bitcast ptr %p to ptr
@@ -19,7 +19,7 @@ define i64 @testGEP(ptr %p) {
   ret i64 %val
 }
 
-; CHECK: define i64 @testGEP(ptr %p) {
+; CHECK: define internal i64 @testGEP(ptr %p) {
 ; CHECK-NEXT:   %1 = bitcast ptr %p to ptr
 ; CHECK-NEXT:   %ptr = getelementptr i32, ptr %1, i32 4
 ; CHECK-NEXT:   %2 = bitcast ptr %p to ptr
diff --git a/llvm/test/CodeGen/DirectX/finalize_linkage.ll b/llvm/test/CodeGen/DirectX/finalize_linkage.ll
@@ -0,0 +1,64 @@
+; RUN: opt -S -dxil-finalize-linkage -mtriple=dxil-unknown-shadermodel6.5-compute %s | FileCheck %s
+; RUN: llc %s --filetype=asm -o - | FileCheck %s --check-prefixes=CHECK-LLC
+
+target triple = "dxilv1.5-pc-shadermodel6.5-compute"
+
+; DXILFinalizeLinkage changes linkage of all functions that are not
+; entry points or exported function to internal.
+
+; CHECK: define internal void @"?f1@@YAXXZ"()
+define void @"?f1@@YAXXZ"() #0 {
+entry:
+  ret void
+}
+
+; CHECK: define internal void @"?f2@@YAXXZ"()
+define void @"?f2@@YAXXZ"() #0 {
+entry:
+  ret void
+}
+
+; CHECK: define internal void @"?f3@@YAXXZ"()
+define void @"?f3@@YAXXZ"() #0 {
+entry:
+  ret void
+}
+
+; CHECK: define internal void @"?foo@@YAXXZ"()
+define void @"?foo@@YAXXZ"() #0 {
+entry:
+  call void @"?f2@@YAXXZ"() #3
+  ret void
+}
+
+; Exported function - do not change linkage
+; CHECK: define void @"?bar@@YAXXZ"()
+define void @"?bar@@YAXXZ"() #1 {
+entry:
+  call void @"?f3@@YAXXZ"() #3
+  ret void
+}
+
+; CHECK: define internal void @"?main@@YAXXZ"() #0
+define internal void @"?main@@YAXXZ"() #0 {
+entry:
+  call void @"?foo@@YAXXZ"() #3
+  call void @"?bar@@YAXXZ"() #3
+  ret void
+}
+
+; Entry point function - do not change linkage
+; CHECK: define void @main() #2
+define void @main() #2 {
+entry:
+  call void @"?main@@YAXXZ"()
+  ret void
+}
+
+attributes #0 = { convergent noinline nounwind optnone}
+attributes #1 = { convergent noinline nounwind optnone "hlsl.export"}
+attributes #2 = { convergent "hlsl.numthreads"="4,1,1" "hlsl.shader"="compute"}
+attributes #3 = { convergent }
+
+; Make sure "hlsl.export" attribute is stripped by llc
+; CHECK-LLC-NOT: "hlsl.export"
diff --git a/llvm/test/CodeGen/DirectX/omit-bitcast-insert.ll b/llvm/test/CodeGen/DirectX/omit-bitcast-insert.ll
@@ -6,7 +6,7 @@ define i64 @test(ptr %p) {
   ret i64 %v
 }
 
-; CHECK: define i64 @test(ptr %p) {
+; CHECK: define internal i64 @test(ptr %p) {
 ; CHECK-NEXT: %v = load i64, ptr %p, align 8
 ; CHECK-NEXT: ret i64 %v
 
@@ -16,7 +16,7 @@ define i64 @test2(ptr %p) {
   ret i64 %v
 }
 
-; CHECK: define i64 @test2(ptr %p) {
+; CHECK: define internal i64 @test2(ptr %p) {
 ; CHECK-NEXT: store i64 0, ptr %p
 ; CHECK-NEXT: %v = load i64, ptr %p, align 8
 ; CHECK-NEXT: ret i64 %v
@@ -27,6 +27,6 @@ define i32 @test3(ptr %0)  {
   ret i32 %3
 }
 
-; CHECK: define i32 @test3(ptr %0)  {
+; CHECK: define internal i32 @test3(ptr %0)  {
 ; CHECK-NEXT: %2 = getelementptr i32, ptr %0, i32 4
 ; CHECK-NEXT: %3 = load i32, ptr %2

Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeDirectXTarget() {`
`48`	`48`	`initializeDXILTranslateMetadataPass(*PR);`
`49`	`49`	`initializeDXILResourceMDWrapperPass(*PR);`
`50`	`50`	`initializeShaderFlagsAnalysisWrapperPass(*PR);`
	`51`	`+ initializeDXILFinalizeLinkageLegacyPass(*PR);`
`51`	`52`	`}`
`52`	`53`
`53`	`54`	`class DXILTargetObjectFile : public TargetLoweringObjectFile {`
`@@ -79,6 +80,7 @@ class DirectXPassConfig : public TargetPassConfig {`
`79`	`80`	`void addCodeGenPrepare() override {`
`80`	`81`	`addPass(createDXILIntrinsicExpansionLegacyPass());`
`81`	`82`	`addPass(createDXILOpLoweringLegacyPass());`
	`83`	`+ addPass(createDXILFinalizeLinkageLegacyPass());`
`82`	`84`	`addPass(createDXILTranslateMetadataPass());`
`83`	`85`	`addPass(createDXILPrepareModulePass());`
`84`	`86`	`}`