[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 ##92071
llvm · Aug 13, 2024 · 2c8bd4a · 2c8bd4a
1 parent 661dda9
commit 2c8bd4a
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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" {{.*}} }