From c178159160404ec25f57ab8dd8ba450de75ce140 Mon Sep 17 00:00:00 2001
From: Natalie Chouinard <chouinard@google.com>
Date: Thu, 16 Jul 2026 17:53:11 -0700
Subject: [PATCH] [tint][msl] Add SwitchReturn IR transform

Adds an MSL transform pass that wraps early returns nested inside switch
blocks with a volatile-guarded conditional. This works around a driver
bug in MTLCompilerService on macOS.

Bug: 508638064
Change-Id: I8c8b7cb0a3710f855499fdfa94872117a8d1bd23
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/325075
Reviewed-by: dan sinclair <dsinclair@chromium.org>
Commit-Queue: Natalie Chouinard <chouinard@google.com>
---

diff --git a/src/dawn/tests/end2end/ComputeFlowControlTests.cpp b/src/dawn/tests/end2end/ComputeFlowControlTests.cpp
index d29b228..cf9ebcc 100644
--- a/src/dawn/tests/end2end/ComputeFlowControlTests.cpp
+++ b/src/dawn/tests/end2end/ComputeFlowControlTests.cpp
@@ -513,6 +513,47 @@
     RunTest(shader, inputs, expected);
 }
 
+// This is a regression test for a crash in MTLCompilerService on macOS (crbug.com/508638064).
+TEST_P(ComputeFlowControlTests, SwitchReturnMTLCompilerServiceCrash) {
+    const char* shader = R"(
+@group(0) @binding(0) var<uniform> u: i32;
+@group(0) @binding(1) var<storage, read_write> s: i32;
+
+fn foo(){
+    switch (s) {
+        case 0i: {}
+        case 2i: {
+            return;
+        }
+        case 4i: {}
+        default: {
+            s = s / u;
+        }
+    }
+    switch (s) {
+        case 0i: {
+            s = 0;
+        }
+        case 1i: {}
+        case 4i: {}
+        default: {
+            s = 0;
+        }
+    }
+}
+
+@compute @workgroup_size(1)
+fn main() {
+    foo();
+    foo();
+}
+)";
+
+    wgpu::ComputePipelineDescriptor csDesc;
+    csDesc.compute.module = utils::CreateShaderModule(device, shader);
+    device.CreateComputePipeline(&csDesc);
+}
+
 DAWN_INSTANTIATE_TEST(ComputeFlowControlTests,
                       D3D11Backend(),
                       D3D12Backend(),
diff --git a/src/tint/lang/msl/writer/raise/BUILD.bazel b/src/tint/lang/msl/writer/raise/BUILD.bazel
index ad45376..d8b1bb8 100644
--- a/src/tint/lang/msl/writer/raise/BUILD.bazel
+++ b/src/tint/lang/msl/writer/raise/BUILD.bazel
@@ -54,6 +54,7 @@
     "raise.cc",
     "shader_io.cc",
     "simd_ballot.cc",
+    "switch_return.cc",
     "validate_subgroup_matrix.cc",
   ],
   hdrs = [
@@ -68,6 +69,7 @@
     "raise.h",
     "shader_io.h",
     "simd_ballot.h",
+    "switch_return.h",
     "validate_subgroup_matrix.h",
   ],
   deps = [
@@ -114,6 +116,7 @@
     "module_scope_vars_test.cc",
     "shader_io_test.cc",
     "simd_ballot_test.cc",
+    "switch_return_test.cc",
     "validate_subgroup_matrix_test.cc",
   ],
   deps = [
diff --git a/src/tint/lang/msl/writer/raise/BUILD.cmake b/src/tint/lang/msl/writer/raise/BUILD.cmake
index d3772ca..db2ac9d 100644
--- a/src/tint/lang/msl/writer/raise/BUILD.cmake
+++ b/src/tint/lang/msl/writer/raise/BUILD.cmake
@@ -63,6 +63,8 @@
   lang/msl/writer/raise/shader_io.h
   lang/msl/writer/raise/simd_ballot.cc
   lang/msl/writer/raise/simd_ballot.h
+  lang/msl/writer/raise/switch_return.cc
+  lang/msl/writer/raise/switch_return.h
   lang/msl/writer/raise/validate_subgroup_matrix.cc
   lang/msl/writer/raise/validate_subgroup_matrix.h
 )
@@ -116,6 +118,7 @@
   lang/msl/writer/raise/module_scope_vars_test.cc
   lang/msl/writer/raise/shader_io_test.cc
   lang/msl/writer/raise/simd_ballot_test.cc
