Re-enable ValidationAfterOverrideStorageSize test

Also move Robustness transform to pipeline creation time on Vulkan
backend.

Bug: tint:1660, dawn:1041
Change-Id: I10220f34119d11f29be86fd29463a282175eccdd
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/103780
Reviewed-by: Austin Eng <enga@chromium.org>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Shrek Shao <shrekshao@google.com>
diff --git a/src/dawn/native/null/DeviceNull.cpp b/src/dawn/native/null/DeviceNull.cpp
index badb570..df2106b 100644
--- a/src/dawn/native/null/DeviceNull.cpp
+++ b/src/dawn/native/null/DeviceNull.cpp
@@ -390,10 +390,12 @@
 
     tint::Program transformedProgram;
     const tint::Program* program;
-    if (!computeStage.metadata->overrides.empty()) {
-        tint::transform::Manager transformManager;
-        tint::transform::DataMap transformInputs;
+    tint::transform::Manager transformManager;
+    tint::transform::DataMap transformInputs;
 
+    transformManager.Add<tint::transform::Robustness>();
+
+    if (!computeStage.metadata->overrides.empty()) {
         transformManager.Add<tint::transform::SingleEntryPoint>();
         transformInputs.Add<tint::transform::SingleEntryPoint::Config>(
             computeStage.entryPoint.c_str());
@@ -403,16 +405,14 @@
         transformManager.Add<tint::transform::SubstituteOverride>();
         transformInputs.Add<tint::transform::SubstituteOverride::Config>(
             BuildSubstituteOverridesTransformConfig(computeStage));
-
-        DAWN_TRY_ASSIGN(transformedProgram,
-                        RunTransforms(&transformManager, computeStage.module->GetTintProgram(),
-                                      transformInputs, nullptr, nullptr));
-
-        program = &transformedProgram;
-    } else {
-        program = computeStage.module->GetTintProgram();
     }
 
+    DAWN_TRY_ASSIGN(transformedProgram,
+                    RunTransforms(&transformManager, computeStage.module->GetTintProgram(),
+                                  transformInputs, nullptr, nullptr));
+
+    program = &transformedProgram;
+
     // Do the workgroup size validation as it is actually backend agnostic.
     const CombinedLimits& limits = GetDevice()->GetLimits();
     Extent3D _;
diff --git a/src/dawn/native/vulkan/ShaderModuleVk.cpp b/src/dawn/native/vulkan/ShaderModuleVk.cpp
index 752338e..b5796d6 100644
--- a/src/dawn/native/vulkan/ShaderModuleVk.cpp
+++ b/src/dawn/native/vulkan/ShaderModuleVk.cpp
@@ -172,19 +172,7 @@
 
 MaybeError ShaderModule::Initialize(ShaderModuleParseResult* parseResult,
                                     OwnedCompilationMessages* compilationMessages) {
-    if (GetDevice()->IsRobustnessEnabled()) {
-        ScopedTintICEHandler scopedICEHandler(GetDevice());
-
-        tint::transform::Robustness robustness;
-        tint::transform::DataMap transformInputs;
-
-        tint::Program program;
-        DAWN_TRY_ASSIGN(program, RunTransforms(&robustness, parseResult->tintProgram.get(),
-                                               transformInputs, nullptr, nullptr));
-        // Rather than use a new ParseResult object, we just reuse the original parseResult
-        parseResult->tintProgram = std::make_unique<tint::Program>(std::move(program));
-    }
-
+    ScopedTintICEHandler scopedICEHandler(GetDevice());
     return InitializeBase(parseResult, compilationMessages);
 }
 
@@ -204,6 +192,7 @@
     X(std::optional<tint::transform::SubstituteOverride::Config>, substituteOverrideConfig) \
     X(LimitsForCompilationRequest, limits)                                                  \
     X(std::string_view, entryPointName)                                                     \
+    X(bool, isRobustnessEnabled)                                                            \
     X(bool, disableWorkgroupInit)                                                           \
     X(bool, useZeroInitializeWorkgroupMemoryExtension)                                      \
     X(CacheKey::UnsafeUnkeyedValue<dawn::platform::Platform*>, tracePlatform)
@@ -284,6 +273,7 @@
     req.bindingPoints = std::move(bindingPoints);
     req.newBindingsMap = std::move(newBindingsMap);
     req.entryPointName = programmableStage.entryPoint;
+    req.isRobustnessEnabled = GetDevice()->IsRobustnessEnabled();
     req.disableWorkgroupInit = GetDevice()->IsToggleEnabled(Toggle::DisableWorkgroupInit);
     req.useZeroInitializeWorkgroupMemoryExtension =
         GetDevice()->IsToggleEnabled(Toggle::VulkanUseZeroInitializeWorkgroupMemoryExtension);
@@ -298,6 +288,9 @@
         spirv, GetDevice(), std::move(req), Spirv::FromBlob,
         [](SpirvCompilationRequest r) -> ResultOrError<Spirv> {
             tint::transform::Manager transformManager;
+            if (r.isRobustnessEnabled) {
+                transformManager.append(std::make_unique<tint::transform::Robustness>());
+            }
             // Many Vulkan drivers can't handle multi-entrypoint shader modules.
             transformManager.append(std::make_unique<tint::transform::SingleEntryPoint>());
             // Run the binding remapper after SingleEntryPoint to avoid collisions with
diff --git a/src/dawn/tests/end2end/ShaderValidationTests.cpp b/src/dawn/tests/end2end/ShaderValidationTests.cpp
index 302a6ba..9f964c9 100644
--- a/src/dawn/tests/end2end/ShaderValidationTests.cpp
+++ b/src/dawn/tests/end2end/ShaderValidationTests.cpp
@@ -330,8 +330,7 @@
 
 // Test workgroup size validation after being overrided with invalid values (storage size limits
 // validation).
-// TODO(tint:1660): re-enable after override can be used as array size.
-TEST_P(WorkgroupSizeValidationTest, DISABLED_ValidationAfterOverrideStorageSize) {
+TEST_P(WorkgroupSizeValidationTest, ValidationAfterOverrideStorageSize) {
     wgpu::Limits supportedLimits = GetSupportedLimits().limits;
 
     constexpr uint32_t kVec4Size = 16;
@@ -347,11 +346,11 @@
         ss << "override b: u32;";
         if (vec4_count > 0) {
             ss << "var<workgroup> vec4_data: array<vec4<f32>, a>;";
-            body << "_ = vec4_data;";
+            body << "_ = vec4_data[0];";
         }
         if (mat4_count > 0) {
             ss << "var<workgroup> mat4_data: array<mat4x4<f32>, b>;";
-            body << "_ = mat4_data;";
+            body << "_ = mat4_data[0];";
         }
         ss << "@compute @workgroup_size(1) fn main() { " << body.str() << " }";