[MLIR][SPIRV] Properly (de-)serialize BranchConditionalOp.

[ergawy]

Summary

Implements proper (de-)serialization logic for BranchConditionalOp when
such ops have true/false target operands.

[antiagainst]

Awesome, thanks for fixing this non-trivial part! :D

mlir/test/Target/SPIRV/phi.mlir

292

Add CHECK-LABELs for these tests?

[ergawy]

Handle review comments.

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

private:
   DenseMap<uint32_t, spirv::FuncOp> funcMap;
 
   // Result <id> to block mapping.
   DenseMap<uint32_t, Block *> blockMap;
 
   // Header block to its merge (and continue) target mapping.
   BlockMergeInfoMap blockMergeInfo;
 
-  // Block to its phi (block argument) mapping.
-  DenseMap<Block *, BlockPhiInfo> blockPhiInfo;
+  // For each pair of {predecessor, target} blocks, maps the pair of blocks to
+  // the list of phi arguments passed from predecessor to target.
+  DenseMap<std::pair<Block * /*predecessor*/, Block * /*target*/>, BlockPhiInfo>
+      blockPhiInfo;
 
   // Result <id> to value mapping.
   DenseMap<uint32_t, Value> valueMap;
 
   // Mapping from result <id> to undef value of a type.
   DenseMap<uint32_t, Type> undefMap;
 
   // Result <id> to name mapping.

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

LLVM_DEBUG(llvm::dbgs() << "[phi] created block argument " << blockArg
                           << " id = " << operands[1] << " of type "
                           << blockArgType << '\n');
 
   // For each (value, predecessor) pair, insert the value to the predecessor's
   // blockPhiInfo entry so later we can fix the block argument there.
   for (unsigned i = 2, e = operands.size(); i < e; i += 2) {
     uint32_t value = operands[i];
     Block *predecessor = getOrCreateBlock(operands[i + 1]);
-    blockPhiInfo[predecessor].push_back(value);
+    std::pair<Block *, Block *> predecessorTargetPair{predecessor, curBlock};
+    blockPhiInfo[predecessorTargetPair].push_back(value);
     LLVM_DEBUG(llvm::dbgs() << "[phi] predecessor @ " << predecessor
                             << " with arg id = " << value << '\n');
   }
 
   return success();
 }
 
 namespace {

 }
 
 LogicalResult spirv::Deserializer::wireUpBlockArgument() {
   LLVM_DEBUG(llvm::dbgs() << "[phi] start wiring up block arguments\n");
 
   OpBuilder::InsertionGuard guard(opBuilder);
 
   for (const auto &info : blockPhiInfo) {
-    Block *block = info.first;
+    Block *block = info.first.first;
+    Block *target = info.first.second;
     const BlockPhiInfo &phiInfo = info.second;
     LLVM_DEBUG(llvm::dbgs() << "[phi] block " << block << "\n");
     LLVM_DEBUG(llvm::dbgs() << "[phi] before creating block argument:\n");
     LLVM_DEBUG(block->getParentOp()->print(llvm::dbgs()));
     LLVM_DEBUG(llvm::dbgs() << '\n');
 
     // Set insertion point to before this block's terminator early because we
     // may materialize ops via getValue() call.
for (uint32_t valueId : phiInfo) {
       }
     }
 
     if (auto branchOp = dyn_cast<spirv::BranchOp>(op)) {
       // Replace the previous branch op with a new one with block arguments.
       opBuilder.create<spirv::BranchOp>(branchOp.getLoc(), branchOp.getTarget(),
                                         blockArgs);
       branchOp.erase();
+    } else if (auto branchCondOp = dyn_cast<spirv::BranchConditionalOp>(op)) {
+      assert((branchCondOp.getTrueBlock() == target ||
+              branchCondOp.getFalseBlock() == target) &&
+             "expected target to be either the true or false target");
+      if (target == branchCondOp.trueTarget())
+        opBuilder.create<spirv::BranchConditionalOp>(
+            branchCondOp.getLoc(), branchCondOp.condition(), blockArgs,
+            branchCondOp.getFalseBlockArguments(),
+            branchCondOp.branch_weightsAttr(), branchCondOp.trueTarget(),
+            branchCondOp.falseTarget());
+      else
+        opBuilder.create<spirv::BranchConditionalOp>(
+            branchCondOp.getLoc(), branchCondOp.condition(),
+            branchCondOp.getTrueBlockArguments(), blockArgs,
+            branchCondOp.branch_weightsAttr(), branchCondOp.getTrueBlock(),
+            branchCondOp.getFalseBlock());
+
+      branchCondOp.erase();
     } else {
       return emitError(unknownLoc, "unimplemented terminator for Phi creation");
     }
 
     LLVM_DEBUG(llvm::dbgs() << "[phi] after creating block argument:\n");
     LLVM_DEBUG(block->getParentOp()->print(llvm::dbgs()));
     LLVM_DEBUG(llvm::dbgs() << '\n');
   }

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

