[spirv] Cleanup and simplify handling of NonUniform.

We used to manually set 'NonUniform' on SpirvEvalInfos, and we had a
limited manual propagation of this attribute. We also had to have
book-keeping for whether we have seen a non-uniform resource index
intrinsic.

With this change, I have removed all the random attribute assignments in
different places, and only assign this attribute in
processIntrinsicNonUniformResourceIndex function. We then use natural
propagation of this attribute, and we emit the NonUniformEXT decoration
in EmitVisitor.

tools/clang/include/clang/SPIRV/EmitVisitor.h

@@ -72,6 +72,13 @@ public:
   // instructions into the annotationsBinary.
   uint32_t emitType(const SpirvType *);
 
+  // Emits OpDecorate (or OpMemberDecorate if memberIndex is non-zero)
+  // targetting the given type. Uses the given decoration kind and its
+  // parameters.
+  void emitDecoration(uint32_t typeResultId, spv::Decoration,
+                      llvm::ArrayRef<uint32_t> decorationParams,
+                      llvm::Optional<uint32_t> memberIndex = llvm::None);
+
   uint32_t getOrCreateConstant(SpirvConstant *);
 
   // Emits an OpConstant instruction and returns its result-id.
@@ -106,13 +113,6 @@ private:
   // result-id. If not, creates a new result-id for such type and returns it.
   uint32_t getResultIdForType(const SpirvType *, bool *alreadyExists);
 
-  // Emits OpDecorate (or OpMemberDecorate if memberIndex is non-zero)
-  // targetting the given type. Uses the given decoration kind and its
-  // parameters.
-  void emitDecoration(uint32_t typeResultId, spv::Decoration,
-                      llvm::ArrayRef<uint32_t> decorationParams,
-                      llvm::Optional<uint32_t> memberIndex = llvm::None);
-
   // Emits an OpName (if memberIndex is not provided) or OpMemberName (if
   // memberIndex is provided) for the given target result-id.
   void emitNameForType(llvm::StringRef name, uint32_t targetTypeId,

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

@@ -229,6 +229,12 @@ public:
       SpirvInstruction *valueToOp, SpirvInstruction *comparator,
       SourceLocation loc = {});
 
