[SPIRV] Add support vk::ext_execution_mode (#4086)

Add support
* vk::ext_execution_mode
* [[vk::ext_storage_class(uint storage_class)]]

This is part of PRs for the
#3919

Author: jiaolu

include/dxc/HlslIntrinsicOp.h

@@ -226,6 +226,9 @@ enum class IntrinsicOp {  IOP_AcceptHitAndEndSearch,
 #endif // ENABLE_SPIRV_CODEGEN
 #ifdef ENABLE_SPIRV_CODEGEN
   IOP_VkReadClock,
+#endif // ENABLE_SPIRV_CODEGEN
+#ifdef ENABLE_SPIRV_CODEGEN
+  IOP_Vkext_execution_mode,
 #endif // ENABLE_SPIRV_CODEGEN
   MOP_Append,
   MOP_RestartStrip,

lib/HLSL/HLOperationLower.cpp

@@ -5662,6 +5662,7 @@ IntrinsicLower gLowerTable[] = {
 #ifdef ENABLE_SPIRV_CODEGEN
     { IntrinsicOp::IOP_VkReadClock, UnsupportedVulkanIntrinsic, DXIL::OpCode::NumOpCodes },
     { IntrinsicOp::IOP_VkRawBufferLoad, UnsupportedVulkanIntrinsic, DXIL::OpCode::NumOpCodes },
+    { IntrinsicOp::IOP_Vkext_execution_mode, UnsupportedVulkanIntrinsic, DXIL::OpCode::NumOpCodes },
 #endif // ENABLE_SPIRV_CODEGEN
     {IntrinsicOp::MOP_Append, StreamOutputLower, DXIL::OpCode::EmitStream},
     {IntrinsicOp::MOP_RestartStrip, StreamOutputLower, DXIL::OpCode::CutStream},

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

@@ -1050,6 +1050,14 @@ def VKExtensionExt : InheritableAttr {
   let Documentation = [Undocumented];
 }
 
+def VKStorageClassExt : InheritableAttr {
+  let Spellings = [CXX11<"vk", "ext_storage_class">];
+  let Subjects = SubjectList<[Var, ParmVar], ErrorDiag>;
+  let Args = [UnsignedArgument<"stclass">];
+  let LangOpts = [SPIRV];
+  let Documentation = [Undocumented];
+}
+
 // Global variables that are of struct type
 def StructGlobalVar : SubsetSubject<Var, [{S->hasGlobalStorage() && S->getType()->isStructureType()}]>;
 

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

@@ -554,8 +554,10 @@ class SpirvBuilder {
                                      llvm::StringRef content = "");
 
   /// \brief Adds an execution mode to the module under construction.
-  inline void addExecutionMode(SpirvFunction *entryPoint, spv::ExecutionMode em,
-                               llvm::ArrayRef<uint32_t> params, SourceLocation);
+  inline SpirvInstruction *addExecutionMode(SpirvFunction *entryPoint,
+                                            spv::ExecutionMode em,
+                                            llvm::ArrayRef<uint32_t> params,
+                                            SourceLocation);
 
   /// \brief Adds an OpModuleProcessed instruction to the module under
   /// construction.
@@ -854,12 +856,15 @@ SpirvBuilder::setDebugSource(uint32_t major, uint32_t minor,
   return mainSource->getFile();
 }
 
-void SpirvBuilder::addExecutionMode(SpirvFunction *entryPoint,
-                                    spv::ExecutionMode em,
-                                    llvm::ArrayRef<uint32_t> params,
-                                    SourceLocation loc) {
-  mod->addExecutionMode(
-      new (context) SpirvExecutionMode(loc, entryPoint, em, params, false));
+SpirvInstruction *
+SpirvBuilder::addExecutionMode(SpirvFunction *entryPoint, spv::ExecutionMode em,
+                               llvm::ArrayRef<uint32_t> params,
+                               SourceLocation loc) {
+  auto mode =
+      new (context) SpirvExecutionMode(loc, entryPoint, em, params, false);
+  mod->addExecutionMode(mode);
+
+  return mode;
 }
 
 } // end namespace spirv

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -1659,8 +1659,11 @@ void SpirvEmitter::doVarDecl(const VarDecl *decl) {
       needsLegalization = true;
   }
 
-  if (var != nullptr) {
+  if (var != nullptr && decl->hasAttrs()) {
     declIdMapper.decorateVariableWithIntrinsicAttrs(decl, var);
+    if (auto attr = decl->getAttr<VKStorageClassExtAttr>()) {
+      var->setStorageClass(static_cast<spv::StorageClass>(attr->getStclass()));
+    }
   }
 
   // All variables that are of opaque struct types should request legalization.
@@ -7637,6 +7640,9 @@ SpirvEmitter::processIntrinsicCallExpr(const CallExpr *callExpr) {
   case hlsl::IntrinsicOp::IOP_VkRawBufferLoad:
     retVal = processRawBufferLoad(callExpr);
     break;
+  case hlsl::IntrinsicOp::IOP_Vkext_execution_mode:
+    retVal = processIntrinsicExecutionMode(callExpr);
+    break;
   case hlsl::IntrinsicOp::IOP_saturate:
     retVal = processIntrinsicSaturate(callExpr);
     break;
@@ -12569,6 +12575,32 @@ SpirvInstruction *SpirvEmitter::processRawBufferLoad(const CallExpr *callExpr) {
   return loadInst;
 }
 
+SpirvInstruction *
+SpirvEmitter::processIntrinsicExecutionMode(const CallExpr *expr) {
+  llvm::SmallVector<uint32_t, 2> execModesParams;
+  uint32_t exeMode = 0;
+  const auto args = expr->getArgs();
+  for (uint32_t i = 0; i < expr->getNumArgs(); ++i) {
+    SpirvConstantInteger *argInst =
+        dyn_cast<SpirvConstantInteger>(doExpr(args[i]));
+    if (argInst == nullptr) {
+      emitError("argument should be constant interger", expr->getExprLoc());
+      return nullptr;
+    }
+    unsigned argInteger = argInst->getValue().getZExtValue();
+    if (i > 0)
+      execModesParams.push_back(argInteger);
+    else
+      exeMode = argInteger;
+  }
+  assert(entryFunction != nullptr);
+  assert(exeMode != 0);
+
+  return spvBuilder.addExecutionMode(entryFunction,
+                                     static_cast<spv::ExecutionMode>(exeMode),
+                                     execModesParams, expr->getExprLoc());
+}
+
 bool SpirvEmitter::spirvToolsValidate(std::vector<uint32_t> *mod,
                                       std::string *messages) {
   spvtools::SpirvTools tools(featureManager.getTargetEnv());

tools/clang/lib/SPIRV/SpirvEmitter.h

@@ -599,6 +599,8 @@ class SpirvEmitter : public ASTConsumer {
 
   /// Custom intrinsic to support basic buffer_reference use case
   SpirvInstruction *processRawBufferLoad(const CallExpr *callExpr);
+  /// Process vk::ext_execution_mode intrinsic
+  SpirvInstruction *processIntrinsicExecutionMode(const CallExpr *expr);
 
 private:
   /// Returns the <result-id> for constant value 0 of the given type.

tools/clang/lib/Sema/SemaHLSL.cpp

@@ -12090,6 +12090,11 @@ void hlsl::HandleDeclAttributeForHLSL(Sema &S, Decl *D, const AttributeList &A,
         ValidateAttributeStringArg(S, A, nullptr, 1),
         A.getAttributeSpellingListIndex());
     break;
+  case AttributeList::AT_VKStorageClassExt:
+    declAttr = ::new (S.Context) VKStorageClassExtAttr(
+        A.getRange(), S.Context, unsigned(ValidateAttributeIntArg(S, A)),
+        A.getAttributeSpellingListIndex());
+    break;
   default:
     Handled = false;
     return;

tools/clang/test/CodeGenSPIRV/spv.intrinsicExecutionMode.hlsl

@@ -0,0 +1,14 @@
+// RUN: %dxc -T ps_6_0 -E main -spirv
+
+
+// CHECK: OpExecutionMode {{%\w+}} StencilRefReplacingEXT
+// CHECK: OpExecutionMode {{%\w+}} SubgroupSize 32
+// CHECK: OpDecorate {{%\w+}} BuiltIn FragStencilRefEXT
+
+[[vk::ext_decorate(11, 5014)]]
+int main() : SV_Target0 {
+  
+  vk::ext_execution_mode(/*StencilRefReplacingEXT*/5027);
+  vk::ext_execution_mode(/*SubgroupSize*/35, 32);
+  return 3;
+}

tools/clang/test/CodeGenSPIRV/spv.intrinsicStorageClass.hlsl

@@ -0,0 +1,15 @@
+// RUN: %dxc -T ps_6_0 -E main -spirv -Vd
+
+//CHECK: [[payloadTy:%\w+]] = OpTypeStruct %v4float
+//CHECK-NEXT: [[payloadTyPtr:%\w+]] = OpTypePointer RayPayloadNV [[payloadTy]]
+//CHECK: [[crossTy:%\w+]] = OpTypePointer CrossWorkgroup %int
+//CHECK: {{%\w+}} = OpVariable [[payloadTyPtr]] RayPayloadNV
+//CHECK: {{%\w+}} = OpVariable [[crossTy]] CrossWorkgroup
+
+[[vk::ext_storage_class(/*RayPayloadNV*/5338)]]
+float4 payload;
+
+int main() : SV_Target0 {
+  [[vk::ext_storage_class(/* CrossWorkgroup */ 5)]] int foo = 3;
+  return foo;
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -1346,6 +1346,8 @@ TEST_F(FileTest, IntrinsicsSpirv) {
   runFileTest("spv.intrinsicInstruction.hlsl");
   runFileTest("spv.intrinsicLiteral.hlsl");
   runFileTest("spv.intrinsicDecorate.hlsl", Expect::Success, false);
+  runFileTest("spv.intrinsicExecutionMode.hlsl", Expect::Success, false);
+  runFileTest("spv.intrinsicStorageClass.hlsl", Expect::Success, false);
   runFileTest("spv.intrinsic.reference.error.hlsl", Expect::Failure);
 }
 TEST_F(FileTest, IntrinsicsVkReadClock) {