Fix PrimitiveId builtin check for Vulkan

According to Vulkan spec 1.1.72:

> The PrimitiveId decoration must be used only within fragment,
> tessellation control, tessellation evaluation, and geometry shaders.

> In a tessellation control or tessellation evaluation shader, any
> variable decorated with PrimitiveId must be declared using the Input
> storage class.

We were enforcing that PrimitiveId can only be used with Output
storage class for TCS and TES before.
diff --git a/source/validate_builtins.cpp b/source/validate_builtins.cpp
index e739603..c494a88 100644
--- a/source/validate_builtins.cpp
+++ b/source/validate_builtins.cpp
@@ -1478,26 +1478,22 @@
              << " " << GetStorageClassDesc(referenced_from_inst);
     }
 
-    if (storage_class == SpvStorageClassInput) {
+    if (storage_class == SpvStorageClassOutput) {
       assert(function_id_ == 0);
       id_to_at_reference_checks_[referenced_from_inst.id()].push_back(std::bind(
           &BuiltInsValidator::ValidateNotCalledWithExecutionModel, this,
           "Vulkan spec doesn't allow BuiltIn PrimitiveId to be used for "
-          "variables with Input storage class if execution model is "
+          "variables with Output storage class if execution model is "
           "TessellationControl.",
           SpvExecutionModelTessellationControl, decoration, built_in_inst,
           referenced_from_inst, std::placeholders::_1));
       id_to_at_reference_checks_[referenced_from_inst.id()].push_back(std::bind(
           &BuiltInsValidator::ValidateNotCalledWithExecutionModel, this,
           "Vulkan spec doesn't allow BuiltIn PrimitiveId to be used for "
-          "variables with Input storage class if execution model is "
+          "variables with Output storage class if execution model is "
           "TessellationEvaluation.",
           SpvExecutionModelTessellationEvaluation, decoration, built_in_inst,
           referenced_from_inst, std::placeholders::_1));
-    }
-
-    if (storage_class == SpvStorageClassOutput) {
-      assert(function_id_ == 0);
       id_to_at_reference_checks_[referenced_from_inst.id()].push_back(std::bind(
           &BuiltInsValidator::ValidateNotCalledWithExecutionModel, this,
           "Vulkan spec doesn't allow BuiltIn PrimitiveId to be used for "
diff --git a/test/val/val_builtins_test.cpp b/test/val/val_builtins_test.cpp
index 9f5b1d5..78e429d 100644
--- a/test/val/val_builtins_test.cpp
+++ b/test/val/val_builtins_test.cpp
@@ -1007,15 +1007,16 @@
 INSTANTIATE_TEST_CASE_P(
     PrimitiveIdInputSuccess,
     ValidateVulkanCombineBuiltInExecutionModelDataTypeResult,
-    Combine(Values("PrimitiveId"), Values("Fragment", "Geometry"),
+    Combine(Values("PrimitiveId"),
+            Values("Fragment", "TessellationControl", "TessellationEvaluation",
+                   "Geometry"),
             Values("Input"), Values("%u32"), Values(TestResult())), );
 
 INSTANTIATE_TEST_CASE_P(
     PrimitiveIdOutputSuccess,
     ValidateVulkanCombineBuiltInExecutionModelDataTypeResult,
-    Combine(Values("PrimitiveId"),
-            Values("Geometry", "TessellationControl", "TessellationEvaluation"),
-            Values("Output"), Values("%u32"), Values(TestResult())), );
+    Combine(Values("PrimitiveId"), Values("Geometry"), Values("Output"),
+            Values("%u32"), Values(TestResult())), );
 
 INSTANTIATE_TEST_CASE_P(
     PrimitiveIdInvalidExecutionModel,
@@ -1028,7 +1029,7 @@
                 "TessellationEvaluation or Geometry execution models"))), );
 
 INSTANTIATE_TEST_CASE_P(
-    PrimitiveIdNotInput,
+    PrimitiveIdFragmentNotInput,
     ValidateVulkanCombineBuiltInExecutionModelDataTypeResult,
     Combine(
         Values("PrimitiveId"), Values("Fragment"), Values("Output"),
@@ -1038,14 +1039,14 @@
                           "which is called with execution model Fragment"))), );
 
 INSTANTIATE_TEST_CASE_P(
-    PrimitiveIdGeometryNotOutput,
+    PrimitiveIdGeometryNotInput,
     ValidateVulkanCombineBuiltInExecutionModelDataTypeResult,
     Combine(Values("PrimitiveId"),
             Values("TessellationControl", "TessellationEvaluation"),
-            Values("Input"), Values("%u32"),
+            Values("Output"), Values("%u32"),
             Values(TestResult(
                 SPV_ERROR_INVALID_DATA,
-                "Input storage class if execution model is Tessellation",
+                "Output storage class if execution model is Tessellation",
                 "which is called with execution model Tessellation"))), );
 
 INSTANTIATE_TEST_CASE_P(