Add validation for SPV_EXT_replicated_composites (KhronosGroup#6583)

Fixes KhronosGroup#6529

* OpConstantCompositeReplicateEXT
* OpSpecConstantCompositeReplicateEXT
* OpCompositeConstructReplicateEXT

Author: alan-baker

source/val/validate_composites.cpp

@@ -409,6 +409,67 @@ spv_result_t ValidateCompositeConstruct(ValidationState_t& _,
   return SPV_SUCCESS;
 }
 
+spv_result_t ValidateCompositeConstructReplicate(ValidationState_t& _,
+                                                 const Instruction* inst) {
+  const auto result_type = _.FindDef(inst->type_id());
+  const uint32_t operand_type = _.GetOperandTypeId(inst, 2);
+
+  switch (result_type->opcode()) {
+    case spv::Op::OpTypeVector:
+    case spv::Op::OpTypeVectorIdEXT:
+    case spv::Op::OpTypeMatrix:
+    case spv::Op::OpTypeArray:
+    case spv::Op::OpTypeCooperativeMatrixKHR:
+    case spv::Op::OpTypeCooperativeMatrixNV: {
+      const auto element_type = result_type->GetOperandAs<uint32_t>(1);
+      if (operand_type != element_type) {
+        return _.diag(SPV_ERROR_INVALID_DATA, inst)
+               << "Expected Value type to be equal to the "
+               << "result's element type";
+      }
+      break;
+    }
+    case spv::Op::OpTypeStruct: {
+      for (uint32_t operand_index = 1;
+           operand_index < result_type->operands().size(); ++operand_index) {
+        const uint32_t member_type =
+            result_type->GetOperandAs<uint32_t>(operand_index);
+        if (operand_type != member_type) {
+          return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                 << "Expected Value type to be equal to the "
+                 << "corresponding member type of the result";
+        }
+      }
+      break;
+    }
+    case spv::Op::OpTypeTensorARM: {
+      const uint32_t component_type = result_type->GetOperandAs<uint32_t>(1);
+      if (operand_type != component_type) {
+        return _.diag(SPV_ERROR_INVALID_DATA, inst)
+               << "Expected Value type to be equal to the result's element "
+                  "type";
+      }
+      if (result_type->operands().size() <= 3) {
+        return _.diag(SPV_ERROR_INVALID_DATA, inst)
+               << "Result tensor type is not a composite type because it lacks "
+                  "a shape operand";
+      }
+      break;
+    }
+    default: {
+      return _.diag(SPV_ERROR_INVALID_DATA, inst)
+             << "Expected Result Type to be a composite type";
+    }
+  }
+
+  if (_.HasCapability(spv::Capability::Shader) &&
+      _.ContainsLimitedUseIntOrFloatType(inst->type_id())) {
+    return _.diag(SPV_ERROR_INVALID_DATA, inst)
+           << "Cannot create a composite containing 8- or 16-bit types";
+  }
+  return SPV_SUCCESS;
+}
+
 spv_result_t ValidateCompositeExtract(ValidationState_t& _,
                                       const Instruction* inst) {
   uint32_t member_type = 0;
@@ -1089,6 +1150,8 @@ spv_result_t CompositesPass(ValidationState_t& _, const Instruction* inst) {
       return ValidateVectorShuffle(_, inst);
     case spv::Op::OpCompositeConstruct:
       return ValidateCompositeConstruct(_, inst);
+    case spv::Op::OpCompositeConstructReplicateEXT:
+      return ValidateCompositeConstructReplicate(_, inst);
     case spv::Op::OpCompositeExtract:
       return ValidateCompositeExtract(_, inst);
     case spv::Op::OpCompositeInsert:

source/val/validate_constants.cpp

@@ -430,6 +430,65 @@ spv_result_t ValidateConstantComposite(ValidationState_t& _,
   return SPV_SUCCESS;
 }
 
+spv_result_t ValidateConstantCompositeReplicate(ValidationState_t& _,
+                                                const Instruction* inst) {
+  std::string opcode_name = std::string("Op") + spvOpcodeString(inst->opcode());
+
+  const auto result_type = _.FindDef(inst->type_id());
+  if (!result_type || !isCompositeType(result_type)) {
+    return _.diag(SPV_ERROR_INVALID_ID, inst)
+           << opcode_name << " Result Type <id> "
+           << _.getIdName(inst->type_id()) << " is not a composite type.";
+  }
+
+  const auto constituent_id = inst->GetOperandAs<uint32_t>(2);
+  const auto constituent = _.FindDef(constituent_id);
+  if (!constituent || !spvOpcodeIsConstantOrUndef(constituent->opcode())) {
+    return _.diag(SPV_ERROR_INVALID_ID, inst)
+           << opcode_name << " Constituent <id> " << _.getIdName(constituent_id)
+           << " is not a constant or undef.";
+  }
+
+  switch (result_type->opcode()) {
+    case spv::Op::OpTypeVector:
+    case spv::Op::OpTypeVectorIdEXT:
+    case spv::Op::OpTypeMatrix:
+    case spv::Op::OpTypeArray:
+    case spv::Op::OpTypeCooperativeMatrixKHR:
+    case spv::Op::OpTypeCooperativeMatrixNV:
+    case spv::Op::OpTypeTensorARM: {
+      const auto component_type = result_type->GetOperandAs<uint32_t>(1);
+      if (component_type != constituent->type_id()) {
+        return _.diag(SPV_ERROR_INVALID_ID, inst)
+               << opcode_name << " Constituent <id> "
+               << _.getIdName(constituent_id)
+               << "s type does not match Result Type <id> "
+               << _.getIdName(result_type->id()) << "s element type.";
+      }
+      break;
+    }
+    case spv::Op::OpTypeStruct: {
+      const auto member_count = result_type->operands().size() - 1;
+      for (uint32_t member_index = 1; member_index <= member_count;
+           member_index++) {
+        const auto member_type_id =
+            result_type->GetOperandAs<uint32_t>(member_index);
+        if (member_type_id != constituent->type_id()) {
+          return _.diag(SPV_ERROR_INVALID_ID, inst)
+                 << opcode_name << " Constituent <id> "
+                 << _.getIdName(constituent_id)
+                 << " type does not match the Result Type <id> "
+                 << _.getIdName(result_type->id()) << "s member type.";
+        }
+      }
+      break;
+    }
+    default:
+      break;
+  }
+  return SPV_SUCCESS;
+}
+
 spv_result_t ValidateConstantSampler(ValidationState_t& _,
                                      const Instruction* inst) {
   const auto result_type = _.FindDef(inst->type_id());
@@ -617,6 +676,11 @@ spv_result_t ConstantPass(ValidationState_t& _, const Instruction* inst) {
     case spv::Op::OpSpecConstantComposite:
       if (auto error = ValidateConstantComposite(_, inst)) return error;
       break;
+    case spv::Op::OpConstantCompositeReplicateEXT:
+    case spv::Op::OpSpecConstantCompositeReplicateEXT:
+      if (auto error = ValidateConstantCompositeReplicate(_, inst))
+        return error;
+      break;
     case spv::Op::OpConstantSampler:
       if (auto error = ValidateConstantSampler(_, inst)) return error;
       break;

test/val/val_composites_test.cpp

@@ -33,15 +33,17 @@ using ValidateComposites = spvtest::ValidateBase<bool>;
 std::string GenerateShaderCode(
     const std::string& body,
     const std::string& capabilities_and_extensions = "",
-    const std::string& execution_model = "Fragment") {
+    const std::string& execution_model = "Fragment",
+    const std::string& extra_types = "",
+    const std::string& memory_model = "Logical GLSL450") {
   std::ostringstream ss;
   ss << R"(
 OpCapability Shader
 OpCapability Float64
 )";
 
   ss << capabilities_and_extensions;
-  ss << "OpMemoryModel Logical GLSL450\n";
+  ss << "OpMemoryModel " << memory_model << "\n";
   ss << "OpEntryPoint " << execution_model << " %main \"main\"\n";
   if (execution_model == "Fragment") {
     ss << "OpExecutionMode %main OriginUpperLeft\n";
@@ -93,7 +95,11 @@ OpCapability Float64
 
 %ptr_big_struct = OpTypePointer Uniform %big_struct
 %var_big_struct = OpVariable %ptr_big_struct Uniform
+)";
+
+  ss << extra_types;
 
+  ss << R"(
 %main = OpFunction %void None %func
 %main_entry = OpLabel
 )";
@@ -536,6 +542,201 @@ TEST_F(ValidateComposites, CompositeConstructStructWrongConstituent) {
                         "corresponding member type of Result Type struct"));
 }
 
+TEST_F(ValidateComposites, CompositeConstructReplicateVectorGood) {
+  const std::string body = R"(
+%val1 = OpCompositeConstructReplicateEXT %f32vec4 %f32_0
+)";
+
+  CompileSuccessfully(
+      GenerateShaderCode(body,
+                         "OpCapability ReplicatedCompositesEXT\nOpExtension "
+                         "\"SPV_EXT_replicated_composites\"\n")
+          .c_str());
+  ASSERT_EQ(SPV_SUCCESS, ValidateInstructions());
+}
+
+TEST_F(ValidateComposites, CompositeConstructReplicateMatrixGood) {
+  const std::string body = R"(
+%val1 = OpCompositeConstructReplicateEXT %f32mat22 %f32vec2_01
+)";
+
+  CompileSuccessfully(
+      GenerateShaderCode(body,
+                         "OpCapability ReplicatedCompositesEXT\nOpExtension "
+                         "\"SPV_EXT_replicated_composites\"\n",
+                         "Fragment")
+          .c_str());
+  ASSERT_EQ(SPV_SUCCESS, ValidateInstructions());
+}
+
+TEST_F(ValidateComposites, CompositeConstructReplicateArrayGood) {
+  const std::string body = R"(
+%val1 = OpCompositeConstructReplicateEXT %f32vec2arr3 %f32vec2_12
+)";
+
+  CompileSuccessfully(
+      GenerateShaderCode(body,
+                         "OpCapability ReplicatedCompositesEXT\nOpExtension "
+                         "\"SPV_EXT_replicated_composites\"\n")
+          .c_str());
+  ASSERT_EQ(SPV_SUCCESS, ValidateInstructions());
+}
+
+TEST_F(ValidateComposites, CompositeConstructReplicateStructGood) {
+  const std::string copy_types = R"(
+%f32struct = OpTypeStruct %f32 %f32 %f32
+)";
+
+  const std::string body = R"(
+%val1 = OpCompositeConstructReplicateEXT %f32struct %f32_0
+)";
+
+  CompileSuccessfully(
+      GenerateShaderCode(body,
+                         "OpCapability ReplicatedCompositesEXT\nOpExtension "
+                         "\"SPV_EXT_replicated_composites\"\n",
+                         "Fragment", copy_types)
+          .c_str());
+  ASSERT_EQ(SPV_SUCCESS, ValidateInstructions());
+}
+
+TEST_F(ValidateComposites, CompositeConstructReplicateCoopMatGood) {
+  const std::string extra_types = R"(
+%u32_8 = OpConstant %u32 8
+%u32_16 = OpConstant %u32 16
+%subgroup = OpConstant %u32 3
+%useA = OpConstant %u32 0
+%f32mat_nv = OpTypeCooperativeMatrixNV %f32 %subgroup %u32_8 %u32_8
+%f32mat_khr = OpTypeCooperativeMatrixKHR %f32 %subgroup %u32_16 %u32_16 %useA
+)";
+
+  const std::string body = R"(
+%val1 = OpCompositeConstructReplicateEXT %f32mat_nv %f32_0
+%val2 = OpCompositeConstructReplicateEXT %f32mat_khr %f32_0
+)";
+
+  CompileSuccessfully(
+      GenerateShaderCode(body,
+                         "OpCapability ReplicatedCompositesEXT\n"
+                         "OpCapability CooperativeMatrixNV\n"
+                         "OpCapability CooperativeMatrixKHR\n"
+                         "OpCapability VulkanMemoryModel\n"
+                         "OpCapability Float16\n"
+                         "OpExtension \"SPV_EXT_replicated_composites\"\n"
+                         "OpExtension \"SPV_NV_cooperative_matrix\"\n"
+                         "OpExtension \"SPV_KHR_cooperative_matrix\"\n"
+                         "OpExtension \"SPV_KHR_vulkan_memory_model\"\n",
+                         "Fragment", extra_types, "Logical Vulkan")
+          .c_str(),
+      SPV_ENV_UNIVERSAL_1_3);
+  ASSERT_EQ(SPV_SUCCESS, ValidateInstructions(SPV_ENV_UNIVERSAL_1_3));
+}
+
+TEST_F(ValidateComposites, CompositeConstructReplicateTensorGood) {
+  const std::string extra_types = R"(
+%arr = OpTypeArray %u32 %u32_1
+%c_arr = OpConstantNull %arr
+%tensor = OpTypeTensorARM %f32 %u32_1 %c_arr
+)";
+
+  const std::string body = R"(
+%val1 = OpCompositeConstructReplicateEXT %tensor %f32_0
+)";
+
+  CompileSuccessfully(
+      GenerateShaderCode(body,
+                         "OpCapability ReplicatedCompositesEXT\nOpCapability "
+                         "TensorsARM\nOpExtension "
+                         "\"SPV_EXT_replicated_composites\"\nOpExtension "
+                         "\"SPV_ARM_tensors\"\n",
+                         "Fragment", extra_types)
+          .c_str());
+  ASSERT_EQ(SPV_SUCCESS, ValidateInstructions());
+}
+
+TEST_F(ValidateComposites, CompositeConstructReplicateCoopMatWrongOperand) {
+  const std::string extra_types = R"(
+%u32_8 = OpConstant %u32 8
+%subgroup = OpConstant %u32 3
+%f32mat_nv = OpTypeCooperativeMatrixNV %f32 %subgroup %u32_8 %u32_8
+)";
+
+  const std::string body = R"(
+%val1 = OpCompositeConstructReplicateEXT %f32mat_nv %u32_0
+)";
+
+  CompileSuccessfully(
+      GenerateShaderCode(body,
+                         "OpCapability ReplicatedCompositesEXT\n"
+                         "OpCapability CooperativeMatrixNV\n"
+                         "OpCapability Float16\n"
+                         "OpExtension \"SPV_EXT_replicated_composites\"\n"
+                         "OpExtension \"SPV_NV_cooperative_matrix\"\n",
+                         "Fragment", extra_types)
+          .c_str());
+  ASSERT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions());
+  EXPECT_THAT(
+      getDiagnosticString(),
+      HasSubstr(
+          "Expected Value type to be equal to the result's element type"));
+}
+
+TEST_F(ValidateComposites, CompositeConstructReplicateTensorWrongOperand) {
+  const std::string extra_types = R"(
+%tensor = OpTypeTensorARM %f32
+)";
+
+  const std::string body = R"(
+%val1 = OpCompositeConstructReplicateEXT %tensor %u32_0
+)";
+
+  CompileSuccessfully(
+      GenerateShaderCode(body,
+                         "OpCapability ReplicatedCompositesEXT\nOpCapability "
+                         "TensorsARM\nOpExtension "
+                         "\"SPV_EXT_replicated_composites\"\nOpExtension "
+                         "\"SPV_ARM_tensors\"\n",
+                         "Fragment", extra_types)
+          .c_str());
+  ASSERT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions());
+  EXPECT_THAT(
+      getDiagnosticString(),
+      HasSubstr(
+          "Expected Value type to be equal to the result's element type"));
+}
+
+TEST_F(ValidateComposites, CompositeConstructReplicateWrongOperandType) {
+  const std::string body = R"(
+%val1 = OpCompositeConstructReplicateEXT %f32vec4 %u32_0
+)";
+
+  CompileSuccessfully(
+      GenerateShaderCode(body,
+                         "OpCapability ReplicatedCompositesEXT\nOpExtension "
+                         "\"SPV_EXT_replicated_composites\"\n")
+          .c_str());
+  ASSERT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions());
+  EXPECT_THAT(
+      getDiagnosticString(),
+      HasSubstr(
+          "Expected Value type to be equal to the result's element type"));
+}
+
+TEST_F(ValidateComposites, CompositeConstructReplicateNotComposite) {
+  const std::string body = R"(
+%val1 = OpCompositeConstructReplicateEXT %f32 %f32_0
+)";
+
+  CompileSuccessfully(
+      GenerateShaderCode(body,
+                         "OpCapability ReplicatedCompositesEXT\nOpExtension "
+                         "\"SPV_EXT_replicated_composites\"\n")
+          .c_str());
+  ASSERT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions());
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("Expected Result Type to be a composite type"));
+}
+
 TEST_F(ValidateComposites, CopyObjectSuccess) {
   const std::string body = R"(
 %val1 = OpCopyObject %f32 %f32_0