Update OpMemoryBarriers rules for WebGPU (#2775)

Part of #2724
diff --git a/source/val/validate_memory_semantics.cpp b/source/val/validate_memory_semantics.cpp
index c8df405..a0de18f 100644
--- a/source/val/validate_memory_semantics.cpp
+++ b/source/val/validate_memory_semantics.cpp
@@ -57,59 +57,77 @@
   }
 
   if (spvIsWebGPUEnv(_.context()->target_env)) {
-    if (inst->opcode() == SpvOpControlBarrier) {
-      if (!(value & SpvMemorySemanticsAcquireReleaseMask)) {
-        return _.diag(SPV_ERROR_INVALID_DATA, inst)
-               << "For WebGPU, AcquireRelease must be set for Memory Semantics "
-                  "of OpControlBarrier.";
-      }
-
-      if (!(value & SpvMemorySemanticsWorkgroupMemoryMask)) {
-        return _.diag(SPV_ERROR_INVALID_DATA, inst)
-               << "For WebGPU, WorkgroupMemory must be set for Memory "
-                  "Semantics of OpControlBarrier.";
-      }
-
-      uint32_t valid_bits = SpvMemorySemanticsAcquireReleaseMask |
-                            SpvMemorySemanticsWorkgroupMemoryMask;
-      if (value & ~valid_bits) {
-        return _.diag(SPV_ERROR_INVALID_DATA, inst)
-               << "For WebGPU only WorkgroupMemory and AcquireRelease may be "
-                  "set for Memory Semantics of OpControlBarrier.";
-      }
-    } else {
-      // TODO(2723): Rewrite this to be in the style of above and simplify.
-      uint32_t valid_bits = SpvMemorySemanticsUniformMemoryMask |
-                            SpvMemorySemanticsWorkgroupMemoryMask |
-                            SpvMemorySemanticsImageMemoryMask |
-                            SpvMemorySemanticsOutputMemoryKHRMask |
-                            SpvMemorySemanticsMakeAvailableKHRMask |
-                            SpvMemorySemanticsMakeVisibleKHRMask;
-      if (!spvOpcodeIsAtomicOp(inst->opcode())) {
-        valid_bits |= SpvMemorySemanticsAcquireReleaseMask;
-      }
-
-      if (value & ~valid_bits) {
-        if (spvOpcodeIsAtomicOp(inst->opcode())) {
+    uint32_t valid_bits;
+    switch (inst->opcode()) {
+      case SpvOpControlBarrier:
+        if (!(value & SpvMemorySemanticsAcquireReleaseMask)) {
           return _.diag(SPV_ERROR_INVALID_DATA, inst)
-                 << "WebGPU spec disallows, for OpAtomic*, any bit masks in "
-                    "Memory Semantics that are not UniformMemory, "
-                    "WorkgroupMemory, ImageMemory, or OutputMemoryKHR";
-        } else {
-          return _.diag(SPV_ERROR_INVALID_DATA, inst)
-                 << "WebGPU spec disallows any bit masks in Memory Semantics "
-                    "that are not AcquireRelease, UniformMemory, "
-                    "WorkgroupMemory, ImageMemory, OutputMemoryKHR, "
-                    "MakeAvailableKHR, or MakeVisibleKHR";
+                 << "For WebGPU, AcquireRelease must be set for Memory "
+                    "Semantics of OpControlBarrier.";
         }
-      }
 
-      if (!spvOpcodeIsAtomicOp(inst->opcode()) &&
-          !(value & SpvMemorySemanticsAcquireReleaseMask)) {
-        return _.diag(SPV_ERROR_INVALID_DATA, inst)
-               << "WebGPU spec requires AcquireRelease to set in Memory "
-                  "Semantics.";
-      }
+        if (!(value & SpvMemorySemanticsWorkgroupMemoryMask)) {
+          return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                 << "For WebGPU, WorkgroupMemory must be set for Memory "
+                    "Semantics of OpControlBarrier.";
+        }
+
+        valid_bits = SpvMemorySemanticsAcquireReleaseMask |
+                     SpvMemorySemanticsWorkgroupMemoryMask;
+        if (value & ~valid_bits) {
+          return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                 << "For WebGPU only WorkgroupMemory and AcquireRelease may be "
+                    "set for Memory Semantics of OpControlBarrier.";
+        }
+        break;
+      case SpvOpMemoryBarrier:
+        if (!(value & SpvMemorySemanticsImageMemoryMask)) {
+          return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                 << "For WebGPU, ImageMemory must be set for Memory Semantics "
+                    "of OpMemoryBarrier.";
+        }
+        valid_bits = SpvMemorySemanticsImageMemoryMask;
+        if (value & ~valid_bits) {
+          return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                 << "For WebGPU only ImageMemory may be set for Memory "
+                    "Semantics "
+                    "of OpMemoryBarrier.";
+        }
+        break;
+      default:
+        // TODO(2723): Rewrite this to be in the style of above and simplify.
+        valid_bits = SpvMemorySemanticsUniformMemoryMask |
+                     SpvMemorySemanticsWorkgroupMemoryMask |
+                     SpvMemorySemanticsImageMemoryMask |
+                     SpvMemorySemanticsOutputMemoryKHRMask |
+                     SpvMemorySemanticsMakeAvailableKHRMask |
+                     SpvMemorySemanticsMakeVisibleKHRMask;
+        if (!spvOpcodeIsAtomicOp(inst->opcode())) {
+          valid_bits |= SpvMemorySemanticsAcquireReleaseMask;
+        }
+
+        if (value & ~valid_bits) {
+          if (spvOpcodeIsAtomicOp(inst->opcode())) {
+            return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                   << "WebGPU spec disallows, for OpAtomic*, any bit masks in "
+                      "Memory Semantics that are not UniformMemory, "
+                      "WorkgroupMemory, ImageMemory, or OutputMemoryKHR";
+          } else {
+            return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                   << "WebGPU spec disallows any bit masks in Memory Semantics "
+                      "that are not AcquireRelease, UniformMemory, "
+                      "WorkgroupMemory, ImageMemory, OutputMemoryKHR, "
+                      "MakeAvailableKHR, or MakeVisibleKHR";
+          }
+        }
+
+        if (!spvOpcodeIsAtomicOp(inst->opcode()) &&
+            !(value & SpvMemorySemanticsAcquireReleaseMask)) {
+          return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                 << "WebGPU spec requires AcquireRelease to set in Memory "
+                    "Semantics.";
+        }
+        break;
     }
   }
 
diff --git a/source/val/validate_scopes.cpp b/source/val/validate_scopes.cpp
index 32a52a8..47c4e45 100644
--- a/source/val/validate_scopes.cpp
+++ b/source/val/validate_scopes.cpp
@@ -228,21 +228,32 @@
 
   // WebGPU specific rules
   if (spvIsWebGPUEnv(_.context()->target_env)) {
-    if (inst->opcode() == SpvOpControlBarrier) {
-      if (value != SpvScopeWorkgroup) {
-        return _.diag(SPV_ERROR_INVALID_DATA, inst)
-               << spvOpcodeString(opcode)
-               << ": in WebGPU environment Memory Scope is limited to "
-               << "Workgroup for OpControlBarrier";
-      }
-    } else {
-      if (value != SpvScopeWorkgroup && value != SpvScopeInvocation &&
-          value != SpvScopeQueueFamilyKHR) {
-        return _.diag(SPV_ERROR_INVALID_DATA, inst)
-               << spvOpcodeString(opcode)
-               << ": in WebGPU environment Memory Scope is limited to "
-               << "Workgroup, Invocation, and QueueFamilyKHR";
-      }
+    switch (inst->opcode()) {
+      case SpvOpControlBarrier:
+        if (value != SpvScopeWorkgroup) {
+          return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                 << spvOpcodeString(opcode)
+                 << ": in WebGPU environment Memory Scope is limited to "
+                 << "Workgroup for OpControlBarrier";
+        }
+        break;
+      case SpvOpMemoryBarrier:
+        if (value != SpvScopeWorkgroup) {
+          return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                 << spvOpcodeString(opcode)
+                 << ": in WebGPU environment Memory Scope is limited to "
+                 << "Workgroup for OpMemoryBarrier";
+        }
+        break;
+      default:
+        if (value != SpvScopeWorkgroup && value != SpvScopeInvocation &&
+            value != SpvScopeQueueFamilyKHR) {
+          return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                 << spvOpcodeString(opcode)
+                 << ": in WebGPU environment Memory Scope is limited to "
+                 << "Workgroup, Invocation, and QueueFamilyKHR";
+        }
+        break;
     }
   }
 
diff --git a/test/val/val_barriers_test.cpp b/test/val/val_barriers_test.cpp
index 1d71e90..18f57f8 100644
--- a/test/val/val_barriers_test.cpp
+++ b/test/val/val_barriers_test.cpp
@@ -86,7 +86,8 @@
 %uniform = OpConstant %u32 64
 %uniform_workgroup = OpConstant %u32 320
 %workgroup_memory = OpConstant %u32 256
-
+%image_memory = OpConstant %u32 2048
+%uniform_image_memory = OpConstant %u32 2112
 
 %main = OpFunction %void None %func
 %main_entry = OpLabel
@@ -708,13 +709,37 @@
   ASSERT_EQ(SPV_SUCCESS, ValidateInstructions(SPV_ENV_VULKAN_1_0));
 }
 
-TEST_F(ValidateBarriers, OpMemoryBarrierWebGPUAcquireReleaseSuccess) {
+TEST_F(ValidateBarriers, OpMemoryBarrierWebGPUImageMemorySuccess) {
+  const std::string body = R"(
+OpMemoryBarrier %workgroup %image_memory
+)";
+
+  CompileSuccessfully(GenerateWebGPUShaderCode(body), SPV_ENV_WEBGPU_0);
+  ASSERT_EQ(SPV_SUCCESS, ValidateInstructions(SPV_ENV_WEBGPU_0));
+}
+
+TEST_F(ValidateBarriers, OpMemoryBarrierWebGPUDeviceFailure) {
+  const std::string body = R"(
+OpMemoryBarrier %subgroup %image_memory
+)";
+
+  CompileSuccessfully(GenerateWebGPUShaderCode(body), SPV_ENV_WEBGPU_0);
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions(SPV_ENV_WEBGPU_0));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("in WebGPU environment Memory Scope is limited to "
+                        "Workgroup for OpMemoryBarrier"));
+}
+
+TEST_F(ValidateBarriers, OpMemoryBarrierWebGPUAcquireReleaseFailure) {
   const std::string body = R"(
 OpMemoryBarrier %workgroup %acquire_release_uniform_workgroup
 )";
 
   CompileSuccessfully(GenerateWebGPUShaderCode(body), SPV_ENV_WEBGPU_0);
