[Mlir-commits] [mlir] [mlir][SPIR-V] Fix deserializer crash on entry block selection header (PR #237137)

[Arseniy Obolenskiy]

https://github.com/aobolensk created https://github.com/llvm/llvm-project/pull/237137

Split the entry block when it also starts a selection, since it cannot move into the region

>From 36339dc2bb05306b306b14433301453de431b25a Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Fri, 30 Oct 2026 09:05:14 +0100
Subject: [PATCH] [mlir][SPIR-V] Fix deserializer crash on entry block
 selection header

Split the entry block when it also starts a selection, since it cannot move into the region
---
 .../SPIRV/Deserialization/Deserializer.cpp    |  5 +-
 .../Target/SPIRV/selection-entry-block.spvasm | 56 +++++++++++++++++++
 2 files changed, 59 insertions(+), 2 deletions(-)
 create mode 100644 mlir/test/Target/SPIRV/selection-entry-block.spvasm

diff --git a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
index 42e5620e308e71..545ece9b7fb0a8 100644
--- a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
+++ b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
@@ -3056,8 +3056,9 @@ LogicalResult spirv::Deserializer::splitSelectionHeader() {
     // Do not split a block that only contains a conditional branch / switch,
     // unless it is also a merge block of another construct - in that case we
     // want to split the block. We do not want two constructs to share header /
-    // merge block.
-    if (!llvm::hasSingleElement(*block) || splitHeaderMergeBlock) {
+    // merge block. The entry block cannot join a region, so split it too.
+    if (!llvm::hasSingleElement(*block) || splitHeaderMergeBlock ||
+        isFnEntryBlock(block)) {
       Block *newBlock = block->splitBlock(terminator);
       OpBuilder builder(block, block->end());
       spirv::BranchOp::create(builder, block->getParent()->getLoc(), newBlock);
diff --git a/mlir/test/Target/SPIRV/selection-entry-block.spvasm b/mlir/test/Target/SPIRV/selection-entry-block.spvasm
new file mode 100644
index 00000000000000..3c3c5fc68f69ee
--- /dev/null
+++ b/mlir/test/Target/SPIRV/selection-entry-block.spvasm
@@ -0,0 +1,56 @@
+; RUN: %if spirv-tools %{ spirv-as --target-env spv1.0 %s -o - | mlir-translate --deserialize-spirv - -o - | FileCheck %s %}
+
+; COM: Check that a selection header that is also the function entry block is
+; COM: split, so the entry block stays outside the selection region.
+
+; CHECK-LABEL: spirv.func @branch_conditional
+; CHECK-NEXT:    spirv.Branch ^[[HEADER:.+]]
+; CHECK-NEXT:  ^[[HEADER]]:
+; CHECK-NEXT:    spirv.mlir.selection {
+; CHECK-NEXT:      spirv.BranchConditional
+; CHECK:           spirv.mlir.merge
+; CHECK-NEXT:    }
+; CHECK-NEXT:    spirv.Return
+
+; CHECK-LABEL: spirv.func @switch
+; CHECK-NEXT:    spirv.Branch ^[[HEADER:.+]]
+; CHECK-NEXT:  ^[[HEADER]]:
+; CHECK-NEXT:    spirv.mlir.selection {
+; CHECK-NEXT:      spirv.Switch
+; CHECK:           spirv.mlir.merge
+; CHECK-NEXT:    }
+; CHECK-NEXT:    spirv.Return
+
+               OpCapability Shader
+               OpCapability Linkage
+               OpMemoryModel Logical GLSL450
+               OpName %branch_conditional "branch_conditional"
+               OpName %switch "switch"
+       %void = OpTypeVoid
+       %bool = OpTypeBool
+        %int = OpTypeInt 32 1
+      %int_1 = OpConstant %int 1
+     %fn_bool = OpTypeFunction %void %bool
+      %fn_int = OpTypeFunction %void %int
+%branch_conditional = OpFunction %void None %fn_bool
+       %cond = OpFunctionParameter %bool
+   %bc_entry = OpLabel
+               OpSelectionMerge %bc_merge None
+               OpBranchConditional %cond %bc_true %bc_merge
+    %bc_true = OpLabel
+     %bc_add = OpIAdd %int %int_1 %int_1
+               OpBranch %bc_merge
+   %bc_merge = OpLabel
+               OpReturn
+               OpFunctionEnd
+     %switch = OpFunction %void None %fn_int
+        %sel = OpFunctionParameter %int
+   %sw_entry = OpLabel
+               OpSelectionMerge %sw_merge None
+               OpSwitch %sel %sw_merge 1 %sw_case
+    %sw_case = OpLabel
+     %sw_add = OpIAdd %int %sel %sel
+               OpBranch %sw_merge
+   %sw_merge = OpLabel
+               OpReturn
+               OpFunctionEnd