From 475b9f3b2b2bd6f6928d4e9b744a40321d8b78db Mon Sep 17 00:00:00 2001
From: Steven Perron <stevenperron@google.com>
Date: Thu, 5 Sep 2024 13:56:56 -0400
Subject: [PATCH] Add linkage capability based on decorations

We currently add the linkage attribute only if there are no extry
pointer. However, there are cases where we may want an extry point with
an exported function.

This commit add the linkage capability any time it sees the linkage
attribute decoration. Note I do not add the linkage attribute all of the
time and leave it to the capability trimming pass because having this
capability with vulkan shaders is generally illegal, and I don't want
the unoptimized code to fail validation.

Fixes #6738
---
 tools/clang/lib/SPIRV/CapabilityVisitor.cpp         | 13 ++++++-------
 tools/clang/test/CodeGenSPIRV/fn.export.hlsl        |  4 ++--
 .../CodeGenSPIRV/fn.export.with.entrypoint.hlsl     | 13 +++++++++++++
 .../test/CodeGenSPIRV/fn.fixfuncall-linkage.hlsl    |  4 ++--
 4 files changed, 23 insertions(+), 11 deletions(-)
 create mode 100644 tools/clang/test/CodeGenSPIRV/fn.export.with.entrypoint.hlsl

diff --git a/tools/clang/lib/SPIRV/CapabilityVisitor.cpp b/tools/clang/lib/SPIRV/CapabilityVisitor.cpp
index 5e632abc3c..f2d2958542 100644
--- a/tools/clang/lib/SPIRV/CapabilityVisitor.cpp
+++ b/tools/clang/lib/SPIRV/CapabilityVisitor.cpp
@@ -387,6 +387,9 @@ bool CapabilityVisitor::visit(SpirvDecoration *decor) {
 
     break;
   }
+  case spv::Decoration::LinkageAttributes:
+    addCapability(spv::Capability::Linkage);
+    break;
   default:
     break;
   }
@@ -847,16 +850,12 @@ bool CapabilityVisitor::visit(SpirvReadClock *inst) {
 }
 
 bool CapabilityVisitor::visit(SpirvModule *, Visitor::Phase phase) {
-  // If there are no entry-points in the module (hence shaderModel is not set),
-  // add the Linkage capability. This allows library shader models to use
-  // 'export' attribute on functions, and generate an "incomplete/partial"
-  // SPIR-V binary.
-  // ExecutionModel::Max means that no entrypoints exist, therefore we should
-  // add the Linkage Capability.
+  // If there are no entry-points in the module add the Shader capability.
+  // This allows library shader models with no entry pointer and just exported
+  // function. ExecutionModel::Max means that no entrypoints exist.
   if (phase == Visitor::Phase::Done &&
       shaderModel == spv::ExecutionModel::Max) {
     addCapability(spv::Capability::Shader);
-    addCapability(spv::Capability::Linkage);
   }
 
   // SPIRV-Tools now has a pass to trim superfluous capabilities. This means we
diff --git a/tools/clang/test/CodeGenSPIRV/fn.export.hlsl b/tools/clang/test/CodeGenSPIRV/fn.export.hlsl
index cc8941b0cc..04022bf340 100644
--- a/tools/clang/test/CodeGenSPIRV/fn.export.hlsl
+++ b/tools/clang/test/CodeGenSPIRV/fn.export.hlsl
@@ -1,7 +1,7 @@
 // RUN: %dxc -T lib_6_3 -fspv-target-env=universal1.5 -fcgl  %s -spirv | FileCheck %s
 
-// CHECK: OpCapability Shader
-// CHECK: OpCapability Linkage
+// CHECK-DAG: OpCapability Shader
+// CHECK-DAG: OpCapability Linkage
 RWBuffer< float4 > output : register(u1);
 
 // CHECK: OpDecorate %main LinkageAttributes "main" Export
diff --git a/tools/clang/test/CodeGenSPIRV/fn.export.with.entrypoint.hlsl b/tools/clang/test/CodeGenSPIRV/fn.export.with.entrypoint.hlsl
new file mode 100644
index 0000000000..da25ead9c1
--- /dev/null
+++ b/tools/clang/test/CodeGenSPIRV/fn.export.with.entrypoint.hlsl
@@ -0,0 +1,13 @@
+// RUN: %dxc -T as_6_6 -E main -fspv-target-env=vulkan1.3 -fcgl  %s -spirv | FileCheck %s
+
+// CHECK: OpCapability Linkage
+// CHECK: OpDecorate %external_function LinkageAttributes "external_function" Export
+export int external_function() {
+  return 1;
+}
+
+[numthreads(8, 8, 1)]
+void main() {
+  external_function();
+	return;
+}
\ No newline at end of file
diff --git a/tools/clang/test/CodeGenSPIRV/fn.fixfuncall-linkage.hlsl b/tools/clang/test/CodeGenSPIRV/fn.fixfuncall-linkage.hlsl
index 7843412c69..5977fc454a 100644
--- a/tools/clang/test/CodeGenSPIRV/fn.fixfuncall-linkage.hlsl
+++ b/tools/clang/test/CodeGenSPIRV/fn.fixfuncall-linkage.hlsl
@@ -1,7 +1,7 @@
 // RUN: %dxc -T lib_6_3 -fspv-target-env=universal1.5 -fspv-fix-func-call-arguments -O0  %s -spirv | FileCheck %s
 
-// CHECK: OpCapability Shader
-// CHECK: OpCapability Linkage
+// CHECK-DAG: OpCapability Shader
+// CHECK-DAG: OpCapability Linkage
 RWStructuredBuffer< float4 > output : register(u1);
 
 // CHECK: OpDecorate %main LinkageAttributes "main" Export