From 2c8bd4a729a042ed823547ea0eacad429c50d7bf Mon Sep 17 00:00:00 2001
From: Helena Kotas <hekotas@microsoft.com>
Date: Tue, 13 Aug 2024 11:09:36 -0700
Subject: [PATCH] [HLSL] Mark exported functions with "hlsl.export" attribute
 (#102275)

Marks exported functions with `"hlsl.export"` attribute. This
information will be later used by DXILFinalizeLinkage pass (coming soon)
to determine which functions should have internal linkage in the final
DXIL code.

Related to #llvm/llvm-project#92071
---
 clang/lib/CodeGen/CGHLSLRuntime.cpp   |  8 ++++++++
 clang/lib/CodeGen/CGHLSLRuntime.h     |  2 +-
 clang/lib/CodeGen/CodeGenFunction.cpp | 10 +++++++---
 clang/test/CodeGenHLSL/export.hlsl    | 10 ++++++----
 4 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/clang/lib/CodeGen/CGHLSLRuntime.cpp b/clang/lib/CodeGen/CGHLSLRuntime.cpp
index a2c3e76f77b7c7..5e59b0f00ebd64 100644
--- a/clang/lib/CodeGen/CGHLSLRuntime.cpp
+++ b/clang/lib/CodeGen/CGHLSLRuntime.cpp
@@ -412,6 +412,14 @@ void CGHLSLRuntime::emitEntryFunction(const FunctionDecl *FD,
   B.CreateRetVoid();
 }
 
+void CGHLSLRuntime::setHLSLFunctionAttributes(const FunctionDecl *FD,
+                                              llvm::Function *Fn) {
+  if (FD->isInExportDeclContext()) {
+    const StringRef ExportAttrKindStr = "hlsl.export";
+    Fn->addFnAttr(ExportAttrKindStr);
+  }
+}
+
 static void gatherFunctions(SmallVectorImpl<Function *> &Fns, llvm::Module &M,
                             bool CtorOrDtor) {
   const auto *GV =
diff --git a/clang/lib/CodeGen/CGHLSLRuntime.h b/clang/lib/CodeGen/CGHLSLRuntime.h
index 527e73a0e21fc4..590d388c456c45 100644
--- a/clang/lib/CodeGen/CGHLSLRuntime.h
+++ b/clang/lib/CodeGen/CGHLSLRuntime.h
@@ -124,7 +124,7 @@ class CGHLSLRuntime {
   void setHLSLEntryAttributes(const FunctionDecl *FD, llvm::Function *Fn);
 
   void emitEntryFunction(const FunctionDecl *FD, llvm::Function *Fn);
-  void setHLSLFunctionAttributes(llvm::Function *, const FunctionDecl *);
+  void setHLSLFunctionAttributes(const FunctionDecl *FD, llvm::Function *Fn);
 
 private:
   void addBufferResourceAnnotation(llvm::GlobalVariable *GV,
diff --git a/clang/lib/CodeGen/CodeGenFunction.cpp b/clang/lib/CodeGen/CodeGenFunction.cpp
index 7adc5a48b7c47b..eff8c9f5694084 100644
--- a/clang/lib/CodeGen/CodeGenFunction.cpp
+++ b/clang/lib/CodeGen/CodeGenFunction.cpp
@@ -1228,9 +1228,13 @@ void CodeGenFunction::StartFunction(GlobalDecl GD, QualType RetTy,
   if (getLangOpts().OpenMP && CurCodeDecl)
     CGM.getOpenMPRuntime().emitFunctionProlog(*this, CurCodeDecl);
 
-  // Handle emitting HLSL entry functions.
-  if (D && D->hasAttr<HLSLShaderAttr>())
-    CGM.getHLSLRuntime().emitEntryFunction(FD, Fn);
+  if (FD && getLangOpts().HLSL) {
+    // Handle emitting HLSL entry functions.
+    if (FD->hasAttr<HLSLShaderAttr>()) {
+      CGM.getHLSLRuntime().emitEntryFunction(FD, Fn);
+    }
+    CGM.getHLSLRuntime().setHLSLFunctionAttributes(FD, Fn);
+  }
 
   EmitFunctionProlog(*CurFnInfo, CurFn, Args);
 
diff --git a/clang/test/CodeGenHLSL/export.hlsl b/clang/test/CodeGenHLSL/export.hlsl
index 53f603739e329b..63f9f9066f9277 100644
--- a/clang/test/CodeGenHLSL/export.hlsl
+++ b/clang/test/CodeGenHLSL/export.hlsl
@@ -2,19 +2,21 @@
 // RUN:   dxil-pc-shadermodel6.3-library %s \
 // RUN:   -emit-llvm -disable-llvm-passes -o - | FileCheck %s
 
-// CHECK: define void @"?f1@@YAXXZ"()
+// CHECK: define void @"?f1@@YAXXZ"() [[Attr:\#[0-9]+]]
 export void f1() {
 }
 
-// CHECK: define void @"?f2@MyNamespace@@YAXXZ"()
+// CHECK: define void @"?f2@MyNamespace@@YAXXZ"() [[Attr]]
 namespace MyNamespace {
   export void f2() {
   }
 }
 
 export {
-// CHECK: define void @"?f3@@YAXXZ"()
-// CHECK: define void @"?f4@@YAXXZ"()    
+// CHECK: define void @"?f3@@YAXXZ"() [[Attr]]
+// CHECK: define void @"?f4@@YAXXZ"() [[Attr]]
     void f3() {}
     void f4() {}
 }
+
+// CHECK: attributes [[Attr]] = { {{.*}} "hlsl.export" {{.*}} }