[mlir][spirv] Create separate selection header block during serialization

[antiagainst]

Summary

The previous "optimization" that tries to reuse existing block for
selection header block can be problematic for deserialization
because it effectively pulls in previous ops in the selection op's
enclosing block into the selection op's header. When deserializing,
those ops will be placed in the selection op's region. If any of
the previous ops has usage after the section, it will break. That
is, the following IR cannot round trip:

mlir
^bb:
  %def = ...
  spv.mlir.selection { ... }
  %use = spv.SomeOp %def

This commit removes the "optimization" to always create new blocks
for the selection header.

Depends On D115560

[Hardcode84]

Was this issue about serializer generating completely invalid blob or valid blob we just can't deserialize (yet?)?
In any case, does it make sense to add sanity check on deserializer side to at least fail gracefully on such blob?

[antiagainst]

Add explicit check and improve error reporting

[antiagainst]

In D115582#3187391, @Hardcode84 wrote:

Was this issue about serializer generating completely invalid blob or valid blob we just can't deserialize (yet?)?

It's more of the the second case. The generated blob is valid for this case per se; although I suspect it can be problematic for other cases by just reusing the existing block for the header block (and it's bad to not round trip anyway) so I fixed the serialization here.

We cannot structurize the control flow for this case (which can also be generated by some other SPIR-V generator) during deserialization because right now we try to "sink" all blocks in the selection construct into a spv.mlir.selection op. If the header block contains a value that are used after the newly formed selection op, then it won't work. It's related to the fact that spv.mlir.selection and spv.mlir.loop do not support yield semantics so values in their region blocks cannot "escape" and have uses outside. (We could consider supporting yielding; but it's a much larger fight that I'd prefer to leave it later. We could also inspect the uses and try to not sink those ones with outside uses but that's nasty again, considering phi values, etc.)

Fundamentally deserialization is solving a harder problem than serialization, given that deserialization is sort of raising the abstraction levels by going from plain basic blocks into regions. Also consider the possible SPIR-V blobs in the wild, the surface area is also huge. So making deserialization robust would require quite some investment. Thus far we have been trying to at least round trip, which is a way to make sure it's not lagging behind too much there (in the hope of eventually we can import in SPIR-V blobs from other generators freely!). Serialization matters more if speaking of compilation from a high level input though. It is lowering so it's easier. It is already relatively robust right now.

In any case, does it make sense to add sanity check on deserializer side to at least fail gracefully on such blob?

Yup. I've added an explicit check for this case and turned some other asserts into proper errors.

[Hardcode84]

Thanks for the explanation

mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

1806

nit: !op.use_empty() ?

mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

LogicalResult ControlFlowStructurizer::structurizeImpl() {
   // SelectionOp/LoopOp resides right now.
   headerBlock->replaceAllUsesWith(mergeBlock);
 
   LLVM_DEBUG(llvm::dbgs() << "[cf] after cloning and fixing references:\n");
   LLVM_DEBUG(llvm::dbgs() << *headerBlock->getParentOp());
   LLVM_DEBUG(llvm::dbgs() << "\n");
 
   if (isLoop) {
-    // The selection/loop header block may have block arguments. Since now
-    // we place the selection/loop op inside the old merge block, we need to
-    // make sure the old merge block has the same block argument list.
-    assert(mergeBlock->args_empty() && "OpPhi in loop merge block unsupported");
+    if (!mergeBlock->args_empty()) {
+      return mergeBlock->getParentOp()->emitError(
+          "OpPhi in loop merge block unsupported");
+    }
+
+    // The loop header block may have block arguments. Since now we place the
+    // loop op inside the old merge block, we need to make sure the old merge
+    // block has the same block argument list.
     for (BlockArgument blockArg : headerBlock->getArguments()) {
       mergeBlock->addArgument(blockArg.getType());
     }
 
-    // If the loop header block has block arguments, make sure the spv.branch op
+    // If the loop header block has block arguments, make sure the spv.Branch op
     // matches.
     SmallVector<Value, 4> blockArgs;
     if (!headerBlock->args_empty())
       blockArgs = {mergeBlock->args_begin(), mergeBlock->args_end()};
 
     // The loop entry block should have a unconditional branch jumping to the
     // loop header block.
     builder.setInsertionPointToEnd(&body.front());
     builder.create<spirv::BranchOp>(location, mapper.lookupOrNull(headerBlock),
                                     ArrayRef<Value>(blockArgs));
   }
 
   // All the blocks cloned into the SelectionOp/LoopOp's region can now be
   // cleaned up.
   LLVM_DEBUG(llvm::dbgs() << "[cf] cleaning up blocks after clone\n");
   // First we need to drop all operands' references inside all blocks. This is
   // needed because we can have blocks referencing SSA values from one another.
   for (auto *block : constructBlocks)
     block->dropAllReferences();
 
+  // Check that whether some op in the to-be-erased blocks still has uses. Those
+  // uses come from blocks that won't be sinked into the SelectionOp/LoopOp's
+  // region. We cannot handle such cases given that once a value is sinked into
+  // the SelectionOp/LoopOp's region, there is no escape for it:
+  // SelectionOp/LooOp does not support yield values right now.
+  for (auto *block : constructBlocks) {
+    for (Operation &op : *block)
+      if (!op.use_empty())
+        return op.emitOpError(
+            "failed control flow structurization: it has uses outside of the "
+            "enclosing selection/loop construct");
+  }
+
   // Then erase all old blocks.
   for (auto *block : constructBlocks) {
     // We've cloned all blocks belonging to this construct into the structured
     // control flow op's region. Among these blocks, some may compose another
     // selection/loop. If so, they will be recorded within blockMergeInfo.
     // We need to update the pointers there to the newly remapped ones so we can
     // continue structurizing them later.
-    // TODO: The asserts in the following assumes input SPIR-V blob
-    // forms correctly nested selection/loop constructs. We should relax this
-    // and support error cases better.
+    // TODO: The asserts in the following assumes input SPIR-V blob forms
+    // correctly nested selection/loop constructs. We should relax this and
+    // support error cases better.
     auto it = blockMergeInfo.find(block);
     if (it != blockMergeInfo.end()) {
+      // Use the original location for nested selection/loop ops.
+      Location loc = it->second.loc;
+
       Block *newHeader = mapper.lookupOrNull(block);
-      assert(newHeader && "nested loop header block should be remapped!");
+      if (!newHeader)
+        return emitError(loc, "failed control flow structurization: nested "
+                              "loop header block should be remapped!");
 
       Block *newContinue = it->second.continueBlock;
       if (newContinue) {
         newContinue = mapper.lookupOrNull(newContinue);
-        assert(newContinue && "nested loop continue block should be remapped!");
+        if (!newContinue)
+          return emitError(loc, "failed control flow structurization: nested "
+                                "loop continue block should be remapped!");
       }
 
       Block *newMerge = it->second.mergeBlock;
       if (Block *mappedTo = mapper.lookupOrNull(newMerge))
         newMerge = mappedTo;
 
-      // Keep original location for nested selection/loop ops.
-      Location loc = it->second.loc;
       // The iterator should be erased before adding a new entry into
       // blockMergeInfo to avoid iterator invalidation.
       blockMergeInfo.erase(it);
       blockMergeInfo.try_emplace(newHeader, loc, it->second.control, newMerge,
                                  newContinue);
     }
 
     // The structured selection/loop's entry block does not have arguments.

mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp

LogicalResult Serializer::processSelectionOp(spirv::SelectionOp selectionOp) {
   // Assign <id>s to all blocks so that branches inside the SelectionOp can
   // resolve properly.
   auto &body = selectionOp.body();
   for (Block &block : body)
     getOrCreateBlockID(&block);
 
   auto *headerBlock = selectionOp.getHeaderBlock();
   auto *mergeBlock = selectionOp.getMergeBlock();
+  auto headerID = getBlockID(headerBlock);
   auto mergeID = getBlockID(mergeBlock);
   auto loc = selectionOp.getLoc();
 
+  // This SelectionOp is in some MLIR block with preceding and following ops. In
+  // the binary format, it should reside in separate SPIR-V blocks from its
+  // preceding and following ops. So we need to emit unconditional branches to
+  // jump to this SelectionOp's SPIR-V blocks and jumping back to the normal
+  // flow afterwards.
+  encodeInstructionInto(functionBody, spirv::Opcode::OpBranch, {headerID});
+
   // Emit the selection header block, which dominates all other blocks, first.
   // We need to emit an OpSelectionMerge instruction before the selection header
   // block's terminator.
   auto emitSelectionMerge = [&]() {
     if (failed(emitDebugLine(functionBody, loc)))
       return failure();
     lastProcessedWasMergeInst = true;
     encodeInstructionInto(
         functionBody, spirv::Opcode::OpSelectionMerge,
         {mergeID, static_cast<uint32_t>(selectionOp.selection_control())});
     return success();
   };
-  // For structured selection, we cannot have blocks in the selection construct
-  // branching to the selection header block. Entering the selection (and
-  // reaching the selection header) must be from the block containing the
-  // spv.mlir.selection op. If there are ops ahead of the spv.mlir.selection op
-  // in the block, we can "merge" them into the selection header. So here we
-  // don't need to emit a separate block; just continue with the existing block.
-  if (failed(processBlock(headerBlock, /*omitLabel=*/true, emitSelectionMerge)))
+  if (failed(
+          processBlock(headerBlock, /*omitLabel=*/false, emitSelectionMerge)))
     return failure();
 
   // Process all blocks with a depth-first visitor starting from the header
   // block. The selection header block and merge block are skipped by this
   // visitor.
   if (failed(visitInPrettyBlockOrder(
           headerBlock, [&](Block *block) { return processBlock(block); },
           /*skipHeader=*/true, /*skipBlocks=*/{mergeBlock})))

mlir/test/Target/SPIRV/loop.mlir

 
 // Loop with selection in its header
 
 spv.module Physical64 OpenCL requires #spv.vce<v1.0, [Kernel, Linkage, Addresses, Int64], []> {
 // CHECK-LABEL:   @kernel
 // CHECK-SAME:    (%[[INPUT0:.+]]: i64)
   spv.func @kernel(%input: i64) "None" {
 // CHECK-NEXT:     %[[VAR:.+]] = spv.Variable : !spv.ptr<i1, Function>
-// CHECK-NEXT:     spv.Branch ^[[BB:.+]](%[[INPUT0]] : i64)
-// CHECK-NEXT:   ^[[BB]](%[[INPUT1:.+]]: i64):
+// CHECK-NEXT:     spv.Branch ^[[BB0:.+]](%[[INPUT0]] : i64)
+// CHECK-NEXT:   ^[[BB0]](%[[INPUT1:.+]]: i64):
     %cst0_i64 = spv.Constant 0 : i64
     %true = spv.Constant true
     %false = spv.Constant false
 // CHECK-NEXT:     spv.mlir.loop {
     spv.mlir.loop {
 // CHECK-NEXT:       spv.Branch ^[[LOOP_HEADER:.+]](%[[INPUT1]] : i64)
       spv.Branch ^loop_header(%input : i64)
 // CHECK-NEXT:     ^[[LOOP_HEADER]](%[[ARG1:.+]]: i64):
     ^loop_header(%arg1: i64):
 // CHECK-NEXT:       spv.Branch ^[[LOOP_BODY:.+]]
 // CHECK-NEXT:     ^[[LOOP_BODY]]:
+// CHECK-NEXT:         %[[C0:.+]] = spv.Constant 0 : i64
       %gt = spv.SGreaterThan %arg1, %cst0_i64 : i64
+// CHECK-NEXT:         %[[GT:.+]] = spv.SGreaterThan %[[ARG1]], %[[C0]] : i64
+// CHECK-NEXT:         spv.Branch ^[[BB1:.+]]
+// CHECK-NEXT:     ^[[BB1]]:
       %var = spv.Variable : !spv.ptr<i1, Function>
 // CHECK-NEXT:       spv.mlir.selection {
       spv.mlir.selection {
-// CHECK-NEXT:         %[[C0:.+]] = spv.Constant 0 : i64
-// CHECK-NEXT:         %[[GT:.+]] = spv.SGreaterThan %[[ARG1]], %[[C0]] : i64
 // CHECK-NEXT:         spv.BranchConditional %[[GT]], ^[[THEN:.+]], ^[[ELSE:.+]]
         spv.BranchConditional %gt, ^then, ^else
 // CHECK-NEXT:       ^[[THEN]]:
       ^then:
 // CHECK-NEXT:         %true = spv.Constant true
 // CHECK-NEXT:         spv.Store "Function" %[[VAR]], %true : i1
         spv.Store "Function" %var, %true : i1
 // CHECK-NEXT:         spv.Branch ^[[SELECTION_MERGE:.+]]

mlir/test/Target/SPIRV/selection.mlir

 // RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 // Selection with both then and else branches
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+// CHECK-LABEL: @selection
   spv.func @selection(%cond: i1) -> () "None" {
-// CHECK:        spv.Branch ^bb1
-// CHECK-NEXT: ^bb1:
+// CHECK-NEXT:   spv.Constant 0
+// CHECK-NEXT:   spv.Variable
+// CHECK:        spv.Branch ^[[BB:.+]]
+// CHECK-NEXT: ^[[BB]]:
     %zero = spv.Constant 0: i32
     %one = spv.Constant 1: i32
     %two = spv.Constant 2: i32
     %var = spv.Variable init(%zero) : !spv.ptr<i32, Function>
 
 // CHECK-NEXT:   spv.mlir.selection control(Flatten)
-// CHECK-NEXT:     spv.Constant 0
-// CHECK-NEXT:     spv.Variable
     spv.mlir.selection control(Flatten) {
-// CHECK-NEXT: spv.BranchConditional %{{.*}} [5, 10], ^bb1, ^bb2
+// CHECK-NEXT: spv.BranchConditional %{{.*}} [5, 10], ^[[THEN:.+]], ^[[ELSE:.+]]
       spv.BranchConditional %cond [5, 10], ^then, ^else
 
-// CHECK-NEXT:   ^bb1:
+// CHECK-NEXT:   ^[[THEN]]:
     ^then:
 // CHECK-NEXT:     spv.Constant 1
 // CHECK-NEXT:     spv.Store
       spv.Store "Function" %var, %one : i32
-// CHECK-NEXT:     spv.Branch ^bb3
+// CHECK-NEXT:     spv.Branch ^[[MERGE:.+]]
       spv.Branch ^merge
 
-// CHECK-NEXT:   ^bb2:
+// CHECK-NEXT:   ^[[ELSE]]:
     ^else:
 // CHECK-NEXT:     spv.Constant 2
 // CHECK-NEXT:     spv.Store
       spv.Store "Function" %var, %two : i32
-// CHECK-NEXT:     spv.Branch ^bb3
+// CHECK-NEXT:     spv.Branch ^[[MERGE]]
       spv.Branch ^merge
 
-// CHECK-NEXT:   ^bb3:
+// CHECK-NEXT:   ^[[MERGE]]:
     ^merge:
 // CHECK-NEXT:     spv.mlir.merge
       spv.mlir.merge
     }
 
     spv.Return
   }
 
   spv.func @main() -> () "None" {
     spv.Return
   }
   spv.EntryPoint "GLCompute" @main
   spv.ExecutionMode @main "LocalSize", 1, 1, 1
 }
 
 // -----
 
 // Selection with only then branch
 // Selection in function entry block
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
-// CHECK:      spv.func @selection(%[[ARG:.*]]: i1
+// CHECK-LABEL: spv.func @selection
+//  CHECK-SAME: (%[[ARG:.*]]: i1)
   spv.func @selection(%cond: i1) -> (i32) "None" {
-// CHECK:        spv.Branch ^bb1
-// CHECK-NEXT: ^bb1:
+// CHECK:        spv.Branch ^[[BB:.+]]
+// CHECK-NEXT: ^[[BB]]:
 // CHECK-NEXT:   spv.mlir.selection
     spv.mlir.selection {
-// CHECK-NEXT: spv.BranchConditional %[[ARG]], ^bb1, ^bb2
+// CHECK-NEXT: spv.BranchConditional %[[ARG]], ^[[THEN:.+]], ^[[ELSE:.+]]
       spv.BranchConditional %cond, ^then, ^merge
 
-// CHECK:        ^bb1:
+// CHECK:        ^[[THEN]]:
     ^then:
       %zero = spv.Constant 0 : i32
       spv.ReturnValue  %zero : i32
 
-// CHECK:        ^bb2:
+// CHECK:        ^[[ELSE]]:
     ^merge:
 // CHECK-NEXT:     spv.mlir.merge
       spv.mlir.merge
     }
 
     %one = spv.Constant 1 : i32
     spv.ReturnValue  %one : i32
   }
 
   spv.func @main() -> () "None" {
     spv.Return
   }
   spv.EntryPoint "GLCompute" @main
   spv.ExecutionMode @main "LocalSize", 1, 1, 1
 }
 
+// -----
+
+// Selection with control flow afterwards
+// SSA value def before selection and use after selection
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+// CHECK-LABEL: @selection_cf()
+  spv.func @selection_cf() -> () "None" {
+    %true = spv.Constant true
+    %false = spv.Constant false
+    %zero = spv.Constant 0 : i32
+    %one = spv.Constant 1 : i32
+// CHECK-NEXT:    %[[VAR:.+]] = spv.Variable
+    %var = spv.Variable : !spv.ptr<i1, Function>
+// CHECK-NEXT:    spv.Branch ^[[BB:.+]]
+// CHECK-NEXT:  ^[[BB]]:
+
+// CHECK-NEXT:    spv.mlir.selection {
+    spv.mlir.selection {
+//      CHECK:      spv.BranchConditional %{{.+}}, ^[[THEN0:.+]], ^[[ELSE0:.+]]
+      spv.BranchConditional %true, ^then0, ^else0
+
+// CHECK-NEXT:    ^[[THEN0]]:
+//      CHECK:      spv.Store "Function" %[[VAR]]
+// CHECK-NEXT:      spv.Branch ^[[MERGE:.+]]
+    ^then0:
+      spv.Store "Function" %var, %true : i1
+      spv.Branch ^merge
+
+// CHECK-NEXT:    ^[[ELSE0]]:
+//      CHECK:      spv.Store "Function" %[[VAR]]
+// CHECK-NEXT:      spv.Branch ^[[MERGE]]
+    ^else0:
+      spv.Store "Function" %var, %false : i1
+      spv.Branch ^merge
+
+// CHECK-NEXT:    ^[[MERGE]]:
+// CHECK-NEXT:      spv.mlir.merge
+    ^merge:
+      spv.mlir.merge
+// CHECK-NEXT:    }
+    }
+
+// CHECK-NEXT:    spv.Load "Function" %[[VAR]]
+    %cond = spv.Load "Function" %var : i1
+//      CHECK:    spv.BranchConditional %1, ^[[THEN1:.+]](%{{.+}} : i32), ^[[ELSE1:.+]](%{{.+}}, %{{.+}} : i32, i32)
+    spv.BranchConditional %cond, ^then1(%one: i32), ^else1(%zero, %zero: i32, i32)
+
+// CHECK-NEXT:  ^[[THEN1]](%{{.+}}: i32):
+// CHECK-NEXT:    spv.Return
+  ^then1(%arg0: i32):
+    spv.Return
+
+// CHECK-NEXT:  ^[[ELSE1]](%{{.+}}: i32, %{{.+}}: i32):
+// CHECK-NEXT:    spv.Return
+  ^else1(%arg1: i32, %arg2: i32):
+    spv.Return
+  }
+}