Allow cap and extention attribute on member variables and type aliases (microsoft#6742)

This implements the changes proposed in
microsoft/hlsl-specs#270.

Author: s-perron

tools/clang/include/clang/Basic/Attr.td

@@ -1317,7 +1317,7 @@ def VKBinding : InheritableAttr {
 
 def VKCapabilityExt : InheritableAttr {
   let Spellings = [CXX11<"vk", "ext_capability">];
-  let Subjects = SubjectList<[Function, Var], ErrorDiag>;
+  let Subjects = SubjectList<[Function, Var, ParmVar, Field, TypedefName], ErrorDiag>;
   let Args = [IntArgument<"capability">];
   let LangOpts = [SPIRV];
   let Documentation = [Undocumented];
@@ -1357,7 +1357,7 @@ def VKDecorateStringExt : InheritableAttr {
 
 def VKExtensionExt : InheritableAttr {
   let Spellings = [CXX11<"vk", "ext_extension">];
-  let Subjects = SubjectList<[Function, Var], ErrorDiag>;
+  let Subjects = SubjectList<[Function, Var, ParmVar, Field, TypedefName], ErrorDiag>;
   let Args = [StringArgument<"name">];
   let LangOpts = [SPIRV];
   let Documentation = [Undocumented];

tools/clang/include/clang/SPIRV/FeatureManager.h

@@ -76,7 +76,7 @@ class FeatureManager {
   /// Allows all extensions to be used in CodeGen.
   void allowAllKnownExtensions();
 
-  /// Rqeusts the given extension for translating the given target feature at
+  /// Requests the given extension for translating the given target feature at
   /// the given source location. Emits an error if the given extension is not
   /// permitted to use.
   bool requestExtension(Extension, llvm::StringRef target, SourceLocation);

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

@@ -764,13 +764,8 @@ class SpirvBuilder {
                                         SpirvInstruction *v);
   SpirvInstruction *getPerVertexStgInput(SpirvInstruction *k);
 
-public:
   std::vector<uint32_t> takeModule();
 
-protected:
-  /// Only friend classes are allowed to add capability/extension to the module
-  /// under construction.
-
   /// \brief Adds the given capability to the module under construction due to
   /// the feature used at the given source location.
   inline void requireCapability(spv::Capability, SourceLocation loc = {});

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -968,13 +968,6 @@ DeclResultIdMapper::getDeclSpirvInfo(const ValueDecl *decl) const {
 SpirvInstruction *DeclResultIdMapper::getDeclEvalInfo(const ValueDecl *decl,
                                                       SourceLocation loc,
                                                       SourceRange range) {
-  if (decl->hasAttr<VKExtensionExtAttr>() ||
-      decl->hasAttr<VKCapabilityExtAttr>()) {
-    theEmitter.createSpirvIntrInstExt(decl->getAttrs(), QualType(),
-                                      /* spvArgs */ {}, /* isInst */ false,
-                                      loc);
-  }
-
   if (auto *builtinAttr = decl->getAttr<VKExtBuiltinInputAttr>()) {
     return getBuiltinVar(spv::BuiltIn(builtinAttr->getBuiltInID()),
                          decl->getType(), spv::StorageClass::Input, loc);
@@ -1640,6 +1633,18 @@ DeclResultIdMapper::createShaderRecordBuffer(const HLSLBufferDecl *decl,
   return bufferVar;
 }
 
+void DeclResultIdMapper::recordsSpirvTypeAlias(const Decl *decl) {
+  auto *typedefDecl = dyn_cast<TypedefNameDecl>(decl);
+  if (!typedefDecl)
+    return;
+
+  if (!typedefDecl->hasAttr<VKCapabilityExtAttr>() &&
+      !typedefDecl->hasAttr<VKExtensionExtAttr>())
+    return;
+
+  typeAliasesWithAttributes.push_back(typedefDecl);
+}
+
 void DeclResultIdMapper::createGlobalsCBuffer(const VarDecl *var) {
   if (astDecls.count(var) != 0)
     return;
@@ -4776,5 +4781,21 @@ void DeclResultIdMapper::storeOutStageVarsToStorage(
   }
 }
 
+void DeclResultIdMapper::registerCapabilitiesAndExtensionsForType(
+    const TypedefType *type) {
+  for (const auto *decl : typeAliasesWithAttributes) {
+    if (type == decl->getTypeForDecl()) {
+      for (auto *attribute : decl->specific_attrs<VKExtensionExtAttr>()) {
+        clang::StringRef extensionName = attribute->getName();
+        spvBuilder.requireExtension(extensionName, decl->getLocation());
+      }
+      for (auto *attribute : decl->specific_attrs<VKCapabilityExtAttr>()) {
+        spv::Capability cap = spv::Capability(attribute->getCapability());
+        spvBuilder.requireCapability(cap, decl->getLocation());
+      }
+    }
+  }
+}
+
 } // end namespace spirv
 } // end namespace clang

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -376,6 +376,10 @@ class DeclResultIdMapper {
   SpirvVariable *createShaderRecordBuffer(const HLSLBufferDecl *decl,
                                           ContextUsageKind kind);
 
+  // Records the TypedefDecl or TypeAliasDecl of vk::SpirvType so that any
+  // required capabilities and extensions can be added if the type is used.
+  void recordsSpirvTypeAlias(const Decl *decl);
+
 private:
   /// The struct containing SPIR-V information of a AST Decl.
   struct DeclSpirvInfo {
@@ -569,6 +573,12 @@ class DeclResultIdMapper {
 
   spv::ExecutionMode getInterlockExecutionMode();
 
+  /// 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
+  /// using `recordSpirvTypeAlias`.
+  void registerCapabilitiesAndExtensionsForType(const TypedefType *type);
+
 private:
   /// \brief Wrapper method to create a fatal error message and report it
   /// in the diagnostic engine associated with this consumer.
@@ -1058,6 +1068,8 @@ class DeclResultIdMapper {
 
   uint32_t perspBaryCentricsIndex, noPerspBaryCentricsIndex;
 
+  llvm::SmallVector<const TypedefNameDecl *, 4> typeAliasesWithAttributes;
+
 public:
   /// The gl_PerVertex structs for both input and output
   GlPerVertex glPerVertex;

tools/clang/lib/SPIRV/LowerTypeVisitor.cpp

@@ -1085,6 +1085,19 @@ LowerTypeVisitor::lowerStructFields(const RecordDecl *decl,
           field->getBitWidthValue(field->getASTContext());
     }
 
+    if (field->hasAttrs()) {
+      for (auto &attr : field->getAttrs()) {
+        if (auto capAttr = dyn_cast<VKCapabilityExtAttr>(attr)) {
+          spvBuilder.requireCapability(
+              static_cast<spv::Capability>(capAttr->getCapability()),
+              capAttr->getLocation());
+        } else if (auto extAttr = dyn_cast<VKExtensionExtAttr>(attr)) {
+          spvBuilder.requireExtension(extAttr->getName(),
+                                      extAttr->getLocation());
+        }
+      }
+    }
+
     fields.push_back(HybridStructType::FieldInfo(
         field->getType(), field->getName(),
         /*vkoffset*/ field->getAttr<VKOffsetAttr>(),

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -986,8 +986,7 @@ void SpirvEmitter::HandleTranslationUnit(ASTContext &context) {
 }
 
 void SpirvEmitter::doDecl(const Decl *decl) {
-  if (isa<EmptyDecl>(decl) || isa<TypedefDecl>(decl) ||
-      isa<TypeAliasDecl>(decl) || isa<TypeAliasTemplateDecl>(decl) ||
+  if (isa<EmptyDecl>(decl) || isa<TypeAliasTemplateDecl>(decl) ||
       isa<VarTemplateDecl>(decl))
     return;
 
@@ -1016,6 +1015,8 @@ void SpirvEmitter::doDecl(const Decl *decl) {
   } else if (const auto *classTemplateDecl =
                  dyn_cast<ClassTemplateDecl>(decl)) {
     doClassTemplateDecl(classTemplateDecl);
+  } else if (isa<TypedefNameDecl>(decl)) {
+    declIdMapper.recordsSpirvTypeAlias(decl);
   } else if (isa<FunctionTemplateDecl>(decl)) {
     // nothing to do.
   } else if (isa<UsingDecl>(decl)) {
@@ -1696,6 +1697,27 @@ bool SpirvEmitter::validateVKAttributes(const NamedDecl *decl) {
   return success;
 }
 
+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());
+  }
+
+  // Now check for any capabilities or extensions that are part of the type.
+  const TypedefType *type = dyn_cast<TypedefType>(varDecl->getType());
+  if (!type)
+    return;
+
+  declIdMapper.registerCapabilitiesAndExtensionsForType(type);
+}
+
 void SpirvEmitter::doHLSLBufferDecl(const HLSLBufferDecl *bufferDecl) {
   // This is a cbuffer/tbuffer decl.
   // Check and emit warnings for member intializers which are not
@@ -1847,6 +1869,8 @@ void SpirvEmitter::doVarDecl(const VarDecl *decl) {
     return;
   }
 
+  registerCapabilitiesAndExtensionsForVarDecl(decl);
+
   // Handle vk::ext_builtin_input and vk::ext_builtin_input by using
   // getBuiltinVar to create the builtin and validate the storage class
   if (decl->hasAttr<VKExtBuiltinInputAttr>()) {

tools/clang/lib/SPIRV/SpirvEmitter.h

@@ -412,6 +412,10 @@ class SpirvEmitter : public ASTConsumer {
   /// errors are found.
   bool validateVKAttributes(const NamedDecl *decl);
 
+  /// Records any Spir-V capabilities and extensions for the given varDecl so
+  /// they will be added to the SPIR-V module.
+  void registerCapabilitiesAndExtensionsForVarDecl(const VarDecl *varDecl);
+
 private:
   /// Converts the given value from the bitwidth of 'fromType' to the bitwidth
   /// of 'toType'. If the two have the same bitwidth, returns the value itself.

tools/clang/test/CodeGenSPIRV/spv.inline.capability.hlsl

@@ -1,12 +1,19 @@
 // RUN: %dxc -T ps_6_0 -E main -fcgl -Vd %s -spirv | FileCheck %s
 
-// CHECK: OpCapability Int8
-// CHECK: OpCapability SampleMaskPostDepthCoverage
+// CHECK-DAG: OpCapability Int8
+// CHECK-DAG: OpCapability SampleMaskPostDepthCoverage
+// CHECK-DAG: OpCapability WorkgroupMemoryExplicitLayoutKHR
 
-[[vk::ext_capability(/* SampleMaskPostDepthCoverageCapability */ 4447)]]
-int val;
 
+// Test that the capability on a typedef is added to the module.
 [[vk::ext_capability(/* Int8 */ 39)]]
+typedef vk::SpirvType</* OpTypeInt */ 21, 8, 8, vk::Literal<vk::integral_constant<uint, 8> >, vk::Literal<vk::integral_constant<bool, false> > > uint8_t;
+
+// Test that the capability on a variable is added to the module.
+[[vk::ext_capability(/* SampleMaskPostDepthCoverageCapability */ 4447)]]
+uint8_t val;
+
+// Test that the capability on the entry point is added to the module.
+[[vk::ext_capability(/* WorkgroupMemoryExplicitLayoutKHR */ 4428)]]
 void main() {
-  int local = val;
 }

tools/clang/test/CodeGenSPIRV/spv.inline.extension.hlsl

@@ -1,13 +1,22 @@
 // RUN: %dxc -T ps_6_0 -E main -fcgl %s -spirv | FileCheck %s
 
-// CHECK: OpExtension "entry_point_extension"
-// CHECK: OpExtension "another_extension"
-// CHECK: OpExtension "some_extension"
+// CHECK-DAG: OpExtension "entry_point_extension"
+// CHECK-DAG: OpExtension "another_extension"
+// CHECK-DAG: OpExtension "some_extension"
+// CHECK-DAG: OpExtension "ext_on_field1"
+// CHECK-DAG: OpExtension "ext_on_field2"
 
 [[vk::ext_extension("some_extension"), vk::ext_extension("another_extension")]]
 int val;
 
+struct T
+{
+  [[vk::ext_extension("ext_on_field1"), vk::ext_extension("ext_on_field2")]]
+  int val;
+};
+
 [[vk::ext_extension("entry_point_extension")]]
 void main() {
-  int local = val;
+  T t;
+  int local = val+t.val;
 }