LogicalResult Serializer::emitPhiForBlockArguments(Block *block) {
   if (block->args_empty() || block->isEntryBlock())
     return success();
 
   // 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;
+  SmallVector<std::pair<Block *, OperandRange>, 4> predecessors;
   for (Block *predecessor : block->getPredecessors()) {
     auto *terminator = predecessor->getTerminator();
     // 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.mlir.selection/spv.mlir.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());
+      predecessors.emplace_back(predecessor, branchOp.getOperands());
+    } else if (auto branchCondOp =
+                   dyn_cast<spirv::BranchConditionalOp>(terminator)) {
+      Optional<OperandRange> blockOperands;
+
+      for (auto successorIdx :
+           llvm::seq<unsigned>(0, predecessor->getNumSuccessors()))
+        if (predecessor->getSuccessors()[successorIdx] == block) {
+          blockOperands = branchCondOp.getSuccessorOperands(successorIdx);
+          break;
+        }
+
+      assert(blockOperands && !blockOperands->empty() &&
+             "expected non-empty block operand range");
+      predecessors.emplace_back(predecessor, *blockOperands);
     } 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);
+      Value value = predecessors[predIndex].second[argIndex];
       uint32_t predBlockId = getOrCreateBlockID(predecessors[predIndex].first);
       LLVM_DEBUG(llvm::dbgs() << "[phi] use predecessor (id = " << predBlockId
                               << ") value " << value << ' ');
       // Each pair is a value <id> ...
       uint32_t valueId = getValueID(value);
       if (valueId == 0) {
         // The op generating this value hasn't been visited yet so we don't have
         // an <id> assigned yet. Record this to fix up later.

mlir/test/Target/SPIRV/phi.mlir

// CHECK-NEXT:   spv.mlir.loop
     }
 
     spv.Return
   }
 
   spv.EntryPoint "GLCompute" @fmul_kernel
   spv.ExecutionMode @fmul_kernel "LocalSize", 32, 1, 1
 }
+
+// -----
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+// CHECK-LABEL: @cond_branch_true_argument
+  spv.func @cond_branch_true_argument() -> () "None" {
+    %true = spv.Constant true
+    %zero = spv.Constant 0 : i32
+    %one = spv.Constant 1 : i32
+// CHECK:   spv.BranchConditional %{{.*}}, ^[[true1:.*]](%{{.*}}, %{{.*}} : i32, i32), ^[[false1:.*]]
+    spv.BranchConditional %true, ^true1(%zero, %zero: i32, i32), ^false1
+// CHECK: [[true1]](%{{.*}}: i32, %{{.*}}: i32)
+  ^true1(%arg0: i32, %arg1: i32):
+    spv.Return
+// CHECK: [[false1]]:
+  ^false1:
+    spv.Return
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+// CHECK-LABEL: @cond_branch_false_argument
+  spv.func @cond_branch_false_argument() -> () "None" {
+    %true = spv.Constant true
+    %zero = spv.Constant 0 : i32
+    %one = spv.Constant 1 : i32
+// CHECK:   spv.BranchConditional %{{.*}}, ^[[true1:.*]], ^[[false1:.*]](%{{.*}}, %{{.*}} : i32, i32)
+    spv.BranchConditional %true, ^true1, ^false1(%zero, %zero: i32, i32)
+// CHECK: [[true1]]:
+  ^true1:
+    spv.Return
+// CHECK: [[false1]](%{{.*}}: i32, %{{.*}}: i32):
+  ^false1(%arg0: i32, %arg1: i32):
+    spv.Return
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+// CHECK-LABEL: @cond_branch_true_and_false_argument
+  spv.func @cond_branch_true_and_false_argument() -> () "None" {
+    %true = spv.Constant true
+    %zero = spv.Constant 0 : i32
+    %one = spv.Constant 1 : i32
+// CHECK:   spv.BranchConditional %{{.*}}, ^[[true1:.*]](%{{.*}} : i32), ^[[false1:.*]](%{{.*}}, %{{.*}} : i32, i32)
+    spv.BranchConditional %true, ^true1(%one: i32), ^false1(%zero, %zero: i32, i32)
+// CHECK: [[true1]](%{{.*}}: i32):
+  ^true1(%arg0: i32):
+    spv.Return
+// CHECK: [[false1]](%{{.*}}: i32, %{{.*}}: i32):
+  ^false1(%arg1: i32, %arg2: i32):
+    spv.Return
+  }
+}