[SPIR-V] Honor inline spir-v attributes on functions

tools/clang/include/clang/SPIRV/SpirvBuilder.h

@@ -801,6 +801,12 @@ class SpirvBuilder {
   void decorateWithLiterals(SpirvInstruction *targetInst, unsigned decorate,
                             llvm::ArrayRef<unsigned> literals, SourceLocation);
 
+  /// \brief Decorates the given function with information from VKDecorateExt.
+  /// Used for [[vk::ext_decorate(...)]] placed directly on a function, which
+  /// emits an OpDecorate targeting the OpFunction.
+  void decorateWithLiterals(SpirvFunction *targetFunc, unsigned decorate,
+                            llvm::ArrayRef<unsigned> literals, SourceLocation);
+
   /// \brief Decorates the given target with result ids of SPIR-V
   /// instructions.
   void decorateWithIds(SpirvInstruction *targetInst, unsigned decorate,

tools/clang/include/clang/SPIRV/SpirvModule.h

@@ -39,7 +39,7 @@ struct DecorationComparisonInfo {
   static inline SpirvDecoration *getEmptyKey() { return nullptr; }
   static inline SpirvDecoration *getTombstoneKey() { return nullptr; }
   static unsigned getHashValue(const SpirvDecoration *decor) {
-    return llvm::hash_combine(decor->getTarget(),
+    return llvm::hash_combine(decor->getTarget(), decor->getTargetFunc(),
                               static_cast<uint32_t>(decor->getDecoration()));
   }
   static bool isEqual(SpirvDecoration *LHS, SpirvDecoration *RHS) {

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -1826,6 +1826,17 @@ SpirvFunction *DeclResultIdMapper::getOrRegisterFn(const FunctionDecl *fn) {
                                spv::LinkageType::Export, fn->getLocation());
   }
 
+  // Honor inline-SPIR-V attributes placed directly on a function. The
+  // entry-point path handles these only for entry functions, and the
+  // vk::ext_instruction path only for functions lowered to an instruction, so a
+  // plain function was previously skipped and these attributes silently
+  // dropped. These reuse the same helpers as the variable/parameter paths:
+  //   [[vk::ext_decorate(d, ...)]]  -> OpDecorate targeting the OpFunction
+  //   [[vk::ext_capability(c)]]     -> OpCapability for the module
+  //   [[vk::ext_extension("...")]]  -> OpExtension for the module
+  decorateWithIntrinsicAttrs(fn, spirvFunction);
+  registerCapabilitiesAndExtensionsForDecl(fn);
+
   // No need to dereference to get the pointer. Function returns that are
   // stand-alone aliases are already pointers to values. All other cases should
   // be normal rvalues.
@@ -5047,6 +5058,34 @@ void DeclResultIdMapper::decorateWithIntrinsicAttrs(
   }
 }
 
+void DeclResultIdMapper::decorateWithIntrinsicAttrs(const NamedDecl *decl,
+                                                    SpirvFunction *targetFunc) {
+  if (!decl->hasAttrs())
+    return;
+
+  for (auto &attr : decl->getAttrs()) {
+    if (auto *decoAttr = dyn_cast<VKDecorateExtAttr>(attr)) {
+      spvBuilder.decorateWithLiterals(
+          targetFunc, decoAttr->getDecorate(),
+          {decoAttr->literals_begin(), decoAttr->literals_end()},
+          decl->getLocation());
+      continue;
+    }
+    // The id/string forms decorate a SpirvInstruction target; there is no
+    // SpirvFunction-target equivalent yet, so reject rather than silently drop.
+    if (isa<VKDecorateIdExtAttr>(attr)) {
+      emitError("vk::ext_decorate_id is not supported on functions",
+                decl->getLocation());
+      continue;
+    }
+    if (isa<VKDecorateStringExtAttr>(attr)) {
+      emitError("vk::ext_decorate_string is not supported on functions",
+                decl->getLocation());
+      continue;
+    }
+  }
+}
+
 void DeclResultIdMapper::decorateStageVarWithIntrinsicAttrs(
     const NamedDecl *decl, StageVar *stageVar, SpirvVariable *varInst) {
   auto checkBuiltInLocationDecoration =
@@ -5136,6 +5175,17 @@ void DeclResultIdMapper::storeOutStageVarsToStorage(
   }
 }
 
+void DeclResultIdMapper::registerCapabilitiesAndExtensionsForDecl(
+    const NamedDecl *decl) {
+  for (auto *attribute : decl->specific_attrs<VKExtensionExtAttr>()) {
+    spvBuilder.requireExtension(attribute->getName(), decl->getLocation());
+  }
+  for (auto *attribute : decl->specific_attrs<VKCapabilityExtAttr>()) {
+    spv::Capability cap = spv::Capability(attribute->getCapability());
+    spvBuilder.requireCapability(cap, decl->getLocation());
+  }
+}
+
 void DeclResultIdMapper::registerCapabilitiesAndExtensionsForType(
     const TypedefType *type) {
   for (const auto *decl : typeAliasesWithAttributes) {

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -576,6 +576,13 @@ class DeclResultIdMapper {
       llvm::function_ref<void(VKDecorateExtAttr *)> extraFunctionForDecoAttr =
           [](VKDecorateExtAttr *) {});
 
+  /// Applies [[vk::ext_decorate(...)]] placed directly on a function to the
+  /// function itself, emitting an OpDecorate that targets the OpFunction. Only
+  /// the literal form is supported; the id/string variants are diagnosed as
+  /// unsupported on functions (they lack a SpirvFunction-target decoration).
+  void decorateWithIntrinsicAttrs(const NamedDecl *decl,
+                                  SpirvFunction *targetFunc);
+
   /// \brief Creates instructions to load the value of output stage variable
   /// defined by outputPatchDecl and store it to ptr. Since the output stage
   /// variable for OutputPatch is an array whose number of elements is the
@@ -589,6 +596,12 @@ class DeclResultIdMapper {
 
   spv::ExecutionMode getInterlockExecutionMode();
 
+  /// Records any Spir-V capabilities and extensions declared directly on the
+  /// given decl via [[vk::ext_capability(...)]] / [[vk::ext_extension(...)]]
+  /// so they will be added to the SPIR-V module. Shared by the variable and
+  /// function declaration paths.
+  void registerCapabilitiesAndExtensionsForDecl(const NamedDecl *decl);
+
   /// Records any Spir-V capabilities and extensions for the given type so
   /// they will be added to the SPIR-V module. The capabilities and extension
   /// required for the type will be sourced from the decls that were recorded

tools/clang/lib/SPIRV/SpirvBuilder.cpp

@@ -1949,6 +1949,16 @@ void SpirvBuilder::decorateWithLiterals(SpirvInstruction *targetInst,
   mod->addDecoration(decor);
 }
 
+void SpirvBuilder::decorateWithLiterals(SpirvFunction *targetFunc,
+                                        unsigned decorate,
+                                        llvm::ArrayRef<unsigned> literals,
+                                        SourceLocation srcLoc) {
+  SpirvDecoration *decor = new (context) SpirvDecoration(
+      srcLoc, targetFunc, static_cast<spv::Decoration>(decorate), literals);
+  assert(decor != nullptr);
+  mod->addDecoration(decor);
+}
+
 void SpirvBuilder::decorateWithIds(SpirvInstruction *targetInst,
                                    unsigned decorate,
                                    llvm::ArrayRef<SpirvInstruction *> ids,

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -1924,16 +1924,8 @@ bool SpirvEmitter::validateVKAttributes(const NamedDecl *decl) {
 
 void SpirvEmitter::registerCapabilitiesAndExtensionsForVarDecl(
     const VarDecl *varDecl) {
-  // First record any extensions that are part of the actual variable
-  // declaration.
-  for (auto *attribute : varDecl->specific_attrs<VKExtensionExtAttr>()) {
-    clang::StringRef extensionName = attribute->getName();
-    spvBuilder.requireExtension(extensionName, varDecl->getLocation());
-  }
-  for (auto *attribute : varDecl->specific_attrs<VKCapabilityExtAttr>()) {
-    spv::Capability cap = spv::Capability(attribute->getCapability());
-    spvBuilder.requireCapability(cap, varDecl->getLocation());
-  }
+  // First record any extensions/capabilities declared on the variable itself.
+  declIdMapper.registerCapabilitiesAndExtensionsForDecl(varDecl);
 
   // Now check for any capabilities or extensions that are part of the type.
   const TypedefType *type = dyn_cast<TypedefType>(varDecl->getType());

tools/clang/lib/SPIRV/SpirvInstruction.cpp

@@ -301,8 +301,9 @@ spv::Op SpirvDecoration::getDecorateStringOpcode(bool isMemberDecoration) {
 }
 
 bool SpirvDecoration::operator==(const SpirvDecoration &that) const {
-  return target == that.target && decoration == that.decoration &&
-         params == that.params && idParams == that.idParams &&
+  return target == that.target && targetFunction == that.targetFunction &&
+         decoration == that.decoration && params == that.params &&
+         idParams == that.idParams &&
          index.hasValue() == that.index.hasValue() &&
          (!index.hasValue() || index.getValue() == that.index.getValue());
 }

tools/clang/test/CodeGenSPIRV/inline-spirv/spv.intrinsicDecorate.function.error.hlsl

@@ -0,0 +1,20 @@
+// RUN: not %dxc -T cs_6_0 -E main -fcgl %s -spirv 2>&1 | FileCheck %s
+
+// vk::ext_decorate_id and vk::ext_decorate_string decorate a value-id target
+// and have no OpFunction-target form, so applying either to a function must be
+// diagnosed rather than silently dropped.
+
+// CHECK: error: vk::ext_decorate_string is not supported on functions
+[[vk::ext_decorate_string(/* UserTypeGOOGLE */ 5636, "myType")]]
+[noinline] uint DecorateString(uint x) { return x; }
+
+// CHECK: error: vk::ext_decorate_id is not supported on functions
+[[vk::ext_decorate_id(/* UniformId */ 27, 13)]]
+[noinline] uint DecorateId(uint x) { return x; }
+
+RWStructuredBuffer<uint> buf;
+
+[numthreads(1, 1, 1)]
+void main(uint3 tid : SV_DispatchThreadID) {
+  buf[0] = DecorateString(tid.x) + DecorateId(tid.x);
+}

tools/clang/test/CodeGenSPIRV/inline-spirv/spv.intrinsicDecorate.function.hlsl

@@ -0,0 +1,32 @@
+// RUN: %dxc -T cs_6_0 -E main -fcgl %s -spirv | FileCheck %s
+
+// Inline-SPIR-V attributes placed directly on a function must be honored:
+// vk::ext_decorate emits an OpDecorate targeting the OpFunction, and
+// vk::ext_capability / vk::ext_extension add the capability / extension to the
+// module. Known SPIR-V enums are used here so the module still validates; the
+// mechanism is identical for vendor-specific decoration numbers.
+
+// CHECK-DAG: OpCapability Int8
+// CHECK-DAG: OpExtension "some_extension"
+
+[[vk::ext_capability(/* Int8 */ 39)]]
+[[vk::ext_extension("some_extension")]]
+[[vk::ext_decorate(/* RelaxedPrecision */ 0)]]
+[noinline] uint Identity(uint x) { return x; }
+
+// Decorate a second function to test the hashing in the module's decoration set.
+[[vk::ext_decorate(/* RelaxedPrecision */ 0)]]
+[noinline] uint Increment(uint x) { return x + 1; }
+
+// CHECK-DAG: OpDecorate %Identity RelaxedPrecision
+// CHECK-DAG: OpDecorate %Increment RelaxedPrecision
+
+// CHECK-DAG: %Identity = OpFunction %uint DontInline
+// CHECK-DAG: %Increment = OpFunction %uint DontInline
+
+RWStructuredBuffer<uint> buf;
+
+[numthreads(1, 1, 1)]
+void main(uint3 tid : SV_DispatchThreadID) {
+  buf[0] = Identity(tid.x) + Increment(tid.x);
+}