[Mlir-commits] [mlir] [mlir][spirv] Split conditional basic blocks during deserialization (PR #127639)

[Jakub Kuderski]

================
@@ -2158,13 +2158,55 @@ LogicalResult spirv::Deserializer::wireUpBlockArgument() {
   return success();
 }
 
+LogicalResult spirv::Deserializer::splitConditionalBlocks() {
+  auto splitBlock = [&](Block *block) {
+    // Do not split loop headers
+    if (auto it = blockMergeInfo.find(block); it != blockMergeInfo.end()) {
+      if (it->second.continueBlock) {
+        return;
+      }
+    }
+
+    if (!block->mightHaveTerminator())
+      return;
+
+    auto terminator = block->getTerminator();
+    assert(terminator != nullptr);
+
+    if (isa<spirv::BranchConditionalOp>(terminator) &&
+        std::distance(block->begin(), block->end()) > 1) {
+      auto newBlock = block->splitBlock(terminator);
+      OpBuilder builder(block, block->end());
+      builder.create<spirv::BranchOp>(block->getParent()->getLoc(), newBlock);
+
+      if (auto it = blockMergeInfo.find(block); it != blockMergeInfo.end()) {
+        auto value = std::move(it->second);
----------------
kuhar wrote:

also here

https://github.com/llvm/llvm-project/pull/127639