[mlir][spirv] Handle debuginfo for control flow ops.

[denis13]

Summary

Handle debuginfo for control flow operations: spv.Selection,
spv.Loop, spv.BranchOp, spv.BranchConditional.

[hanchung]

Still working on it, will reply once it's ready to review.

[hanchung]

In D79931#2037189, @hanchung wrote:

Still working on it, will reply once it's ready to review.

Oh, I left the comment on a wrong patch, please ignore this. Sorry about that.

[antiagainst]

Awesome, thanks Denis!

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

395

Super nit: can this be placed after the emitDebugInfo field? It's logically a field necessitated by emitDebugInfo so appearing after it would be nice. :)

[denis13]

Addressed review comments.

[denis13]

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

395

Sure, thanks for review!

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

 /// A struct for containing a header block's merge and continue targets.
 ///
 /// This struct is used to track original structured control flow info from
 /// SPIR-V blob. This info will be used to create spv.selection/spv.loop
 /// later.
 struct BlockMergeInfo {
   Block *mergeBlock;
   Block *continueBlock; // nullptr for spv.selection
+  Location loc;
 
-  BlockMergeInfo() : mergeBlock(nullptr), continueBlock(nullptr) {}
-  BlockMergeInfo(Block *m, Block *c = nullptr)
-      : mergeBlock(m), continueBlock(c) {}
+  BlockMergeInfo(Location location)
+      : mergeBlock(nullptr), continueBlock(nullptr), loc(location) {}
+  BlockMergeInfo(Location location, Block *m, Block *c = nullptr)
+      : mergeBlock(m), continueBlock(c), loc(location) {}
 };
 
 /// A struct for containing OpLine instruction information.
 struct DebugLine {
   uint32_t fileID;
   uint32_t line;
   uint32_t col;
 
if (!curBlock) {
     return emitError(unknownLoc, "OpBranch must appear inside a block");
   }
 
   if (operands.size() != 1) {
     return emitError(unknownLoc, "OpBranch must take exactly one target label");
   }
 
   auto *target = getOrCreateBlock(operands[0]);
-  opBuilder.create<spirv::BranchOp>(unknownLoc, target);
+  auto loc = createFileLineColLoc(opBuilder);
+  // The preceding instruction for the OpBranch instruction could be an
+  // OpLoopMerge or an OpSelectionMerge instruction, in this case they will have
+  // the same OpLine information.
+  opBuilder.create<spirv::BranchOp>(loc, target);
 
   clearDebugLine();
   return success();
 }
 
 LogicalResult
 Deserializer::processBranchConditional(ArrayRef<uint32_t> operands) {
   if (!curBlock) {
Deserializer::processBranchConditional(ArrayRef<uint32_t> operands) {
   auto condition = getValue(operands[0]);
   auto *trueBlock = getOrCreateBlock(operands[1]);
   auto *falseBlock = getOrCreateBlock(operands[2]);
 
   Optional<std::pair<uint32_t, uint32_t>> weights;
   if (operands.size() == 5) {
     weights = std::make_pair(operands[3], operands[4]);
   }
-
+  // The preceding instruction for the OpBranchConditional instruction could be
+  // an OpSelectionMerge instruction, in this case they will have the same
+  // OpLine information.
+  auto loc = createFileLineColLoc(opBuilder);
   opBuilder.create<spirv::BranchConditionalOp>(
-      unknownLoc, condition, trueBlock,
+      loc, condition, trueBlock,
       /*trueArguments=*/ArrayRef<Value>(), falseBlock,
       /*falseArguments=*/ArrayRef<Value>(), weights);
 
   clearDebugLine();
   return success();
 }
 
 LogicalResult Deserializer::processLabel(ArrayRef<uint32_t> operands) {
LogicalResult Deserializer::processSelectionMerge(ArrayRef<uint32_t> operands) {
 
   if (static_cast<uint32_t>(spirv::SelectionControl::None) != operands[1]) {
     return emitError(unknownLoc,
                      "unimplmented OpSelectionMerge selection control: ")
            << operands[2];
   }
 
   auto *mergeBlock = getOrCreateBlock(operands[0]);
+  auto loc = createFileLineColLoc(opBuilder);
 
-  if (!blockMergeInfo.try_emplace(curBlock, mergeBlock).second) {
+  if (!blockMergeInfo.try_emplace(curBlock, loc, mergeBlock).second) {
     return emitError(
         unknownLoc,
         "a block cannot have more than one OpSelectionMerge instruction");
   }
 
   return success();
 }
 
LogicalResult Deserializer::processLoopMerge(ArrayRef<uint32_t> operands) {
 
   if (static_cast<uint32_t>(spirv::LoopControl::None) != operands[2]) {
     return emitError(unknownLoc, "unimplmented OpLoopMerge loop control: ")
            << operands[2];
   }
 
   auto *mergeBlock = getOrCreateBlock(operands[0]);
   auto *continueBlock = getOrCreateBlock(operands[1]);
+  auto loc = createFileLineColLoc(opBuilder);
 
-  if (!blockMergeInfo.try_emplace(curBlock, mergeBlock, continueBlock).second) {
+  if (!blockMergeInfo.try_emplace(curBlock, loc, mergeBlock, continueBlock)
+           .second) {
     return emitError(
         unknownLoc,
         "a block cannot have more than one OpLoopMerge instruction");
   }
 
   return success();
 }
 
if (!isFnEntryBlock(block)) {
         mapper.map(blockArg, newArg);
         LLVM_DEBUG(llvm::dbgs() << "[cf] remapped block argument " << blockArg
                                 << " to " << newArg << '\n');
       }
     } else {
       LLVM_DEBUG(llvm::dbgs()
                  << "[cf] block " << block << " is a function entry block\n");
     }
-
     for (auto &op : *block)
       newBlock->push_back(op.clone(mapper));
   }
 
   // Go through all ops and remap the operands.
   auto remapOperands = [&](Operation *op) {
     for (auto &operand : op->getOpOperands())
       if (auto mappedOp = mapper.lookupOrNull(operand.get()))
if (it != blockMergeInfo.end()) {
         newContinue = mapper.lookupOrNull(newContinue);
         assert(newContinue && "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, newMerge, newContinue);
+      blockMergeInfo.try_emplace(newHeader, loc, newMerge, newContinue);
     }
 
     // The structured selection/loop's entry block does not have arguments.
     // If the function's header block is also part of the structured control
     // flow, we cannot just simply erase it because it may contain arguments
     // matching the function signature and used by the cloned blocks.
     if (isFnEntryBlock(block)) {
       LLVM_DEBUG(llvm::dbgs() << "[cf] changing entry block " << block
while (!blockMergeInfo.empty()) {
     LLVM_DEBUG(mergeBlock->print(llvm::dbgs()));
 
     auto *continueBlock = mergeInfo.continueBlock;
     if (continueBlock) {
       LLVM_DEBUG(llvm::dbgs()
                  << "[cf] continue block " << continueBlock << ":\n");
       LLVM_DEBUG(continueBlock->print(llvm::dbgs()));
     }
-
     // Erase this case before calling into structurizer, who will update
     // blockMergeInfo.
     blockMergeInfo.erase(blockMergeInfo.begin());
-    if (failed(ControlFlowStructurizer::structurize(unknownLoc, blockMergeInfo,
-                                                    headerBlock, mergeBlock,
-                                                    continueBlock)))
+    if (failed(ControlFlowStructurizer::structurize(mergeInfo.loc,
+                                                    blockMergeInfo, headerBlock,
+                                                    mergeBlock, continueBlock)))
       return failure();
   }
 
   LLVM_DEBUG(llvm::dbgs() << "[cf] completed structurizing control flow\n");
   return success();
 }
 
 //===----------------------------------------------------------------------===//

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

private:
   spirv::ModuleOp module;
 
   /// An MLIR builder for getting MLIR constructs.
   mlir::Builder mlirBuilder;
 
   /// A flag which indicates if the debuginfo should be emitted.
   bool emitDebugInfo = false;
 
+  /// A flag which indicates if the last processed instruction was a merge
+  /// instruction.
+  /// According to SPIR-V spec: "If a branch merge instruction is used, the last
+  /// OpLine in the block must be before its merge instruction".
+  bool lastProcessedWasMergeInst = false;
+
   /// The <id> of the OpString instruction, which specifies a file name, for
   /// use by other debug instructions.
   uint32_t fileID = 0;
 
   /// The next available result <id>.
   uint32_t nextID = 1;
 
   // The following are for different SPIR-V instruction sections. They follow
LogicalResult Serializer::processSelectionOp(spirv::SelectionOp selectionOp) {
   // resolve properly.
   auto &body = selectionOp.body();
   for (Block &block : body)
     getOrCreateBlockID(&block);
 
   auto *headerBlock = selectionOp.getHeaderBlock();
   auto *mergeBlock = selectionOp.getMergeBlock();
   auto mergeID = getBlockID(mergeBlock);
+  auto loc = selectionOp.getLoc();
 
   // 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 = [&]() {
+    emitDebugLine(functionBody, loc);
+    lastProcessedWasMergeInst = true;
     // TODO(antiagainst): properly support selection control here
     encodeInstructionInto(
         functionBody, spirv::Opcode::OpSelectionMerge,
         {mergeID, static_cast<uint32_t>(spirv::SelectionControl::None)});
   };
   // 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
for (Block &block :
     getOrCreateBlockID(&block);
   }
   auto *headerBlock = loopOp.getHeaderBlock();
   auto *continueBlock = loopOp.getContinueBlock();
   auto *mergeBlock = loopOp.getMergeBlock();
   auto headerID = getBlockID(headerBlock);
   auto continueID = getBlockID(continueBlock);
   auto mergeID = getBlockID(mergeBlock);
+  auto loc = loopOp.getLoc();
 
   // 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});
 
   // 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 = [&]() {
+    emitDebugLine(functionBody, loc);
+    lastProcessedWasMergeInst = true;
     // TODO(antiagainst): properly support loop control here
     encodeInstructionInto(
         functionBody, spirv::Opcode::OpLoopMerge,
         {mergeID, continueID, static_cast<uint32_t>(spirv::LoopControl::None)});
   };
   if (failed(processBlock(headerBlock, /*omitLabel=*/false, emitLoopMerge)))
     return failure();
 
LogicalResult Serializer::processBranchConditionalOp(
   auto falseLabelID = getOrCreateBlockID(condBranchOp.getFalseBlock());
   SmallVector<uint32_t, 5> arguments{conditionID, trueLabelID, falseLabelID};
 
   if (auto weights = condBranchOp.branch_weights()) {
     for (auto val : weights->getValue())
       arguments.push_back(val.cast<IntegerAttr>().getInt());
   }
 
+  emitDebugLine(functionBody, condBranchOp.getLoc());
   return encodeInstructionInto(functionBody, spirv::Opcode::OpBranchConditional,
                                arguments);
 }
 
 LogicalResult Serializer::processBranchOp(spirv::BranchOp branchOp) {
+  emitDebugLine(functionBody, branchOp.getLoc());
   return encodeInstructionInto(functionBody, spirv::Opcode::OpBranch,
                                {getOrCreateBlockID(branchOp.getTarget())});
 }
 
 //===----------------------------------------------------------------------===//
 // Operation
 //===----------------------------------------------------------------------===//
 
LogicalResult Serializer::emitDecoration(uint32_t target,
   return success();
 }
 
 LogicalResult Serializer::emitDebugLine(SmallVectorImpl<uint32_t> &binary,
                                         Location loc) {
   if (!emitDebugInfo)
     return success();
 
+  if (lastProcessedWasMergeInst) {
+    lastProcessedWasMergeInst = false;
+    return success();
+  }
+
   auto fileLoc = loc.dyn_cast<FileLineColLoc>();
   if (fileLoc)
     encodeInstructionInto(binary, spirv::Opcode::OpLine,
                           {fileID, fileLoc.getLine(), fileLoc.getColumn()});
   return success();
 }
 
 LogicalResult spirv::serialize(spirv::ModuleOp module,

mlir/test/Dialect/SPIRV/Serialization/debug.mlir

spv.func @memory_accesses(%arg0 : !spv.ptr<!spv.array<4x!spv.array<4xf32>>, StorageBuffer>, %arg1 : i32, %arg2 : i32) "None" {
     %2 = spv.AccessChain %arg0[%arg1, %arg2] : !spv.ptr<!spv.array<4x!spv.array<4xf32>>, StorageBuffer>
     // CHECK: loc({{".*debug.mlir"}}:63:10)
     %3 = spv.Load "StorageBuffer" %2 : f32
     // CHECK: loc({{.*debug.mlir"}}:65:5)
     spv.Store "StorageBuffer" %2, %3 : f32
     // CHECK: loc({{".*debug.mlir"}}:67:5)
     spv.Return
   }
+
+  spv.func @loop(%count : i32) -> () "None" {
+    %zero = spv.constant 0: i32
+    %one = spv.constant 1: i32
+    %ivar = spv.Variable init(%zero) : !spv.ptr<i32, Function>
+    %jvar = spv.Variable init(%zero) : !spv.ptr<i32, Function>
+    spv.loop {
+      // CHECK: loc({{".*debug.mlir"}}:75:5)
+      spv.Branch ^header
+    ^header:
+      %ival0 = spv.Load "Function" %ivar : i32
+      %icmp = spv.SLessThan %ival0, %count : i32
+      // CHECK: loc({{".*debug.mlir"}}:75:5)
+      spv.BranchConditional %icmp, ^body, ^merge
+    ^body:
+      spv.Store "Function" %jvar, %zero : i32
+      spv.loop {
+        // CHECK: loc({{".*debug.mlir"}}:85:7)
+        spv.Branch ^header
+      ^header:
+        %jval0 = spv.Load "Function" %jvar : i32
+        %jcmp = spv.SLessThan %jval0, %count : i32
+        // CHECK: loc({{".*debug.mlir"}}:85:7)
+        spv.BranchConditional %jcmp, ^body, ^merge
+      ^body:
+        // CHECK: loc({{".*debug.mlir"}}:95:9)
+        spv.Branch ^continue
+      ^continue:
+        %jval1 = spv.Load "Function" %jvar : i32
+        %add = spv.IAdd %jval1, %one : i32
+        spv.Store "Function" %jvar, %add : i32
+        // CHECK: loc({{".*debug.mlir"}}:101:9)
+        spv.Branch ^header
+      ^merge:
+        // CHECK: loc({{".*debug.mlir"}}:85:7)
+        spv._merge
+        // CHECK: loc({{".*debug.mlir"}}:85:7)
+      }
+      // CHECK: loc({{".*debug.mlir"}}:108:7)
+      spv.Branch ^continue
+    ^continue:
+      %ival1 = spv.Load "Function" %ivar : i32
+      %add = spv.IAdd %ival1, %one : i32
+      spv.Store "Function" %ivar, %add : i32
+      // CHECK: loc({{".*debug.mlir"}}:114:7)
+      spv.Branch ^header
+    ^merge:
+      // CHECK: loc({{".*debug.mlir"}}:75:5)
+      spv._merge
+    // CHECK: loc({{".*debug.mlir"}}:75:5)
+    }
+    spv.Return
+  }
+
+  spv.func @selection(%cond: i1) -> () "None" {
+    %zero = spv.constant 0: i32
+    %one = spv.constant 1: i32
+    %two = spv.constant 2: i32
+    %var = spv.Variable init(%zero) : !spv.ptr<i32, Function>
+    spv.selection {
+      // CHECK: loc({{".*debug.mlir"}}:128:5)
+      spv.BranchConditional %cond [5, 10], ^then, ^else
+    ^then:
+      spv.Store "Function" %var, %one : i32
+      // CHECK: loc({{".*debug.mlir"}}:134:7)
+      spv.Branch ^merge
+    ^else:
+      spv.Store "Function" %var, %two : i32
+      // CHECK: loc({{".*debug.mlir"}}:138:7)
+      spv.Branch ^merge
+    ^merge:
+      // CHECK: loc({{".*debug.mlir"}}:128:5)
+      spv._merge
+    // CHECK: loc({{".*debug.mlir"}}:128:5)
+    }
+    spv.Return
+  }
 }