[mlir][spirv] Fix wrong Phi parent block for back-to-back loops

[antiagainst]

Summary

If we have two back-to-back loops with block arguments, the OpPhi
instructions generated for the second loop's block arguments should
have use the merge block of the first SPIR-V loop structure as
their incoming parent block.

[mravishankar]

mlir/lib/Dialect/SPIRV/Serialization/Serializer.cpp

99

Just to clarify, this loop executes at most 1 iteration? If so can we avoid having a while loop? That makes it more confusing IMO

[antiagainst]

mlir/lib/Dialect/SPIRV/Serialization/Serializer.cpp

99

Not necessary. This actually scans all previous ops before this loop op. For example, you can have the following:

spv.loop { ... } // loop1
spv.Load
spv.Bla
spv.Store
spv.loop { ... } // loop2

Here if loop2 has some block arguments that require phi nodes, the incoming parent block should be loop1's merge block. loop1 is multiple ops away so we need while loop here.

mlir/lib/Dialect/SPIRV/Serialization/Serializer.cpp

for (Block *block : llvm::depth_first_ext(headerBlock, doneBlocks)) {
     if (skipHeader && block == headerBlock)
       continue;
     if (failed(blockHandler(block)))
       return failure();
   }
   return success();
 }
 
-/// Returns the last structured control flow op's merge block if the given
-/// `block` contains any structured control flow op. Otherwise returns nullptr.
-static Block *getLastStructuredControlFlowOpMergeBlock(Block *block) {
-  for (Operation &op : llvm::reverse(block->getOperations())) {
-    if (auto selectionOp = dyn_cast<spirv::SelectionOp>(op))
-      return selectionOp.getMergeBlock();
-    if (auto loopOp = dyn_cast<spirv::LoopOp>(op))
-      return loopOp.getMergeBlock();
-  }
+/// Returns the merge block if the given `op` is a structured control flow op.
+/// Otherwise returns nullptr.
+static Block *getStructuredControlFlowOpMergeBlock(Operation *op) {
   return nullptr;
 }
 
+/// Given a predecessor `block` for a block with arguments, returns the block
+/// that should be used as the parent block for SPIR-V OpPhi instructions
+/// corresponding to the block arguments.
+static Block *getPhiIncomingBlock(Block *block) {
+  // If the predecessor block in question is the entry block for a spv.loop,
+  // we jump to this spv.loop from its enclosing block.
+  if (block->isEntryBlock()) {
+    if (auto loopOp = dyn_cast<spirv::LoopOp>(block->getParentOp())) {
+      // Then the incoming parent block for OpPhi should be the merge block of
+      // the structured control flow op before this loop.
+      Operation *op = loopOp.getOperation();
+      while ((op = op->getPrevNode()) != nullptr)
+        if (Block *incomingBlock = getStructuredControlFlowOpMergeBlock(op))
+          return incomingBlock;
+      // Or the enclosing block itself if no structured control flow ops
+      // exists before this loop.
+      return loopOp.getOperation()->getBlock();
+    }
+  }
+
+  // Otherwise, we jump from the given predecessor block. Try to see if there is
+  // a structured control flow op inside it.
+  for (Operation &op : llvm::reverse(block->getOperations())) {
+    if (Block *incomingBlock = getStructuredControlFlowOpMergeBlock(&op))
+      return incomingBlock;
+  }
+  return block;
+}
+
 namespace {
 
 /// A SPIR-V module serializer.
 ///
 /// A SPIR-V binary module is a single linear stream of instructions; each
 /// instruction is composed of 32-bit words with the layout:
 ///
 ///   | <word-count>|<opcode> |  <operand>   |  <operand>   | ... |
LogicalResult Serializer::emitPhiForBlockArguments(Block *block) {
   // If the block has arguments, we need to create SPIR-V OpPhi instructions.
   // A SPIR-V OpPhi instruction is of the syntax:
   //   OpPhi | result type | result <id> | (value <id>, parent block <id>) pair
   // So we need to collect all predecessor blocks and the arguments they send
   // to this block.
   SmallVector<std::pair<Block *, Operation::operand_iterator>, 4> predecessors;
   for (Block *predecessor : block->getPredecessors()) {
     auto *terminator = predecessor->getTerminator();
-    // Check whether this predecessor block contains a structured control flow
-    // op. If so, the structured control flow op will be serialized to multiple
-    // SPIR-V blocks. The branch op jumping to the OpPhi's block then resides in
-    // the last structured control flow op's merge block.
-    if (auto *merge = getLastStructuredControlFlowOpMergeBlock(predecessor))
-      predecessor = merge;
+    // The predecessor here is the immediate one according to MLIR's IR
+    // structure. It does not directly map to the incoming parent block for the
+    // OpPhi instructions at SPIR-V binary level. This is because structured
+    // control flow ops are serialized to multiple SPIR-V blocks. If there is a
+    // spv.selection/spv.loop op in the MLIR predecessor block, the branch op
+    // jumping to the OpPhi's block then resides in the previous structured
+    // control flow op's merge block.
+    predecessor = getPhiIncomingBlock(predecessor);
     if (auto branchOp = dyn_cast<spirv::BranchOp>(terminator)) {
       predecessors.emplace_back(predecessor, branchOp.operand_begin());
     } else {
       return terminator->emitError("unimplemented terminator for Phi creation");
     }
   }
 
   // Then create OpPhi instruction for each of the block argument.
   for (auto argIndex : llvm::seq<unsigned>(0, block->getNumArguments())) {
     BlockArgument arg = block->getArgument(argIndex);
 
     // Get the type <id> and result <id> for this OpPhi instruction.
     uint32_t phiTypeID = 0;
     if (failed(processType(arg.getLoc(), arg.getType(), phiTypeID)))
       return failure();
     uint32_t phiID = getNextID();
 
     LLVM_DEBUG(llvm::dbgs() << "[phi] for block argument #" << argIndex << ' '
                             << arg << " (id = " << phiID << ")\n");
 
+    // Prepare the (value <id>, parent block <id>) pairs.
     SmallVector<uint32_t, 8> phiArgs;
     phiArgs.push_back(phiTypeID);
     phiArgs.push_back(phiID);
 
     for (auto predIndex : llvm::seq<unsigned>(0, predecessors.size())) {
       Value value = *(predecessors[predIndex].second + argIndex);
       uint32_t predBlockId = getOrCreateBlockID(predecessors[predIndex].first);
       LLVM_DEBUG(llvm::dbgs() << "[phi] use predecessor (id = " << predBlockId
LogicalResult Serializer::processLoopOp(spirv::LoopOp loopOp) {
 
   // This LoopOp 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 LoopOp's SPIR-V blocks and jumping back to the normal flow
   // afterwards.
   encodeInstructionInto(functionBody, spirv::Opcode::OpBranch, {headerID});
 
-  // We omit the LoopOp's entry block and start serialization from the loop
-  // header block. The entry block should not contain any additional ops other
-  // than a single spv.Branch that jumps to the loop header block. However,
-  // the spv.Branch can contain additional block arguments. Those block
-  // arguments must come from out of the loop using implicit capture. We will
-  // need to query the <id> for the value sent and the <id> for the incoming
-  // parent block. For the latter, we need to make sure this block is
-  // registered. The value sent should come from the block this loop resides in.
-  blockIDMap[loopOp.getEntryBlock()] =
-      getBlockID(loopOp.getOperation()->getBlock());
+  // LoopOp's entry block is just there for satisfying MLIR's structural
+  // requirements so we omit it and start serialization from the loop header
+  // block.
 
   // Emit the loop header block, which dominates all other blocks, first. We
   // need to emit an OpLoopMerge instruction before the loop header block's
   // terminator.
   auto emitLoopMerge = [&]() {
     // TODO(antiagainst): properly support loop control here
     encodeInstructionInto(
         functionBody, spirv::Opcode::OpLoopMerge,

mlir/test/Dialect/SPIRV/Serialization/phi.mlir

// CHECK:     spv._merge
       spv._merge
     }
     spv.Return
   }
 
   spv.EntryPoint "GLCompute" @fmul_kernel, @__builtin_var_WorkgroupId__, @__builtin_var_NumWorkgroups__
   spv.ExecutionMode @fmul_kernel "LocalSize", 32, 1, 1
 }
+
+// -----
+
+// Test back-to-back loops with block arguments
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+  spv.func @fmul_kernel() "None" {
+    %cst4 = spv.constant 4 : i32
+
+    %val1 = spv.constant 43 : i32
+    %val2 = spv.constant 44 : i32
+
+// CHECK:        spv.constant 43
+// CHECK-NEXT:   spv.Branch ^[[BB1:.+]](%{{.+}} : i32)
+// CHECK-NEXT: ^[[BB1]](%{{.+}}: i32):
+// CHECK-NEXT:   spv.loop
+    spv.loop { // loop 1
+      spv.Branch ^bb1(%val1 : i32)
+    ^bb1(%loop1_bb_arg: i32):
+      %loop1_lt = spv.SLessThan %loop1_bb_arg, %cst4 : i32
+      spv.BranchConditional %loop1_lt, ^bb2, ^bb3
+    ^bb2:
+      %loop1_add = spv.IAdd %loop1_bb_arg, %cst4 : i32
+      spv.Branch ^bb1(%loop1_add : i32)
+    ^bb3:
+      spv._merge
+    }
+
+// CHECK:        spv.constant 44
+// CHECK-NEXT:   spv.Branch ^[[BB2:.+]](%{{.+}} : i32)
+// CHECK-NEXT: ^[[BB2]](%{{.+}}: i32):
+// CHECK-NEXT:   spv.loop
+    spv.loop { // loop 2
+      spv.Branch ^bb1(%val2 : i32)
+    ^bb1(%loop2_bb_arg: i32):
+      %loop2_lt = spv.SLessThan %loop2_bb_arg, %cst4 : i32
+      spv.BranchConditional %loop2_lt, ^bb2, ^bb3
+    ^bb2:
+      %loop2_add = spv.IAdd %loop2_bb_arg, %cst4 : i32
+      spv.Branch ^bb1(%loop2_add : i32)
+    ^bb3:
+      spv._merge
+    }
+
+    spv.Return
+  }
+
+  spv.EntryPoint "GLCompute" @fmul_kernel
+  spv.ExecutionMode @fmul_kernel "LocalSize", 32, 1, 1
+}