+  lang/msl/writer/raise/switch_return_test.cc
   lang/msl/writer/raise/validate_subgroup_matrix_test.cc
 )
 
diff --git a/src/tint/lang/msl/writer/raise/BUILD.gn b/src/tint/lang/msl/writer/raise/BUILD.gn
index bb7d507..85f27b8 100644
--- a/src/tint/lang/msl/writer/raise/BUILD.gn
+++ b/src/tint/lang/msl/writer/raise/BUILD.gn
@@ -67,6 +67,8 @@
       "shader_io.h",
       "simd_ballot.cc",
       "simd_ballot.h",
+      "switch_return.cc",
+      "switch_return.h",
       "validate_subgroup_matrix.cc",
       "validate_subgroup_matrix.h",
     ]
@@ -113,6 +115,7 @@
         "module_scope_vars_test.cc",
         "shader_io_test.cc",
         "simd_ballot_test.cc",
+        "switch_return_test.cc",
         "validate_subgroup_matrix_test.cc",
       ]
       deps = [
diff --git a/src/tint/lang/msl/writer/raise/raise.cc b/src/tint/lang/msl/writer/raise/raise.cc
index e340bcf..71400bc 100644
--- a/src/tint/lang/msl/writer/raise/raise.cc
+++ b/src/tint/lang/msl/writer/raise/raise.cc
@@ -72,6 +72,7 @@
 #include "src/tint/lang/msl/writer/raise/module_scope_vars.h"
 #include "src/tint/lang/msl/writer/raise/shader_io.h"
 #include "src/tint/lang/msl/writer/raise/simd_ballot.h"
+#include "src/tint/lang/msl/writer/raise/switch_return.h"
 #include "src/tint/lang/msl/writer/raise/validate_subgroup_matrix.h"
 
 namespace tint::msl::writer {
@@ -302,6 +303,8 @@
         options.workarounds.disable_module_constant_f16};
     TINT_CHECK_RESULT(raise::ModuleConstant(module, module_const_config));
 
+    TINT_CHECK_RESULT(raise::SwitchReturn(module));
+
     // These transforms need to be run last as various transforms introduce terminator arguments,
     // naming conflicts, and expressions that need to be explicitly not inlined.
     TINT_CHECK_RESULT(core::ir::transform::RemoveTerminatorArgs(module));
diff --git a/src/tint/lang/msl/writer/raise/switch_return.cc b/src/tint/lang/msl/writer/raise/switch_return.cc
new file mode 100644
index 0000000..ca90580
--- /dev/null
+++ b/src/tint/lang/msl/writer/raise/switch_return.cc
@@ -0,0 +1,97 @@
+// Copyright 2026 The Dawn & Tint Authors
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+//    list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+//    this list of conditions and the following disclaimer in the documentation
+//    and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+//    contributors may be used to endorse or promote products derived from
+//    this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "src/tint/lang/msl/writer/raise/switch_return.h"
+
+#include "src/tint/lang/core/ir/builder.h"
+#include "src/tint/lang/core/ir/module.h"
+#include "src/tint/lang/core/ir/validator.h"
+#include "src/tint/lang/core/type/u32.h"
+#include "src/tint/lang/msl/ir/builtin_call.h"
+

diff --git a/src/dawn/tests/end2end/ComputeFlowControlTests.cpp b/src/dawn/tests/end2end/ComputeFlowControlTests.cpp
index d29b228..cf9ebcc 100644
--- a/src/dawn/tests/end2end/ComputeFlowControlTests.cpp
+++ b/src/dawn/tests/end2end/ComputeFlowControlTests.cpp
@@ -513,6 +513,47 @@
     RunTest(shader, inputs, expected);
 }
 
