[spirv] use OpIsHelperInvocation for [[vk::HelperInvocation]] (micros…

…oft#4172)

For Vulkan 1.2 or less, we must not use HelperInvocation BuiltIn
decoration. Instread, for a variable with `[[vk::HelperInvocation]]`
attribute, we have to use `OpIsHelperInvocation` instruction:
- Create a variable (with Input storage class)
- In the module initialization, execute `OpIsHelperInvocation`
  instruction and store the result in the variable

docs/SPIR-V.rst

@@ -263,7 +263,13 @@ language. To support them, ``[[vk::builtin("<builtin>")]]`` is introduced.
 Right now the following ``<builtin>`` are supported:
 
 * ``PointSize``: The GLSL equivalent is ``gl_PointSize``.
-* ``HelperInvocation``: The GLSL equivalent is ``gl_HelperInvocation``.
+* ``HelperInvocation``: For Vulkan 1.3 or above, we use its GLSL equivalent
+  ``gl_HelperInvocation`` and decorate it with ``HelperInvocation`` builtin
+  since Vulkan 1.3 or above supports ``Volatile`` decoration for builtin
+  variables. For Vulkan 1.2 or earlier, we do not create a builtin variable for
+  ``HelperInvocation``. Instead, we create a variable with ``Private`` storage
+  class and set its value as the result of `OpIsHelperInvocationEXT <https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/EXT/SPV_EXT_demote_to_helper_invocation.html#OpIsHelperInvocationEXT>`_
+  instruction.
 * ``BaseVertex``: The GLSL equivalent is ``gl_BaseVertexARB``.
   Need ``SPV_KHR_shader_draw_parameters`` extension.
 * ``BaseInstance``: The GLSL equivalent is ``gl_BaseInstanceARB``.

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

@@ -9,6 +9,7 @@
 #ifndef LLVM_CLANG_SPIRV_SPIRVBUILDER_H
 #define LLVM_CLANG_SPIRV_SPIRVBUILDER_H
 
+#include "clang/SPIRV/FeatureManager.h"
 #include "clang/SPIRV/SpirvBasicBlock.h"
 #include "clang/SPIRV/SpirvContext.h"
 #include "clang/SPIRV/SpirvFunction.h"
@@ -49,7 +50,8 @@ class SpirvBuilder {
   friend class CapabilityVisitor;
 
 public:
-  SpirvBuilder(ASTContext &ac, SpirvContext &c, const SpirvCodeGenOptions &);
+  SpirvBuilder(ASTContext &ac, SpirvContext &c, const SpirvCodeGenOptions &,
+               FeatureManager &featureMgr);
   ~SpirvBuilder() = default;
 
   // Forbid copy construction and assignment
@@ -472,6 +474,10 @@ class SpirvBuilder {
   /// \brief Creates an OpDemoteToHelperInvocation instruction.
   SpirvInstruction *createDemoteToHelperInvocation(SourceLocation);
 
+  /// \brief Creates an OpIsHelperInvocationEXT instruction.
+  SpirvInstruction *createIsHelperInvocationEXT(QualType type,
+                                                SourceLocation loc);
+
   // === SPIR-V Rich Debug Info Creation ===
   SpirvDebugSource *createDebugSource(llvm::StringRef file,
                                       llvm::StringRef text = "");
@@ -550,6 +556,16 @@ class SpirvBuilder {
   ///   3. Use the clone variable in all the places
   SpirvInstruction *initializeCloneVarForFxcCTBuffer(SpirvInstruction *instr);
 
+  /// \brief Adds a module variable with the Private storage class for a
+  /// stage variable with [[vk::builtin(HelperInvocation)]] attribute and
+  /// initializes it as the result of OpIsHelperInvocationEXT instruction.
+  ///
+  /// Note that we must not use it for Vulkan 1.3 or above. Vulkan 1.3 or
+  /// above allows us to use HelperInvocation Builtin decoration for stage
+  /// variables.
+  SpirvVariable *addVarForHelperInvocation(QualType type, bool isPrecise,
+                                           SourceLocation loc);
+
   // === SPIR-V Module Structure ===
   inline void setMemoryModel(spv::AddressingModel, spv::MemoryModel);
 
@@ -773,6 +789,7 @@ class SpirvBuilder {
 private:
   ASTContext &astContext;
   SpirvContext &context; ///< From which we allocate various SPIR-V object
+  FeatureManager &featureManager;
 
   std::unique_ptr<SpirvModule> mod; ///< The current module being built
   SpirvFunction *function;          ///< The current function being built

tools/clang/include/clang/SPIRV/SpirvInstruction.h

@@ -100,6 +100,7 @@ class SpirvInstruction {
     IK_CompositeInsert,          // OpCompositeInsert
     IK_CopyObject,               // OpCopyObject
     IK_DemoteToHelperInvocation, // OpDemoteToHelperInvocation
+    IK_IsHelperInvocationEXT,    // OpIsHelperInvocationEXT
     IK_ExtInst,                  // OpExtInst
     IK_FunctionCall,             // OpFunctionCall
 
@@ -2041,6 +2042,25 @@ class SpirvDemoteToHelperInvocation : public SpirvInstruction {
   bool invokeVisitor(Visitor *v) override;
 };
 
+/// \brief OpIsHelperInvocationEXT instruction.
+/// Result is true if the invocation is currently a helper invocation, otherwise
+/// result is false. An invocation is currently a helper invocation if it was
+/// originally invoked as a helper invocation or if it has been demoted to a
+/// helper invocation by OpDemoteToHelperInvocationEXT.
+class SpirvIsHelperInvocationEXT : public SpirvInstruction {
+public:
+  SpirvIsHelperInvocationEXT(QualType, SourceLocation);
+
+  DEFINE_RELEASE_MEMORY_FOR_CLASS(SpirvIsHelperInvocationEXT)
+
+  // For LLVM-style RTTI
+  static bool classof(const SpirvInstruction *inst) {
+    return inst->getKind() == IK_IsHelperInvocationEXT;
+  }
+
+  bool invokeVisitor(Visitor *v) override;
+};
+
 // A class keeping information of [[vk::ext_instruction(uint opcode,
 // string extended_instruction_set)]] attribute. The attribute allows users to
 // emit an arbitrary SPIR-V instruction by adding it to a function declaration.

tools/clang/include/clang/SPIRV/SpirvVisitor.h

@@ -115,6 +115,7 @@ class Visitor {
   DEFINE_VISIT_METHOD(SpirvArrayLength)
   DEFINE_VISIT_METHOD(SpirvRayTracingOpNV)
   DEFINE_VISIT_METHOD(SpirvDemoteToHelperInvocation)
+  DEFINE_VISIT_METHOD(SpirvIsHelperInvocationEXT)
   DEFINE_VISIT_METHOD(SpirvDebugInfoNone)
   DEFINE_VISIT_METHOD(SpirvDebugSource)
   DEFINE_VISIT_METHOD(SpirvDebugCompilationUnit)

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -659,6 +659,14 @@ bool CapabilityVisitor::visit(SpirvDemoteToHelperInvocation *inst) {
   return true;
 }
 
+bool CapabilityVisitor::visit(SpirvIsHelperInvocationEXT *inst) {
+  addCapability(spv::Capability::DemoteToHelperInvocation,
+                inst->getSourceLocation());
+  addExtension(Extension::EXT_demote_to_helper_invocation,
+               "[[vk::HelperInvocation]]", inst->getSourceLocation());
+  return true;
+}
+
 bool CapabilityVisitor::visit(SpirvReadClock *inst) {
   auto loc = inst->getSourceLocation();
   addCapabilityForType(inst->getResultType(), loc, inst->getStorageClass());

tools/clang/lib/SPIRV/CapabilityVisitor.h

@@ -22,10 +22,10 @@ class SpirvBuilder;
 class CapabilityVisitor : public Visitor {
 public:
   CapabilityVisitor(ASTContext &astCtx, SpirvContext &spvCtx,
-                    const SpirvCodeGenOptions &opts, SpirvBuilder &builder)
+                    const SpirvCodeGenOptions &opts, SpirvBuilder &builder,
+                    FeatureManager &featureMgr)
       : Visitor(opts, spvCtx), spvBuilder(builder),
-        shaderModel(spv::ExecutionModel::Max),
-        featureManager(astCtx.getDiagnostics(), opts) {}
+        shaderModel(spv::ExecutionModel::Max), featureManager(featureMgr) {}
 
   bool visit(SpirvModule *, Phase) override;
 
@@ -39,6 +39,7 @@ class CapabilityVisitor : public Visitor {
   bool visit(SpirvExtInst *) override;
   bool visit(SpirvAtomic *) override;
   bool visit(SpirvDemoteToHelperInvocation *) override;
+  bool visit(SpirvIsHelperInvocationEXT *) override;
   bool visit(SpirvReadClock *) override;
 
   using Visitor::visit;

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -1592,6 +1592,8 @@ DeclResultIdMapper::collectStageVars(SpirvFunction *entryPoint) const {
     if (var.getEntryPoint() && var.getEntryPoint() != entryPoint)
       continue;
     auto *instr = var.getSpirvInstr();
+    if (instr->getStorageClass() == spv::StorageClass::Private)
+      continue;
     if (seenVars.count(instr) == 0) {
       vars.push_back(instr);
       seenVars.insert(instr);
@@ -3145,8 +3147,12 @@ DeclResultIdMapper::invertWIfRequested(SpirvInstruction *position,
 void DeclResultIdMapper::decorateInterpolationMode(const NamedDecl *decl,
                                                    QualType type,
                                                    SpirvVariable *varInstr) {
-  const auto loc = decl->getLocation();
+  if (varInstr->getStorageClass() != spv::StorageClass::Input &&
+      varInstr->getStorageClass() != spv::StorageClass::Output) {
+    return;
+  }
 
+  const auto loc = decl->getLocation();
   if (isUintOrVecMatOfUintType(type) || isSintOrVecMatOfSintType(type) ||
       isBoolOrVecMatOfBoolType(type)) {
     // TODO: Probably we can call hlsl::ValidateSignatureElement() for the
@@ -3272,6 +3278,14 @@ SpirvVariable *DeclResultIdMapper::createSpirvStageVar(
             .Default(BuiltIn::Max);
 
     assert(spvBuiltIn != BuiltIn::Max); // The frontend should guarantee this.
+    if (spvBuiltIn == BuiltIn::HelperInvocation &&
+        !featureManager.isTargetEnvVulkan1p3OrAbove()) {
+      // If [[vk::HelperInvocation]] is used for Vulkan 1.2 or less, we enable
+      // SPV_EXT_demote_to_helper_invocation extension to use
+      // OpIsHelperInvocationEXT instruction.
+      featureManager.allowExtension("SPV_EXT_demote_to_helper_invocation");
+      return spvBuilder.addVarForHelperInvocation(type, isPrecise, srcLoc);
+    }
     return spvBuilder.addStageBuiltinVar(type, sc, spvBuiltIn, isPrecise,
                                          srcLoc);
   }

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -857,6 +857,7 @@ class DeclResultIdMapper {
 private:
   SpirvBuilder &spvBuilder;
   SpirvEmitter &theEmitter;
+  FeatureManager &featureManager;
   const SpirvCodeGenOptions &spirvOptions;
   ASTContext &astContext;
   SpirvContext &spvContext;
@@ -983,8 +984,8 @@ DeclResultIdMapper::DeclResultIdMapper(ASTContext &context,
                                        SpirvEmitter &emitter,
                                        FeatureManager &features,
                                        const SpirvCodeGenOptions &options)
-    : spvBuilder(spirvBuilder), theEmitter(emitter), spirvOptions(options),
-      astContext(context), spvContext(spirvContext),
+    : spvBuilder(spirvBuilder), theEmitter(emitter), featureManager(features),
+      spirvOptions(options), astContext(context), spvContext(spirvContext),
       diags(context.getDiagnostics()), entryFunction(nullptr),
       needsLegalization(false), needsFlatteningCompositeResources(false),
       glPerVertex(context, spirvContext, spirvBuilder) {}

tools/clang/lib/SPIRV/EmitVisitor.cpp

@@ -1350,6 +1350,14 @@ bool EmitVisitor::visit(SpirvDemoteToHelperInvocation *inst) {
   return true;
 }
 
+bool EmitVisitor::visit(SpirvIsHelperInvocationEXT *inst) {
+  initInstruction(inst);
+  curInst.push_back(inst->getResultTypeId());
+  curInst.push_back(getOrAssignResultId<SpirvInstruction>(inst));
+  finalizeInstruction(&mainBinary);
+  return true;
+}
+
 bool EmitVisitor::visit(SpirvDebugInfoNone *inst) {
   initInstruction(inst);
   curInst.push_back(inst->getResultTypeId());

tools/clang/lib/SPIRV/EmitVisitor.h

@@ -47,13 +47,12 @@ class EmitTypeHandler {
 
 public:
   EmitTypeHandler(ASTContext &astCtx, SpirvContext &spvContext,
-                  const SpirvCodeGenOptions &opts,
+                  const SpirvCodeGenOptions &opts, FeatureManager &featureMgr,
                   std::vector<uint32_t> *debugVec,
                   std::vector<uint32_t> *decVec,
                   std::vector<uint32_t> *typesVec,
                   const std::function<uint32_t()> &takeNextIdFn)
-      : astContext(astCtx), context(spvContext),
-        featureManager(astCtx.getDiagnostics(), opts),
+      : astContext(astCtx), context(spvContext), featureManager(featureMgr),
         debugVariableBinary(debugVec), annotationsBinary(decVec),
         typeConstantBinary(typesVec), takeNextIdFunction(takeNextIdFn),
         emittedConstantInts({}), emittedConstantFloats({}),
@@ -200,9 +199,9 @@ class EmitVisitor : public Visitor {
 
 public:
   EmitVisitor(ASTContext &astCtx, SpirvContext &spvCtx,
-              const SpirvCodeGenOptions &opts)
+              const SpirvCodeGenOptions &opts, FeatureManager &featureMgr)
       : Visitor(opts, spvCtx), astContext(astCtx), id(0),
-        typeHandler(astCtx, spvCtx, opts, &debugVariableBinary,
+        typeHandler(astCtx, spvCtx, opts, featureMgr, &debugVariableBinary,
                     &annotationsBinary, &typeConstantBinary,
                     [this]() -> uint32_t { return takeNextId(); }),
         debugMainFileId(0), debugInfoExtInstId(0), debugLineStart(0),
@@ -274,6 +273,7 @@ class EmitVisitor : public Visitor {
   bool visit(SpirvArrayLength *) override;
   bool visit(SpirvRayTracingOpNV *) override;
   bool visit(SpirvDemoteToHelperInvocation *) override;
+  bool visit(SpirvIsHelperInvocationEXT *) override;
   bool visit(SpirvRayQueryOpKHR *) override;
   bool visit(SpirvReadClock *) override;
   bool visit(SpirvRayTracingTerminateOpKHR *) override;

tools/clang/lib/SPIRV/RemoveBufferBlockVisitor.h

@@ -22,8 +22,9 @@ class SpirvContext;
 class RemoveBufferBlockVisitor : public Visitor {
 public:
   RemoveBufferBlockVisitor(ASTContext &astCtx, SpirvContext &spvCtx,
-                           const SpirvCodeGenOptions &opts)
-      : Visitor(opts, spvCtx), featureManager(astCtx.getDiagnostics(), opts) {}
+                           const SpirvCodeGenOptions &opts,
+                           FeatureManager &featureMgr)
+      : Visitor(opts, spvCtx), featureManager(featureMgr) {}
 
   bool visit(SpirvModule *, Phase) override;
   bool visit(SpirvFunction *, Phase) override;

tools/clang/lib/SPIRV/SpirvBuilder.cpp

@@ -25,11 +25,13 @@ namespace clang {
 namespace spirv {
 
 SpirvBuilder::SpirvBuilder(ASTContext &ac, SpirvContext &ctx,
-                           const SpirvCodeGenOptions &opt)
-    : astContext(ac), context(ctx), mod(llvm::make_unique<SpirvModule>()),
-      function(nullptr), moduleInit(nullptr), moduleInitInsertPoint(nullptr),
-      spirvOptions(opt), builtinVars(), debugNone(nullptr),
-      nullDebugExpr(nullptr), stringLiterals() {}
+                           const SpirvCodeGenOptions &opt,
+                           FeatureManager &featureMgr)
+    : astContext(ac), context(ctx), featureManager(featureMgr),
+      mod(llvm::make_unique<SpirvModule>()), function(nullptr),
+      moduleInit(nullptr), moduleInitInsertPoint(nullptr), spirvOptions(opt),
+      builtinVars(), debugNone(nullptr), nullDebugExpr(nullptr),
+      stringLiterals() {}
 
 SpirvFunction *SpirvBuilder::createSpirvFunction(QualType returnType,
                                                  SourceLocation loc,
@@ -877,6 +879,14 @@ SpirvBuilder::createDemoteToHelperInvocation(SourceLocation loc) {
   return inst;
 }
 
+SpirvInstruction *
+SpirvBuilder::createIsHelperInvocationEXT(QualType type, SourceLocation loc) {
+  assert(insertPoint && "null insert point");
+  auto *inst = new (context) SpirvIsHelperInvocationEXT(type, loc);
+  insertPoint->addInstruction(inst);
+  return inst;
+}
+
 SpirvDebugSource *SpirvBuilder::createDebugSource(llvm::StringRef file,
                                                   llvm::StringRef text) {
   auto *inst = new (context) SpirvDebugSource(file, text);
@@ -1247,6 +1257,22 @@ SpirvVariable *SpirvBuilder::addStageIOVar(QualType type,
   return var;
 }
 
+SpirvVariable *SpirvBuilder::addVarForHelperInvocation(QualType type,
+                                                       bool isPrecise,
+                                                       SourceLocation loc) {
+  SpirvVariable *var = addModuleVar(type, spv::StorageClass::Private, isPrecise,
+                                    "HelperInvocation", llvm::None, loc);
+
+  auto *oldInsertPoint = insertPoint;
+  switchInsertPointToModuleInit();
+
+  SpirvInstruction *isHelperInvocation = createIsHelperInvocationEXT(type, loc);
+  createStore(var, isHelperInvocation, loc, SourceRange());
+
+  insertPoint = oldInsertPoint;
+  return var;
+}
+
 SpirvVariable *SpirvBuilder::addStageBuiltinVar(QualType type,
                                                 spv::StorageClass storageClass,
                                                 spv::BuiltIn builtin,
@@ -1601,13 +1627,14 @@ std::vector<uint32_t> SpirvBuilder::takeModule() {
   // Run necessary visitor passes first
   LiteralTypeVisitor literalTypeVisitor(astContext, context, spirvOptions);
   LowerTypeVisitor lowerTypeVisitor(astContext, context, spirvOptions);
-  CapabilityVisitor capabilityVisitor(astContext, context, spirvOptions, *this);
+  CapabilityVisitor capabilityVisitor(astContext, context, spirvOptions, *this,
+                                      featureManager);
   RelaxedPrecisionVisitor relaxedPrecisionVisitor(context, spirvOptions);
   PreciseVisitor preciseVisitor(context, spirvOptions);
   NonUniformVisitor nonUniformVisitor(context, spirvOptions);
-  RemoveBufferBlockVisitor removeBufferBlockVisitor(astContext, context,
-                                                    spirvOptions);
-  EmitVisitor emitVisitor(astContext, context, spirvOptions);
+  RemoveBufferBlockVisitor removeBufferBlockVisitor(
+      astContext, context, spirvOptions, featureManager);
+  EmitVisitor emitVisitor(astContext, context, spirvOptions, featureManager);
 
   mod->invokeVisitor(&literalTypeVisitor, true);
 

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -507,7 +507,7 @@ SpirvEmitter::SpirvEmitter(CompilerInstance &ci)
       spirvOptions(ci.getCodeGenOpts().SpirvOptions),
       entryFunctionName(ci.getCodeGenOpts().HLSLEntryFunction), spvContext(),
       featureManager(diags, spirvOptions),
-      spvBuilder(astContext, spvContext, spirvOptions),
+      spvBuilder(astContext, spvContext, spirvOptions, featureManager),
       declIdMapper(astContext, spvContext, spvBuilder, *this, featureManager,
                    spirvOptions),
       entryFunction(nullptr), curFunction(nullptr), curThis(nullptr),

tools/clang/lib/SPIRV/SpirvInstruction.cpp

@@ -83,6 +83,7 @@ DEFINE_INVOKE_VISITOR_FOR_CLASS(SpirvVectorShuffle)
 DEFINE_INVOKE_VISITOR_FOR_CLASS(SpirvArrayLength)
 DEFINE_INVOKE_VISITOR_FOR_CLASS(SpirvRayTracingOpNV)
 DEFINE_INVOKE_VISITOR_FOR_CLASS(SpirvDemoteToHelperInvocation)
+DEFINE_INVOKE_VISITOR_FOR_CLASS(SpirvIsHelperInvocationEXT)
 DEFINE_INVOKE_VISITOR_FOR_CLASS(SpirvDebugInfoNone)
 DEFINE_INVOKE_VISITOR_FOR_CLASS(SpirvDebugSource)
 DEFINE_INVOKE_VISITOR_FOR_CLASS(SpirvDebugCompilationUnit)
@@ -878,6 +879,11 @@ SpirvDemoteToHelperInvocation::SpirvDemoteToHelperInvocation(SourceLocation loc)
                        spv::Op::OpDemoteToHelperInvocation, /*QualType*/ {},
                        loc) {}
 
+SpirvIsHelperInvocationEXT::SpirvIsHelperInvocationEXT(QualType resultType,
+                                                       SourceLocation loc)
+    : SpirvInstruction(IK_IsHelperInvocationEXT,
+                       spv::Op::OpIsHelperInvocationEXT, resultType, loc) {}
+
 // Note: we are using a null result type in the constructor. All debug
 // instructions should later get OpTypeVoid as their result type.
 SpirvDebugInstruction::SpirvDebugInstruction(Kind kind, uint32_t opcode)

tools/clang/test/CodeGenSPIRV/spirv.builtin.helper-invocation.hlsl

@@ -1,18 +1,19 @@
 // RUN: %dxc -T ps_6_0 -E main
 
-// CHECK:      OpEntryPoint Fragment
-// CHECK-SAME: %gl_HelperInvocation
+// CHECK-NOT: OpDecorate {{%\w+}} BuiltIn HelperInvocation
 
-// CHECK:      OpDecorate %gl_HelperInvocation BuiltIn HelperInvocation
-
-// CHECK:      %gl_HelperInvocation = OpVariable %_ptr_Input_bool Input
+// CHECK: %HelperInvocation = OpVariable %_ptr_Private_bool Private
 
 float4 main([[vk::builtin("HelperInvocation")]] bool isHI : HI) : SV_Target {
-// CHECK:      [[val:%\d+]] = OpLoad %bool %gl_HelperInvocation
+// CHECK:      [[val:%\d+]] = OpLoad %bool %HelperInvocation
 // CHECK-NEXT: OpStore %param_var_isHI [[val]]
     float ret = 1.0;
 
     if (isHI) ret = 2.0;
 
     return ret;
 }
+// CHECK:      %module_init = OpFunction %void None
+// CHECK:      %module_init_bb = OpLabel
+// CHECK:      [[HelperInvocation:%\d+]] = OpIsHelperInvocationEXT %bool
+// CHECK-NEXT: OpStore %HelperInvocation [[HelperInvocation]]