-  ASSERT_EQ(SPV_SUCCESS, ValidateInstructions(SPV_ENV_WEBGPU_0));
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions(SPV_ENV_WEBGPU_0));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("ImageMemory must be set for Memory Semantics of "
+                        "OpMemoryBarrier"));
 }
 
 TEST_F(ValidateBarriers, OpMemoryBarrierWebGPURelaxedFailure) {
@@ -725,7 +750,8 @@
   CompileSuccessfully(GenerateWebGPUShaderCode(body), SPV_ENV_WEBGPU_0);
   EXPECT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions(SPV_ENV_WEBGPU_0));
   EXPECT_THAT(getDiagnosticString(),
-              HasSubstr("WebGPU spec requires AcquireRelease to set"));
+              HasSubstr("ImageMemory must be set for Memory Semantics of "
+                        "OpMemoryBarrier"));
 }
 
 TEST_F(ValidateBarriers, OpMemoryBarrierWebGPUAcquireFailure) {
@@ -735,9 +761,9 @@
 
   CompileSuccessfully(GenerateWebGPUShaderCode(body), SPV_ENV_WEBGPU_0);
   EXPECT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions(SPV_ENV_WEBGPU_0));
-  EXPECT_THAT(
-      getDiagnosticString(),
-      HasSubstr("WebGPU spec disallows any bit masks in Memory Semantics"));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("ImageMemory must be set for Memory Semantics of "
+                        "OpMemoryBarrier"));
 }
 
 TEST_F(ValidateBarriers, OpMemoryBarrierWebGPUReleaseFailure) {
@@ -747,9 +773,21 @@
 
   CompileSuccessfully(GenerateWebGPUShaderCode(body), SPV_ENV_WEBGPU_0);
   EXPECT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions(SPV_ENV_WEBGPU_0));
-  EXPECT_THAT(
-      getDiagnosticString(),
-      HasSubstr("WebGPU spec disallows any bit masks in Memory Semantics"));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("ImageMemory must be set for Memory Semantics of "
+                        "OpMemoryBarrier"));
+}
+
+TEST_F(ValidateBarriers, OpMemoryBarrierWebGPUUniformFailure) {
+  const std::string body = R"(
+OpMemoryBarrier %workgroup %uniform_image_memory
+)";
+
+  CompileSuccessfully(GenerateWebGPUShaderCode(body), SPV_ENV_WEBGPU_0);
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA, ValidateInstructions(SPV_ENV_WEBGPU_0));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("only ImageMemory may be set for Memory Semantics of "
+                        "OpMemoryBarrier"));
 }
 
 TEST_F(ValidateBarriers, OpMemoryBarrierFloatMemoryScope) {