+// This is a regression test for a crash in MTLCompilerService on macOS (crbug.com/508638064).
+TEST_P(ComputeFlowControlTests, SwitchReturnMTLCompilerServiceCrash) {
+    const char* shader = R"(
+@group(0) @binding(0) var<uniform> u: i32;
+@group(0) @binding(1) var<storage, read_write> s: i32;
+
+fn foo(){
+    switch (s) {
+        case 0i: {}
+        case 2i: {
+            return;
+        }
+        case 4i: {}
+        default: {
+            s = s / u;
+        }
+    }
+    switch (s) {
+        case 0i: {
+            s = 0;
+        }
+        case 1i: {}
+        case 4i: {}
+        default: {
+            s = 0;
+        }
+    }
+}
+
+@compute @workgroup_size(1)
+fn main() {
+    foo();
+    foo();
+}
+)";
+
+    wgpu::ComputePipelineDescriptor csDesc;
+    csDesc.compute.module = utils::CreateShaderModule(device, shader);
+    device.CreateComputePipeline(&csDesc);
+}
+
 DAWN_INSTANTIATE_TEST(ComputeFlowControlTests,
                       D3D11Backend(),
                       D3D12Backend(),
diff --git a/src/tint/lang/msl/writer/raise/switch_return_test.cc b/src/tint/lang/msl/writer/raise/switch_return_test.cc
new file mode 100644
index 0000000..ca6f943
--- /dev/null
+++ b/src/tint/lang/msl/writer/raise/switch_return_test.cc
@@ -0,0 +1,152 @@
+// Copyright 2026 The Dawn & Tint Authors
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+//    list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+//    this list of conditions and the following disclaimer in the documentation
+//    and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+//    contributors may be used to endorse or promote products derived from
+//    this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "src/tint/lang/msl/writer/raise/switch_return.h"
+
+#include <utility>
+
+#include "src/tint/lang/core/fluent_types.h"
+#include "src/tint/lang/core/ir/transform/helper_test.h"
+
+using namespace tint::core::fluent_types;     // NOLINT
+using namespace tint::core::number_suffixes;  // NOLINT
+
+namespace tint::msl::writer::raise {
+namespace {
+
+using MslWriter_SwitchReturnTest = core::ir::transform::TransformTest;
+
+TEST_F(MslWriter_SwitchReturnTest, ReturnInsideSwitch) {
+    auto* func = b.Function("foo", ty.void_());
+    b.Append(func->Block(), [&] {
+        auto* sw = b.Switch(1_i);
+        b.Append(b.Case(sw, {b.Constant(0_i)}), [&] { b.Return(func); });
+        b.Append(b.DefaultCase(sw), [&] { b.ExitSwitch(sw); });
+        b.Return(func);
+    });
+
+    auto* src = R"(
+%foo = func():void {
+  $B1: {
+    switch 1i [c: (0i, $B2), c: (default, $B3)] {  # switch_1
+      $B2: {  # case
+        ret
+      }
+      $B3: {  # case
+        exit_switch  # switch_1
+      }
+    }
+    ret
+  }
+}
+)";
+    EXPECT_EQ(src, str());
+
+    auto* expect = R"(
+%foo = func():void {
+  $B1: {
+    switch 1i [c: (0i, $B2), c: (default, $B3)] {  # switch_1
+      $B2: {  # case
+        %2:u32 = msl.volatile_zero
+        %3:bool = eq %2, 0u
+        if %3 [t: $B4] {  # if_1
+          $B4: {  # true
+            ret
+          }
+        }
+        exit_switch  # switch_1
+      }
+      $B3: {  # case
+        exit_switch  # switch_1
+      }
+    }
+    ret
+  }
+}
+)";
+
+    Run(SwitchReturn);
+
+    EXPECT_EQ(expect, str());
+}
+
+TEST_F(MslWriter_SwitchReturnTest, ReturnValueInsideSwitch) {
+    auto* func = b.Function("foo", ty.i32());
+    b.Append(func->Block(), [&] {
+        auto* sw = b.Switch(1_i);
+        b.Append(b.Case(sw, {b.Constant(0_i)}), [&] { b.Return(func, 42_i); });
+        b.Append(b.DefaultCase(sw), [&] { b.ExitSwitch(sw); });
+        b.Return(func, 0_i);
+    });
+
+    auto* src = R"(
+%foo = func():i32 {
+  $B1: {
+    switch 1i [c: (0i, $B2), c: (default, $B3)] {  # switch_1
+      $B2: {  # case
+        ret 42i
+      }
+      $B3: {  # case
+        exit_switch  # switch_1
+      }
+    }
+    ret 0i
+  }
+}
+)";
+    EXPECT_EQ(src, str());
+
+    auto* expect = R"(
+%foo = func():i32 {
+  $B1: {
+    switch 1i [c: (0i, $B2), c: (default, $B3)] {  # switch_1
+      $B2: {  # case
+        %2:u32 = msl.volatile_zero
+        %3:bool = eq %2, 0u
+        if %3 [t: $B4] {  # if_1
+          $B4: {  # true
+            ret 42i
+          }
+        }
+        exit_switch  # switch_1
+      }
+      $B3: {  # case
+        exit_switch  # switch_1
+      }
+    }
+    ret 0i
+  }
+}
+)";
+
+    Run(SwitchReturn);
+
+    EXPECT_EQ(expect, str());
+}
+
+}  // namespace
+}  // namespace tint::msl::writer::raise