+  /// \brief Creates an OpSampledImage SPIR-V instruction with proper
+  /// decorations for the given parameters.
+  SpirvSampledImage *createSampledImage(QualType, SpirvInstruction *image,
+                                        SpirvInstruction *sampler,
+                                        SourceLocation loc = {});
+
   /// \brief Creates an OpImageTexelPointer SPIR-V instruction with the given
   /// parameters.
   SpirvImageTexelPointer *createImageTexelPointer(QualType resultType,
@@ -259,9 +265,8 @@ public:
   SpirvInstruction *
   createImageSample(QualType texelType, QualType imageType,
                     SpirvInstruction *image, SpirvInstruction *sampler,
-                    bool isNonUniform, SpirvInstruction *coordinate,
-                    SpirvInstruction *compareVal, SpirvInstruction *bias,
-                    SpirvInstruction *lod,
+                    SpirvInstruction *coordinate, SpirvInstruction *compareVal,
+                    SpirvInstruction *bias, SpirvInstruction *lod,
                     std::pair<SpirvInstruction *, SpirvInstruction *> grad,
                     SpirvInstruction *constOffset, SpirvInstruction *varOffset,
                     SpirvInstruction *constOffsets, SpirvInstruction *sample,
@@ -302,11 +307,11 @@ public:
   SpirvInstruction *
   createImageGather(QualType texelType, QualType imageType,
                     SpirvInstruction *image, SpirvInstruction *sampler,
-                    bool isNonUniform, SpirvInstruction *coordinate,
-                    SpirvInstruction *component, SpirvInstruction *compareVal,
-                    SpirvInstruction *constOffset, SpirvInstruction *varOffset,
-                    SpirvInstruction *constOffsets, SpirvInstruction *sample,
-                    SpirvInstruction *residencyCode, SourceLocation loc = {});
+                    SpirvInstruction *coordinate, SpirvInstruction *component,
+                    SpirvInstruction *compareVal, SpirvInstruction *constOffset,
+                    SpirvInstruction *varOffset, SpirvInstruction *constOffsets,
+                    SpirvInstruction *sample, SpirvInstruction *residencyCode,
+                    SourceLocation loc = {});
 
   /// \brief Creates an OpImageSparseTexelsResident SPIR-V instruction for the
   /// given Resident Code and returns the instruction pointer.

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

@@ -161,7 +161,7 @@ public:
   void setRelaxedPrecision() { isRelaxedPrecision_ = true; }
   bool isRelaxedPrecision() const { return isRelaxedPrecision_; }
 
-  void setNonUniform(bool nu = true) { isNonUniform_ = true; }
+  void setNonUniform(bool nu = true) { isNonUniform_ = nu; }
   bool isNonUniform() const { return isNonUniform_; }
 
 protected:

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -381,6 +381,9 @@ bool CapabilityVisitor::visitInstruction(SpirvInstruction *instr) {
 
   // Add NonUniform capabilities if necessary
   if (instr->isNonUniform()) {
+    spvBuilder.addExtension(Extension::EXT_descriptor_indexing, "NonUniformEXT",
+                            instr->getSourceLocation());
+    spvBuilder.requireCapability(spv::Capability::ShaderNonUniformEXT);
     spvBuilder.requireCapability(getNonUniformCapability(resultType));
   }
 

tools/clang/lib/SPIRV/EmitVisitor.cpp

@@ -91,6 +91,12 @@ void EmitVisitor::initInstruction(SpirvInstruction *inst) {
     inst->setResultTypeId(resultTypeId);
   }
 
+  // Emit NonUniformEXT decoration (if any).
+  if (inst->isNonUniform()) {
+    typeHandler.emitDecoration(getOrAssignResultId<SpirvInstruction>(inst),
+                               spv::Decoration::NonUniformEXT, {});
+  }
+
   // Initialize the current instruction for emitting.
   curInst.clear();
   curInst.push_back(static_cast<uint32_t>(inst->getopcode()));

tools/clang/lib/SPIRV/SPIRVEmitter.cpp

@@ -549,8 +549,7 @@ SPIRVEmitter::SPIRVEmitter(CompilerInstance &ci)
       declIdMapper(shaderModel, astContext, spvContext, spvBuilder, *this,
                    featureManager, spirvOptions),
       entryFunction(nullptr), curFunction(nullptr), curThis(nullptr),
-      seenPushConstantAt(), isSpecConstantMode(false),
-      foundNonUniformResourceIndex(false), needsLegalization(false),
+      seenPushConstantAt(), isSpecConstantMode(false), needsLegalization(false),
       mainSourceFile(nullptr) {
   if (shaderModel.GetKind() == hlsl::ShaderModel::Kind::Invalid)
     emitError("unknown shader module: %0", {}) << shaderModel.GetName();
@@ -903,12 +902,6 @@ SpirvInstruction *SPIRVEmitter::loadIfGLValue(const Expr *expr,
   }
   assert(loadedInstr);
 
-  // Decorate with NonUniformEXT if loading from a pointer with that property.
-  // We are likely loading an element from the resource array here.
-  if (info->isNonUniform()) {
-    spvBuilder.decorateNonUniformEXT(loadedInstr);
-  }
-
   // Special-case: According to the SPIR-V Spec: There is no physical size or
   // bit pattern defined for boolean type. Therefore an unsigned integer is used
   // to represent booleans when layout is required. In such cases, after loading
@@ -1946,19 +1939,11 @@ void SPIRVEmitter::doSwitchStmt(const SwitchStmt *switchStmt,
 
 SpirvInstruction *
 SPIRVEmitter::doArraySubscriptExpr(const ArraySubscriptExpr *expr) {
-  // Make sure we don't have previously unhandled NonUniformResourceIndex()
-  assert(!foundNonUniformResourceIndex);
-
   llvm::SmallVector<SpirvInstruction *, 4> indices;
   const auto *base = collectArrayStructIndices(
       expr, /*rawIndex*/ false, /*rawIndices*/ nullptr, &indices);
   auto *info = loadIfAliasVarRef(base);
 
-  if (foundNonUniformResourceIndex) {
-    info->setNonUniform(); // Carry forward the NonUniformEXT decoration
-    foundNonUniformResourceIndex = false;
-  }
-
   if (!indices.empty()) {
     info = turnIntoElementPtr(base->getType(), info, expr->getType(), indices);
   }
@@ -3004,14 +2989,8 @@ SPIRVEmitter::processTextureLevelOfDetail(const CXXMemberCallExpr *expr,
   auto *samplerState = doExpr(expr->getArg(0));
   auto *coordinate = doExpr(expr->getArg(1));
 
-  auto *sampledImage = spvBuilder.createBinaryOp(
-      spv::Op::OpSampledImage, object->getType(), objectInfo, samplerState);
-
-  if (objectInfo->isNonUniform() || samplerState->isNonUniform()) {
-    // The sampled image will be used to access resource's memory, so we need
-    // to decorate it with NonUniformEXT.
-    spvBuilder.decorateNonUniformEXT(sampledImage);
-  }
+  auto *sampledImage = spvBuilder.createSampledImage(
+      object->getType(), objectInfo, samplerState, expr->getExprLoc());
 
   // The result type of OpImageQueryLod must be a float2.
   const QualType queryResultType =
@@ -3100,7 +3079,6 @@ SpirvInstruction *SPIRVEmitter::processTextureGatherRGBACmpRGBA(
   }
 
   auto *status = hasStatusArg ? doExpr(expr->getArg(numArgs - 1)) : nullptr;
-  const bool isNonUniform = image->isNonUniform() || sampler->isNonUniform();
 
   if (needsEmulation) {
     const auto elemType = hlsl::GetHLSLVecElementType(callee->getReturnType());
@@ -3109,7 +3087,7 @@ SpirvInstruction *SPIRVEmitter::processTextureGatherRGBACmpRGBA(
     for (uint32_t i = 0; i < 4; ++i) {
       varOffset = doExpr(expr->getArg(2 + isCmp + i));
       auto *gatherRet = spvBuilder.createImageGather(
-          retType, imageType, image, sampler, isNonUniform, coordinate,
+          retType, imageType, image, sampler, coordinate,
           spvBuilder.getConstantInt(astContext.IntTy,
                                     llvm::APInt(32, component, true)),
           compareVal,
@@ -3122,7 +3100,7 @@ SpirvInstruction *SPIRVEmitter::processTextureGatherRGBACmpRGBA(
   }
 
   return spvBuilder.createImageGather(
-      retType, imageType, image, sampler, isNonUniform, coordinate,
+      retType, imageType, image, sampler, coordinate,
       spvBuilder.getConstantInt(astContext.IntTy,
                                 llvm::APInt(32, component, true)),
       compareVal, constOffset, varOffset, constOffsets,
@@ -3173,8 +3151,7 @@ SPIRVEmitter::processTextureGatherCmp(const CXXMemberCallExpr *expr) {
       hasStatusArg ? doExpr(expr->getArg(numArgs - 1)) : nullptr;
 
   return spvBuilder.createImageGather(
-      retType, imageType, image, sampler,
-      image->isNonUniform() || sampler->isNonUniform(), coordinate,
+      retType, imageType, image, sampler, coordinate,
       /*component*/ nullptr, comparator, constOffset, varOffset,
       /*constOffsets*/ nullptr,
       /*sampleNumber*/ nullptr, status);
@@ -3197,11 +3174,6 @@ SpirvInstruction *SPIRVEmitter::processBufferTextureLoad(
 
   auto *objectInfo = loadIfGLValue(object);
 
-  if (objectInfo->isNonUniform()) {
-    // Decoreate the image handle for OpImageFetch/OpImageRead
-    spvBuilder.decorateNonUniformEXT(objectInfo);
-  }
-
   // For Texture2DMS and Texture2DMSArray, Sample must be used rather than Lod.
   SpirvInstruction *sampleNumber = nullptr;
   if (TypeTranslator::isTextureMS(type) ||
@@ -3966,7 +3938,7 @@ SPIRVEmitter::processIntrinsicMemberCall(const CXXMemberCallExpr *expr,
 
 SpirvInstruction *SPIRVEmitter::createImageSample(
     QualType retType, QualType imageType, SpirvInstruction *image,
-    SpirvInstruction *sampler, bool isNonUniform, SpirvInstruction *coordinate,
+    SpirvInstruction *sampler, SpirvInstruction *coordinate,
     SpirvInstruction *compareVal, SpirvInstruction *bias, SpirvInstruction *lod,
     std::pair<SpirvInstruction *, SpirvInstruction *> grad,
     SpirvInstruction *constOffset, SpirvInstruction *varOffset,
@@ -3976,10 +3948,10 @@ SpirvInstruction *SPIRVEmitter::createImageSample(
   // SampleDref* instructions in SPIR-V always return a scalar.
   // They also have the correct type in HLSL.
   if (compareVal) {
-    return spvBuilder.createImageSample(
-        retType, imageType, image, sampler, isNonUniform, coordinate,
-        compareVal, bias, lod, grad, constOffset, varOffset, constOffsets,
-        sample, minLod, residencyCodeId);
+    return spvBuilder.createImageSample(retType, imageType, image, sampler,
+                                        coordinate, compareVal, bias, lod, grad,
+                                        constOffset, varOffset, constOffsets,
+                                        sample, minLod, residencyCodeId);
   }
 
   // Non-Dref Sample instructions in SPIR-V must always return a vec4.
@@ -4003,9 +3975,9 @@ SpirvInstruction *SPIRVEmitter::createImageSample(
     needsLegalization = true;
 
   auto *retVal = spvBuilder.createImageSample(
-      texelType, imageType, image, sampler, isNonUniform, coordinate,
-      compareVal, bias, lod, grad, constOffset, varOffset, constOffsets, sample,
-      minLod, residencyCodeId);
+      texelType, imageType, image, sampler, coordinate, compareVal, bias, lod,
+      grad, constOffset, varOffset, constOffsets, sample, minLod,
+      residencyCodeId);
 
   // Extract smaller vector from the vec4 result if necessary.
   if (texelType != retType) {
@@ -4071,19 +4043,18 @@ SPIRVEmitter::processTextureSampleGather(const CXXMemberCallExpr *expr,
   SpirvInstruction *constOffset = nullptr, *varOffset = nullptr;
   if (hasOffsetArg)
     handleOffsetInMethodCall(expr, 2, &constOffset, &varOffset);
-  const bool isNonUniform = image->isNonUniform() || sampler->isNonUniform();
 
   const auto retType = expr->getDirectCallee()->getReturnType();
   if (isSample) {
-    return createImageSample(
-        retType, imageType, image, sampler, isNonUniform, coordinate,
-        /*compareVal*/ nullptr, /*bias*/ nullptr, /*lod*/ nullptr,
-        std::make_pair(nullptr, nullptr), constOffset, varOffset,
-        /*constOffsets*/ nullptr, /*sampleNumber*/ nullptr,
-        /*minLod*/ clamp, status);
+    return createImageSample(retType, imageType, image, sampler, coordinate,
+                             /*compareVal*/ nullptr, /*bias*/ nullptr,
+                             /*lod*/ nullptr, std::make_pair(nullptr, nullptr),
+                             constOffset, varOffset,
+                             /*constOffsets*/ nullptr, /*sampleNumber*/ nullptr,
+                             /*minLod*/ clamp, status);
   } else {
     return spvBuilder.createImageGather(
-        retType, imageType, image, sampler, isNonUniform, coordinate,
+        retType, imageType, image, sampler, coordinate,
         // .Gather() doc says we return four components of red data.
         spvBuilder.getConstantInt(astContext.IntTy, llvm::APInt(32, 0)),
         /*compareVal*/ nullptr, constOffset, varOffset,
@@ -4161,13 +4132,12 @@ SPIRVEmitter::processTextureSampleBiasLevel(const CXXMemberCallExpr *expr,
 
   const auto retType = expr->getDirectCallee()->getReturnType();
 
-  return createImageSample(
-      retType, imageType, image, sampler,
-      image->isNonUniform() || sampler->isNonUniform(), coordinate,
-      /*compareVal*/ nullptr, bias, lod, std::make_pair(nullptr, nullptr),
-      constOffset, varOffset,
-      /*constOffsets*/ nullptr, /*sampleNumber*/ nullptr, /*minLod*/ clamp,
-      status);
+  return createImageSample(retType, imageType, image, sampler, coordinate,
+                           /*compareVal*/ nullptr, bias, lod,
+                           std::make_pair(nullptr, nullptr), constOffset,
+                           varOffset,
+                           /*constOffsets*/ nullptr, /*sampleNumber*/ nullptr,
+                           /*minLod*/ clamp, status);
 }
 
 SpirvInstruction *
@@ -4219,13 +4189,12 @@ SPIRVEmitter::processTextureSampleGrad(const CXXMemberCallExpr *expr) {
     handleOffsetInMethodCall(expr, 4, &constOffset, &varOffset);
 
   const auto retType = expr->getDirectCallee()->getReturnType();
-  return createImageSample(
-      retType, imageType, image, sampler,
-      image->isNonUniform() || sampler->isNonUniform(), coordinate,
-      /*compareVal*/ nullptr, /*bias*/ nullptr, /*lod*/ nullptr,
-      std::make_pair(ddx, ddy), constOffset, varOffset,
-      /*constOffsets*/ nullptr, /*sampleNumber*/ nullptr,
-      /*minLod*/ clamp, status);
+  return createImageSample(retType, imageType, image, sampler, coordinate,
+                           /*compareVal*/ nullptr, /*bias*/ nullptr,
+                           /*lod*/ nullptr, std::make_pair(ddx, ddy),
+                           constOffset, varOffset,
+                           /*constOffsets*/ nullptr, /*sampleNumber*/ nullptr,
+                           /*minLod*/ clamp, status);
 }
 
 SpirvInstruction *
@@ -4311,8 +4280,7 @@ SPIRVEmitter::processTextureSampleCmpCmpLevelZero(const CXXMemberCallExpr *expr,
   const auto imageType = imageExpr->getType();
 
   return createImageSample(
-      retType, imageType, image, sampler,
-      image->isNonUniform() || sampler->isNonUniform(), coordinate, compareVal,
+      retType, imageType, image, sampler, coordinate, compareVal,
       /*bias*/ nullptr, lod, std::make_pair(nullptr, nullptr), constOffset,
       varOffset,
       /*constOffsets*/ nullptr, /*sampleNumber*/ nullptr, /*minLod*/ clamp,
@@ -5260,9 +5228,6 @@ SpirvInstruction *SPIRVEmitter::processBinaryOp(
     // Propagate RelaxedPrecision
     if (lhsVal->isRelaxedPrecision() || rhsVal->isRelaxedPrecision())
       val->setRelaxedPrecision();
-    // Propagate NonUniformEXT
-    if (lhsVal->isNonUniform() || rhsVal->isNonUniform())
-      val->setNonUniform();
 
     return val;
   }
@@ -5797,10 +5762,6 @@ SPIRVEmitter::tryToAssignToRWBufferRWTexture(const Expr *lhs,
     auto *baseInfo = doExpr(baseExpr);
     auto *image = spvBuilder.createLoad(imageType, baseInfo);
     spvBuilder.createImageWrite(imageType, image, loc, rhs);
-    if (baseInfo->isNonUniform()) {
-      // Decorate the image handle for OpImageWrite
-      spvBuilder.decorateNonUniformEXT(image);
-    }
     return rhs;
   }
   return nullptr;
@@ -6841,11 +6802,6 @@ SPIRVEmitter::processIntrinsicInterlockedMethod(const CallExpr *expr,
       auto *coordInstr = doExpr(index);
       ptr = spvBuilder.createImageTexelPointer(baseType, baseInstr, coordInstr,
                                                zero);
-      if (baseInstr->isNonUniform()) {
-        // Image texel pointer will used to access image memory. Vulkan requires
-        // it to be decorated with NonUniformEXT.
-        spvBuilder.decorateNonUniformEXT(ptr);
-      }
     }
   }
   if (!ptr) {
@@ -6902,10 +6858,7 @@ SPIRVEmitter::processIntrinsicInterlockedMethod(const CallExpr *expr,
 
 SpirvInstruction *
 SPIRVEmitter::processIntrinsicNonUniformResourceIndex(const CallExpr *expr) {
-  foundNonUniformResourceIndex = true;
   auto *index = doExpr(expr->getArg(0));
-  index->setNonUniform();
-
   // Decorate the expression in NonUniformResourceIndex() with NonUniformEXT.
   // Aside from this, we also need to eventually populate the NonUniformEXT
   // status to the usage of this expression: the "pointer" operand to a memory
@@ -6915,8 +6868,7 @@ SPIRVEmitter::processIntrinsicNonUniformResourceIndex(const CallExpr *expr) {
   // image instructions) and the resource descriptor being accessed is not
   // dynamically uniform, then the operand corresponding to that resource (e.g.
   // the pointer or sampled image operand) must be decorated with NonUniformEXT.
-  spvBuilder.decorateNonUniformEXT(index);
-
+  index->setNonUniform();
   return index;
 }
 
@@ -8901,7 +8853,7 @@ SpirvConstant *SPIRVEmitter::translateAPValue(const APValue &value,
 
 SpirvConstant *SPIRVEmitter::translateAPInt(const llvm::APInt &intValue,
                                             QualType targetType) {
-  targetType = typeTranslator.getIntendedLiteralType(targetType);	
+  targetType = typeTranslator.getIntendedLiteralType(targetType);
   return spvBuilder.getConstantInt(targetType, intValue, isSpecConstantMode);
 }
 
@@ -8974,7 +8926,7 @@ SPIRVEmitter::tryToEvaluateAsFloat32(const llvm::APFloat &floatValue) {
 
 SpirvConstant *SPIRVEmitter::translateAPFloat(llvm::APFloat floatValue,
                                               QualType targetType) {
-  targetType = typeTranslator.getIntendedLiteralType(targetType);	
+  targetType = typeTranslator.getIntendedLiteralType(targetType);
   return spvBuilder.getConstantFloat(targetType, floatValue,
                                      isSpecConstantMode);
 }

tools/clang/lib/SPIRV/SPIRVEmitter.h

@@ -902,15 +902,16 @@ private:
   /// HLSL image sampling methods may return a scalar, vec1, vec2, vec3, or
   /// vec4. But non-Dref image sampling instructions in SPIR-V must always
   /// return a vec4. As a result, an extra processing step is necessary.
-  SpirvInstruction *createImageSample(
-      QualType retType, QualType imageType, SpirvInstruction *image,
-      SpirvInstruction *sampler, bool isNonUniform,
-      SpirvInstruction *coordinate, SpirvInstruction *compareVal,
-      SpirvInstruction *bias, SpirvInstruction *lod,
-      std::pair<SpirvInstruction *, SpirvInstruction *> grad,
-      SpirvInstruction *constOffset, SpirvInstruction *varOffset,
-      SpirvInstruction *constOffsets, SpirvInstruction *sample,
-      SpirvInstruction *minLod, SpirvInstruction *residencyCodeId);
+  SpirvInstruction *
+  createImageSample(QualType retType, QualType imageType,
+                    SpirvInstruction *image, SpirvInstruction *sampler,
+                    SpirvInstruction *coordinate, SpirvInstruction *compareVal,
+                    SpirvInstruction *bias, SpirvInstruction *lod,
+                    std::pair<SpirvInstruction *, SpirvInstruction *> grad,
+                    SpirvInstruction *constOffset, SpirvInstruction *varOffset,
+                    SpirvInstruction *constOffsets, SpirvInstruction *sample,
+                    SpirvInstruction *minLod,
+                    SpirvInstruction *residencyCodeId);
 
   /// \brief Emit an OpLine instruction for the given source location.
   void emitDebugLine(SourceLocation);
@@ -995,13 +996,6 @@ private:
   /// all 32-bit scalar constants will be translated into OpSpecConstant.
   bool isSpecConstantMode;
 
-  /// Indicates that we have found a NonUniformResourceIndex call when
-  /// traversing.
-  /// This field is used to convery information in a bottom-up manner; if we
-  /// have something like `aResource[NonUniformResourceIndex(aIndex)]`, we need
-  /// to attach `aResource` with proper decorations.
-  bool foundNonUniformResourceIndex;
-
   /// Whether the translated SPIR-V binary needs legalization.
   ///
   /// The following cases will require legalization:

tools/clang/lib/SPIRV/SpirvBuilder.cpp

@@ -186,6 +186,7 @@ SpirvLoad *SpirvBuilder::createLoad(QualType resultType,
   instruction->setStorageClass(pointer->getStorageClass());
   instruction->setRValue();
   instruction->setLayoutRule(pointer->getLayoutRule());
+  instruction->setNonUniform(pointer->isNonUniform());
   insertPoint->addInstruction(instruction);
   return instruction;
 }
@@ -200,6 +201,7 @@ SpirvLoad *SpirvBuilder::createLoad(const SpirvType *resultType,
   instruction->setStorageClass(pointer->getStorageClass());
   instruction->setRValue();
   instruction->setLayoutRule(pointer->getLayoutRule());
+  instruction->setNonUniform(pointer->isNonUniform());
   insertPoint->addInstruction(instruction);
   return instruction;
 }
@@ -232,6 +234,10 @@ SpirvBuilder::createAccessChain(QualType resultType, SpirvInstruction *base,
       new (context) SpirvAccessChain(resultType, /*id*/ 0, loc, base, indexes);
   instruction->setStorageClass(base->getStorageClass());
   instruction->setLayoutRule(base->getLayoutRule());
+  bool isNonUniform = base->isNonUniform();
+  for (auto *index : indexes)
+    isNonUniform = isNonUniform || index->isNonUniform();
+  instruction->setNonUniform(isNonUniform);
   insertPoint->addInstruction(instruction);
   return instruction;
 }
@@ -245,6 +251,10 @@ SpirvAccessChain *SpirvBuilder::createAccessChain(
   instruction->setResultType(resultType);
   instruction->setStorageClass(base->getStorageClass());
   instruction->setLayoutRule(base->getLayoutRule());
+  bool isNonUniform = base->isNonUniform();
+  for (auto *index : indexes)
+    isNonUniform = isNonUniform || index->isNonUniform();
+  instruction->setNonUniform(isNonUniform);
   insertPoint->addInstruction(instruction);
   return instruction;
 }
@@ -255,6 +265,7 @@ SpirvUnaryOp *SpirvBuilder::createUnaryOp(spv::Op op, QualType resultType,
   assert(insertPoint && "null insert point");
   auto *instruction =
       new (context) SpirvUnaryOp(op, resultType, /*id*/ 0, loc, operand);
+  instruction->setNonUniform(operand->isNonUniform());
   insertPoint->addInstruction(instruction);
   return instruction;
 }
@@ -266,6 +277,7 @@ SpirvBinaryOp *SpirvBuilder::createBinaryOp(spv::Op op, QualType resultType,
   assert(insertPoint && "null insert point");
   auto *instruction =
       new (context) SpirvBinaryOp(op, resultType, /*id*/ 0, loc, lhs, rhs);
+  instruction->setNonUniform(lhs->isNonUniform() || rhs->isNonUniform());
   insertPoint->addInstruction(instruction);
   return instruction;
 }
@@ -279,6 +291,7 @@ SpirvBinaryOp *SpirvBuilder::createBinaryOp(spv::Op op,
   auto *instruction =
       new (context) SpirvBinaryOp(op, /*QualType*/ {}, /*id*/ 0, loc, lhs, rhs);
   instruction->setResultType(resultType);
+  instruction->setNonUniform(lhs->isNonUniform() || rhs->isNonUniform());
   insertPoint->addInstruction(instruction);
   return instruction;
 }
@@ -351,12 +364,25 @@ SpirvAtomic *SpirvBuilder::createAtomicCompareExchange(
   return instruction;
 }
 
+SpirvSampledImage *SpirvBuilder::createSampledImage(QualType imageType,
+                                                    SpirvInstruction *image,
+                                                    SpirvInstruction *sampler,
+                                                    SourceLocation loc) {
+  assert(insertPoint && "null insert point");
+  auto *sampledImage =
+      new (context) SpirvSampledImage(imageType, /*id*/ 0, loc, image, sampler);
+  sampledImage->setNonUniform(image->isNonUniform() || sampler->isNonUniform());
+  insertPoint->addInstruction(sampledImage);
+  return sampledImage;
+}
+
 SpirvImageTexelPointer *SpirvBuilder::createImageTexelPointer(
     QualType resultType, SpirvInstruction *image, SpirvInstruction *coordinate,
     SpirvInstruction *sample, SourceLocation loc) {
   assert(insertPoint && "null insert point");
   auto *instruction = new (context) SpirvImageTexelPointer(
       resultType, /*id*/ 0, loc, image, coordinate, sample);
+  instruction->setNonUniform(image->isNonUniform());
   insertPoint->addInstruction(instruction);
   return instruction;
 }
@@ -401,7 +427,7 @@ spv::ImageOperandsMask SpirvBuilder::composeImageOperandsMask(
 
 SpirvInstruction *SpirvBuilder::createImageSample(
     QualType texelType, QualType imageType, SpirvInstruction *image,
-    SpirvInstruction *sampler, bool isNonUniform, SpirvInstruction *coordinate,
+    SpirvInstruction *sampler, SpirvInstruction *coordinate,
     SpirvInstruction *compareVal, SpirvInstruction *bias, SpirvInstruction *lod,
     std::pair<SpirvInstruction *, SpirvInstruction *> grad,
     SpirvInstruction *constOffset, SpirvInstruction *varOffset,
@@ -434,15 +460,7 @@ SpirvInstruction *SpirvBuilder::createImageSample(
   assert(lod == nullptr || minLod == nullptr);
 
   // An OpSampledImage is required to do the image sampling.
-  auto *sampledImage =
-      new (context) SpirvSampledImage(imageType, /*id*/ 0, loc, image, sampler);
-  insertPoint->addInstruction(sampledImage);
-
-  if (isNonUniform) {
-    // The sampled image will be used to access resource's memory, so we need
-    // to decorate it with NonUniformEXT.
-    decorateNonUniformEXT(sampledImage, loc);
-  }
+  auto *sampledImage = createSampledImage(imageType, image, sampler, loc);
 
   const auto mask = composeImageOperandsMask(
       bias, lod, grad, constOffset, varOffset, constOffsets, sample, minLod);
@@ -520,7 +538,7 @@ void SpirvBuilder::createImageWrite(QualType imageType, SpirvInstruction *image,
 
 SpirvInstruction *SpirvBuilder::createImageGather(
     QualType texelType, QualType imageType, SpirvInstruction *image,
-    SpirvInstruction *sampler, bool isNonUniform, SpirvInstruction *coordinate,
+    SpirvInstruction *sampler, SpirvInstruction *coordinate,
     SpirvInstruction *component, SpirvInstruction *compareVal,
     SpirvInstruction *constOffset, SpirvInstruction *varOffset,
     SpirvInstruction *constOffsets, SpirvInstruction *sample,
@@ -528,15 +546,7 @@ SpirvInstruction *SpirvBuilder::createImageGather(
   assert(insertPoint && "null insert point");
 
   // An OpSampledImage is required to do the image sampling.
-  auto *sampledImage =
-      new (context) SpirvSampledImage(imageType, /*id*/ 0, loc, image, sampler);
-  insertPoint->addInstruction(sampledImage);
-
-  if (isNonUniform) {
-    // The sampled image will be used to access resource's memory, so we need
-    // to decorate it with NonUniformEXT.
-    decorateNonUniformEXT(sampledImage, loc);
-  }
+  auto *sampledImage = createSampledImage(imageType, image, sampler, loc);
 
   // TODO: Update ImageGather to accept minLod if necessary.
   const auto mask = composeImageOperandsMask(

tools/clang/test/CodeGenSPIRV/intrinsics.non-uniform-resource-index.hlsl

@@ -1,5 +1,6 @@
 // Run: %dxc -T ps_6_0 -E main
 
+// CHECK: OpCapability RuntimeDescriptorArrayEXT
 // CHECK: OpCapability ShaderNonUniformEXT
 // CHECK: OpCapability SampledImageArrayNonUniformIndexingEXT
 // CHECK: OpCapability StorageImageArrayNonUniformIndexingEXT
@@ -14,44 +15,53 @@
 // CHECK: OpDecorate [[nu3:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu4:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu5:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu6:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu7:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu8:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu9:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu10:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu11:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu12:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu13:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu14:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu15:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu16:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu17:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu18:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu19:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu20:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu21:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu22:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu23:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu24:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu25:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu26:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu27:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu28:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu29:%\d+]] NonUniformEXT
-
 // CHECK: OpDecorate [[nu30:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu31:%\d+]] NonUniformEXT
 // CHECK: OpDecorate [[nu32:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu33:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu34:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu35:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu36:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu37:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu38:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu39:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu40:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu41:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu42:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu43:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu44:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu45:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu46:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu47:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu48:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu49:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu50:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu51:%\d+]] NonUniformEXT
+// CHECK: OpDecorate [[nu52:%\d+]] NonUniformEXT
 
 Texture2D           gTextures[32];
 SamplerState        gSamplers[];
@@ -62,79 +72,99 @@ SubpassInput        gSubpassInputs[32];
 
 float4 main(uint index : A, float2 loc : B, int2 offset : C) : SV_Target {
 // CHECK: [[nu1]] = OpLoad %uint %index
-// CHECK: [[nu2]] = OpLoad %type_2d_image
-// CHECK: [[nu3]] = OpIAdd %uint {{%\d+}} %uint_1
-// CHECK: [[nu4]] = OpLoad %type_sampler
-// CHECK: [[nu5]] = OpSampledImage %type_sampled_image
+// CHECK: [[nu2]] = OpAccessChain %_ptr_UniformConstant_type_2d_image %gTextures [[nu1]]
+// CHECK: [[nu3]] = OpLoad %type_2d_image
+// CHECK: [[nu4]] = OpIAdd %uint {{%\d+}} %uint_1
+// CHECK: [[nu5]] = OpAccessChain %_ptr_UniformConstant_type_sampler %gSamplers [[nu4]]
+// CHECK: [[nu6]] = OpLoad %type_sampler
+// CHECK: [[nu7]] = OpSampledImage %type_sampled_image
 // CHECK:           OpImageSampleImplicitLod
     float4 v1 = gTextures[NonUniformResourceIndex(index)].Sample(
         gSamplers[NonUniformResourceIndex(index + 1)], loc);
 
-// CHECK: [[nu6]] = OpLoad %uint %index
-// CHECK:           OpIAdd %uint {{%\d+}} %uint_1
-// CHECK: [[nu7]] = OpLoad %type_sampler
-// CHECK: [[nu8]] = OpSampledImage %type_sampled_image
-// CHECK:           OpImageSampleImplicitLod
+// CHECK:  [[nu8]] = OpLoad %uint %index
+// CHECK:            OpIAdd %uint {{%\d+}} %uint_1
+// CHECK:  [[nu9]] = OpAccessChain %_ptr_UniformConstant_type_sampler %gSamplers [[nu8]]
+// CHECK: [[nu10]] = OpLoad %type_sampler
+// CHECK: [[nu11]] = OpSampledImage %type_sampled_image
+// CHECK:            OpImageSampleImplicitLod
     float4 v2 = gTextures[0].Sample(
         gSamplers[NonUniformResourceIndex(index++)], loc);
 
 // CHECK:            OpLoad %uint %index
 // CHECK:            OpISub %uint {{%\d+}} %uint_1
-// CHECK:  [[nu9]] = OpLoad %uint %index
-// CHECK: [[nu10]] = OpLoad %type_2d_image
-// CHECK: [[nu11]] = OpSampledImage %type_sampled_image
+// CHECK: [[nu12]] = OpLoad %uint %index
+// CHECK: [[nu13]] = OpAccessChain %_ptr_UniformConstant_type_2d_image %gTextures [[nu12]]
+// CHECK: [[nu14]] = OpLoad %type_2d_image
+// CHECK: [[nu15]] = OpSampledImage %type_sampled_image
 // CHECK:            OpImageSampleImplicitLod
     float4 v3 = gTextures[NonUniformResourceIndex(--index)].Sample(
         gSamplers[0], loc);
 
-// CHECK: [[nu12]] = OpIMul %uint
-// CHECK: [[nu13]] = OpLoad %type_2d_image_0
+// CHECK: [[nu16]] = OpIMul %uint
+// CHECK: [[nu17]] = OpAccessChain %_ptr_UniformConstant_type_2d_image_0 %gRWTextures [[nu16]]
+// CHECK: [[nu18]] = OpLoad %type_2d_image_0
 // CHECK:            OpImageRead
     float4 v4 = gRWTextures[NonUniformResourceIndex(index * index)].Load(loc);
 
-// CHECK: [[nu14]] = OpLoad %uint %index
-// CHECK:            OpUMod %uint {{%\d+}} %uint_3
-// CHECK: [[nu15]] = OpLoad %type_2d_image_0
+// CHECK: [[nu19]] = OpLoad %uint %index
+// CHECK: [[nu20]] = OpUMod %uint [[nu19]] %uint_3
+// CHECK: [[nu21]] = OpAccessChain %_ptr_UniformConstant_type_2d_image_0 %gRWTextures [[nu20]]
+// CHECK: [[nu22]] = OpLoad %type_2d_image_0
 // CHECK:            OpImageWrite
     gRWTextures[NonUniformResourceIndex(index) % 3][loc] = 4;
 
-// CHECK: [[nu16]] = OpLoad %uint %index
+// CHECK: [[nu23]] = OpLoad %uint %index
 // CHECK:            OpLoad %uint %index
-// CHECK: [[nu17]] = OpLoad %type_buffer_image
+// CHECK: [[nu24]] = OpIMul %uint [[nu23]] {{%\d+}}
+// CHECK: [[nu25]] = OpAccessChain %_ptr_UniformConstant_type_buffer_image %gBuffers [[nu24]]
+// CHECK: [[nu26]] = OpLoad %type_buffer_image
 // CHECK:            OpImageFetch
     float4 v5 = gBuffers[NonUniformResourceIndex(index) * index][5];
 
-// CHECK: [[nu18]] = OpLoad %uint %index
-// CHECK: [[nu19]] = OpLoad %type_buffer_image_0
+// CHECK: [[nu27]] = OpLoad %uint %index
+// CHECK: [[nu28]] = OpAccessChain %_ptr_UniformConstant_type_buffer_image_0 %gRWBuffers [[nu27]]
+// CHECK: [[nu29]] = OpLoad %type_buffer_image_0
 // CHECK:            OpImageRead
     float4 v6 = gRWBuffers[NonUniformResourceIndex(index)].Load(6);
 
-// CHECK: [[nu20]] = OpLoad %uint %index
-// CHECK: [[nu21]] = OpLoad %type_buffer_image_0
+// CHECK: [[nu30]] = OpLoad %uint %index
+// CHECK: [[nu31]] = OpAccessChain %_ptr_UniformConstant_type_buffer_image_0 %gRWBuffers [[nu30]]
+// CHECK: [[nu32]] = OpLoad %type_buffer_image_0
 // CHECK:            OpImageWrite
     gRWBuffers[NonUniformResourceIndex(index)][8] = 9;
 
-// CHECK: [[nu22]] = OpLoad %uint %index
-// CHECK: [[nu23]] = OpLoad %uint %index
-// CHECK: [[nu24]] = OpImageTexelPointer %_ptr_Image_uint {{%\d+}} %uint_10 %uint_0
+// CHECK: [[nu33]] = OpLoad %uint %index
+// CHECK: [[nu34]] = OpLoad %uint %index
+// CHECK: [[nu35]] = OpIMul %uint [[nu33]] [[nu34]]
+// CHECK: [[nu36]] = OpAccessChain %_ptr_UniformConstant_type_buffer_image_0 %gRWBuffers [[nu35]]
+// CHECK: [[nu37]] = OpImageTexelPointer %_ptr_Image_uint {{%\d+}} %uint_10 %uint_0
 // CHECK:            OpAtomicIAdd
     uint old = 0;
     InterlockedAdd(gRWBuffers[NonUniformResourceIndex(index) * NonUniformResourceIndex(index)][10], 1, old);
 
-// CHECK: [[nu25]] = OpLoad %uint %index
-// CHECK: [[nu26]] = OpLoad %type_subpass_image
+// CHECK: [[nu38]] = OpLoad %uint %index
+// CHECK: [[nu39]] = OpAccessChain %_ptr_UniformConstant_type_subpass_image %gSubpassInputs [[nu38]]
+// CHECK: [[nu40]] = OpLoad %type_subpass_image
 // CHECK:            OpImageRead
     float4 v7 = gSubpassInputs[NonUniformResourceIndex(index)].SubpassLoad();
 
-// CHECK: [[nu27]] = OpLoad %uint %index
-// CHECK: [[nu28]] = OpLoad %type_2d_image
-// CHECK: [[nu29]] = OpSampledImage %type_sampled_image
+// CHECK: [[nu41]] = OpLoad %uint %index
+// CHECK: [[nu42]] = OpAccessChain %_ptr_UniformConstant_type_2d_image %gTextures [[nu41]]
+// CHECK: [[nu43]] = OpLoad %type_2d_image
+// CHECK: [[nu44]] = OpSampledImage %type_sampled_image
 // CHECK:            OpImageGather
     float4 v8 = gTextures[NonUniformResourceIndex(index)].Gather(gSamplers[0], loc, offset);
 
-// CHECK: [[nu30]] = OpLoad %uint %index
-// CHECK: [[nu31]] = OpLoad %type_2d_image
-// CHECK: [[nu32]] = OpSampledImage %type_sampled_image
+// CHECK: [[nu45]] = OpLoad %uint %index
+// CHECK: [[nu46]] = OpAccessChain %_ptr_UniformConstant_type_2d_image %gTextures [[nu45]]
+// CHECK: [[nu47]] = OpLoad %type_2d_image
+// CHECK: [[nu48]] = OpSampledImage %type_sampled_image
+// CHECK:            OpImageQueryLod
+// CHECK: [[nu49]] = OpLoad %uint %index
+// CHECK: [[nu50]] = OpAccessChain %_ptr_UniformConstant_type_2d_image %gTextures [[nu49]]
+// CHECK: [[nu51]] = OpLoad %type_2d_image
+// CHECK: [[nu52]] = OpSampledImage %type_sampled_image
 // CHECK:            OpImageQueryLod
     float  v9 = gTextures[NonUniformResourceIndex(index)].CalculateLevelOfDetail(gSamplers[0], 0.5);