[SPIR-V] Fix an issue where OpConstantComposite is being used with spec constants (#5263)

Fixes #4849

Author: godlikepanos

tools/clang/lib/SPIRV/EmitVisitor.cpp

@@ -1963,7 +1963,7 @@ bool EmitVisitor::visit(SpirvIntrinsicInstruction *inst) {
   return true;
 }
 
-bool EmitVisitor::visit(SpirvEmitMeshTasksEXT *inst) { 
+bool EmitVisitor::visit(SpirvEmitMeshTasksEXT *inst) {
   initInstruction(inst);
 
   curInst.push_back(getOrAssignResultId<SpirvInstruction>(inst->getXDimension()));
@@ -2310,7 +2310,8 @@ EmitTypeHandler::getOrCreateConstantComposite(SpirvConstantComposite *inst) {
   } else {
     // Constant wasn't emitted in the past.
     const uint32_t typeId = emitType(inst->getResultType());
-    initTypeInstruction(spv::Op::OpConstantComposite);
+    initTypeInstruction(isSpecConst ? spv::Op::OpSpecConstantComposite
+                                    : spv::Op::OpConstantComposite);
     curTypeInst.push_back(typeId);
     curTypeInst.push_back(getOrAssignResultId<SpirvInstruction>(inst));
     for (auto constituent : inst->getConstituents())

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -6938,7 +6938,8 @@ SpirvInstruction *SpirvEmitter::createVectorSplat(const Expr *scalarExpr,
   // Should find a more meaningful one.
   if (auto *constVal = dyn_cast<SpirvConstant>(scalarVal)) {
     llvm::SmallVector<SpirvConstant *, 4> elements(size_t(size), constVal);
-    auto *value = spvBuilder.getConstantComposite(vecType, elements);
+    const bool isSpecConst = constVal->getopcode() == spv::Op::OpSpecConstant;
+    auto *value = spvBuilder.getConstantComposite(vecType, elements, isSpecConst);
     value->setRValue();
     return value;
   } else {
@@ -12018,7 +12019,7 @@ SpirvEmitter::getSpirvShaderStage(hlsl::ShaderModel::Kind smk, bool extMeshShadi
     return spv::ExecutionModel::CallableNV;
   case hlsl::ShaderModel::Kind::Mesh:
     return extMeshShading ?
-           spv::ExecutionModel::MeshEXT: 
+           spv::ExecutionModel::MeshEXT:
            spv::ExecutionModel::MeshNV;
   case hlsl::ShaderModel::Kind::Amplification:
     return extMeshShading ?

tools/clang/test/CodeGenSPIRV/vk.spec-constant.composite.hlsl

@@ -0,0 +1,12 @@
+// RUN: %dxc -T cs_6_0 -E main
+
+[[vk::constant_id(0)]] const uint kConst = 1;
+
+RWTexture2D<float3> img;
+
+[numthreads(1, 1, 1)]
+void main(uint3 svDispatchThreadId : SV_DISPATCHTHREADID)
+{
+	// CHECK: OpSpecConstantComposite %v2uint %kConst %kConst
+    img[svDispatchThreadId.xy * kConst] = 0;
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -2264,6 +2264,9 @@ TEST_F(FileTest, VulkanSpecConstantError5) {
 TEST_F(FileTest, VulkanSpecConstantErrorNotSegfault) {
   runFileTest("vk.spec-constant.error.not.segfault.hlsl", Expect::Failure);
 }
+TEST_F(FileTest, VulkanSpecConstantComposite) {
+  runFileTest("vk.spec-constant.composite.hlsl");
+}
 
 TEST_F(FileTest, VulkanLayoutCBufferMatrixZpr) {
   runFileTest("vk.layout.cbuffer.zpr.hlsl");