[MLIR][SPIRV] Control attributes parsing/printing for loop and selection

[georgemitenkov]

Summary

This patch handles loopControl and selectionControl in parsing and
printing. In order to reuse the functionality, and avoid handling cases when
{ of the region is parsed as a dictionary attribute, control keyword was
introduced.None is a default control attribute.

This functionality can be later extended to spv.func.

[antiagainst]

Awesome! Can we also support (de)serialization for this?

[georgemitenkov]

Added support of (de)serialization of loop and selection control.

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

 static constexpr const char kMemoryAccessAttrName[] = "memory_access";
 static constexpr const char kSourceMemoryAccessAttrName[] =
     "source_memory_access";
 static constexpr const char kAlignmentAttrName[] = "alignment";
 static constexpr const char kSourceAlignmentAttrName[] = "source_alignment";
 static constexpr const char kBranchWeightAttrName[] = "branch_weights";
 static constexpr const char kCallee[] = "callee";
 static constexpr const char kClusterSize[] = "cluster_size";
+static constexpr const char kControl[] = "control";
 static constexpr const char kDefaultValueAttrName[] = "default_value";
 static constexpr const char kExecutionScopeAttrName[] = "execution_scope";
 static constexpr const char kEqualSemanticsAttrName[] = "equal_semantics";
 static constexpr const char kFnNameAttrName[] = "fn";
 static constexpr const char kGroupOperationAttrName[] = "group_operation";
 static constexpr const char kIndicesAttrName[] = "indices";
 static constexpr const char kInitializerAttrName[] = "initializer";
 static constexpr const char kInterfaceAttrName[] = "interface";
parseEnumKeywordAttr(EnumClass &value, OpAsmParser &parser,
   if (parseEnumKeywordAttr(value, parser)) {
     return failure();
   }
   state.addAttribute(attrName, parser.getBuilder().getI32IntegerAttr(
                                    llvm::bit_cast<int32_t>(value)));
   return success();
 }
 
+/// Parses Function, Selection and Loop control attributes. If no control is
+/// specified, "None" is used as a default.
+template <typename EnumClass>
+static ParseResult
+parseControlAttribute(OpAsmParser &parser, OperationState &state,
+                      StringRef attrName = spirv::attributeName<EnumClass>()) {
+  if (succeeded(parser.parseOptionalKeyword(kControl))) {
+    EnumClass control;
+    if (parser.parseLParen() || parseEnumKeywordAttr(control, parser, state) ||
+        parser.parseRParen())
+      return failure();
+    return success();
+  }
+  // Set control to "None" otherwise.
+  Builder builder = parser.getBuilder();
+  state.addAttribute(attrName, builder.getI32IntegerAttr(0));
+  return success();
+}
+
 /// Parses optional memory access attributes attached to a memory access
 /// operand/pointer. Specifically, parses the following syntax:
 ///     (`[` memory-access `]`)?
 /// where:
 ///     memory-access ::= `"None"` | `"Volatile"` | `"Aligned", `
 ///         integer-literal | `"NonTemporal"`
 static ParseResult parseMemoryAccessAttributes(OpAsmParser &parser,
                                                OperationState &state) {

 void spirv::LoopOp::build(OpBuilder &builder, OperationState &state) {
   state.addAttribute("loop_control",
                      builder.getI32IntegerAttr(
                          static_cast<uint32_t>(spirv::LoopControl::None)));
   state.addRegion();
 }
 
 static ParseResult parseLoopOp(OpAsmParser &parser, OperationState &state) {
-  // TODO: support loop control properly
-  Builder builder = parser.getBuilder();
-  state.addAttribute("loop_control",
-                     builder.getI32IntegerAttr(
-                         static_cast<uint32_t>(spirv::LoopControl::None)));
-
+  if (parseControlAttribute<spirv::LoopControl>(parser, state))
+    return failure();
   return parser.parseRegion(*state.addRegion(), /*arguments=*/{},
                             /*argTypes=*/{});
 }
 
 static void print(spirv::LoopOp loopOp, OpAsmPrinter &printer) {
   auto *op = loopOp.getOperation();
 
   printer << spirv::LoopOp::getOperationName();
+  auto control = loopOp.loop_control();
+  if (control != spirv::LoopControl::None)
+    printer << " control(" << spirv::stringifyLoopControl(control) << ")";
   printer.printRegion(op->getRegion(0), /*printEntryBlockArgs=*/false,
                       /*printBlockTerminators=*/true);
 }
 
 /// Returns true if the given `srcBlock` contains only one `spv.Branch` to the
 /// given `dstBlock`.
 static inline bool hasOneBranchOpTo(Block &srcBlock, Block &dstBlock) {
   // Check that there is only one op in the `srcBlock`.

 }
 
 //===----------------------------------------------------------------------===//
 // spv.selection
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseSelectionOp(OpAsmParser &parser,
                                     OperationState &state) {
-  // TODO: support selection control properly
-  Builder builder = parser.getBuilder();
-  state.addAttribute("selection_control",
-                     builder.getI32IntegerAttr(
-                         static_cast<uint32_t>(spirv::SelectionControl::None)));
-
+  if (parseControlAttribute<spirv::SelectionControl>(parser, state))
+    return failure();
   return parser.parseRegion(*state.addRegion(), /*arguments=*/{},
                             /*argTypes=*/{});
 }
 
 static void print(spirv::SelectionOp selectionOp, OpAsmPrinter &printer) {
   auto *op = selectionOp.getOperation();
 
   printer << spirv::SelectionOp::getOperationName();
+  auto control = selectionOp.selection_control();
+  if (control != spirv::SelectionControl::None)
+    printer << " control(" << spirv::stringifySelectionControl(control) << ")";
   printer.printRegion(op->getRegion(0), /*printEntryBlockArgs=*/false,
                       /*printBlockTerminators=*/true);
 }
 
 static LogicalResult verify(spirv::SelectionOp selectionOp) {
   auto *op = selectionOp.getOperation();
 
   // We need to verify that the blocks follow the following layout:

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

 ///
 /// 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;
+  uint32_t control;
 
-  BlockMergeInfo(Location location)
-      : mergeBlock(nullptr), continueBlock(nullptr), loc(location) {}
-  BlockMergeInfo(Location location, Block *m, Block *c = nullptr)
-      : mergeBlock(m), continueBlock(c), loc(location) {}
+  BlockMergeInfo(Location location, uint32_t control)
+      : mergeBlock(nullptr), continueBlock(nullptr), loc(location),
+        control(control) {}
+  BlockMergeInfo(Location location, uint32_t control, Block *m,
+                 Block *c = nullptr)
+      : mergeBlock(m), continueBlock(c), loc(location), control(control) {}
 };
 
 /// A struct for containing OpLine instruction information.
 struct DebugLine {
   uint32_t fileID;
   uint32_t line;
   uint32_t col;
 
LogicalResult Deserializer::processSelectionMerge(ArrayRef<uint32_t> operands) {
   }
 
   if (operands.size() < 2) {
     return emitError(
         unknownLoc,
         "OpSelectionMerge must specify merge target and selection control");
   }
 
-  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);
+  auto selectionControl = operands[1];
 
-  if (!blockMergeInfo.try_emplace(curBlock, loc, mergeBlock).second) {
+  if (!blockMergeInfo.try_emplace(curBlock, loc, selectionControl, mergeBlock)
+           .second) {
     return emitError(
         unknownLoc,
         "a block cannot have more than one OpSelectionMerge instruction");
   }
 
   return success();
 }
 
 LogicalResult Deserializer::processLoopMerge(ArrayRef<uint32_t> operands) {
   if (!curBlock) {
     return emitError(unknownLoc, "OpLoopMerge must appear in a block");
   }
 
   if (operands.size() < 3) {
     return emitError(unknownLoc, "OpLoopMerge must specify merge target, "
                                  "continue target and loop control");
   }
 
-  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);
+  uint32_t loopControl = operands[2];
 
-  if (!blockMergeInfo.try_emplace(curBlock, loc, mergeBlock, continueBlock)
+  if (!blockMergeInfo
+           .try_emplace(curBlock, loc, loopControl, mergeBlock, continueBlock)
            .second) {
     return emitError(
         unknownLoc,
         "a block cannot have more than one OpLoopMerge instruction");
   }
 
   return success();
 }

 public:
   /// Structurizes the loop at the given `headerBlock`.
   ///
   /// This method will create an spv.loop op in the `mergeBlock` and move all
   /// blocks in the structured loop into the spv.loop's region. All branches to
   /// the `headerBlock` will be redirected to the `mergeBlock`.
   /// This method will also update `mergeInfo` by remapping all blocks inside to
   /// the newly cloned ones inside structured control flow op's regions.
-  static LogicalResult structurize(Location loc, BlockMergeInfoMap &mergeInfo,
+  static LogicalResult structurize(Location loc, uint32_t control,
+                                   BlockMergeInfoMap &mergeInfo,
                                    Block *headerBlock, Block *mergeBlock,
                                    Block *continueBlock) {
-    return ControlFlowStructurizer(loc, mergeInfo, headerBlock, mergeBlock,
-                                   continueBlock)
+    return ControlFlowStructurizer(loc, control, mergeInfo, headerBlock,
+                                   mergeBlock, continueBlock)
         .structurizeImpl();
   }
 
 private:
-  ControlFlowStructurizer(Location loc, BlockMergeInfoMap &mergeInfo,
-                          Block *header, Block *merge, Block *cont)
-      : location(loc), blockMergeInfo(mergeInfo), headerBlock(header),
-        mergeBlock(merge), continueBlock(cont) {}
+  ControlFlowStructurizer(Location loc, uint32_t control,
+                          BlockMergeInfoMap &mergeInfo, Block *header,
+                          Block *merge, Block *cont)
+      : location(loc), control(control), blockMergeInfo(mergeInfo),
+        headerBlock(header), mergeBlock(merge), continueBlock(cont) {}
 
   /// Creates a new spv.selection op at the beginning of the `mergeBlock`.
-  spirv::SelectionOp createSelectionOp();
+  spirv::SelectionOp createSelectionOp(uint32_t selectionControl);
 
   /// Creates a new spv.loop op at the beginning of the `mergeBlock`.
-  spirv::LoopOp createLoopOp();
+  spirv::LoopOp createLoopOp(uint32_t loopControl);
 
   /// Collects all blocks reachable from `headerBlock` except `mergeBlock`.
   void collectBlocksInConstruct();
 
   LogicalResult structurizeImpl();
 
   Location location;
+  uint32_t control;
 
   BlockMergeInfoMap &blockMergeInfo;
 
   Block *headerBlock;
   Block *mergeBlock;
   Block *continueBlock; // nullptr for spv.selection
 
   llvm::SetVector<Block *> constructBlocks;
 };
 } // namespace
 
-spirv::SelectionOp ControlFlowStructurizer::createSelectionOp() {
+spirv::SelectionOp
+ControlFlowStructurizer::createSelectionOp(uint32_t selectionControl) {
   // Create a builder and set the insertion point to the beginning of the
   // merge block so that the newly created SelectionOp will be inserted there.
   OpBuilder builder(&mergeBlock->front());
 
-  auto control = builder.getI32IntegerAttr(
-      static_cast<uint32_t>(spirv::SelectionControl::None));
+  auto control = builder.getI32IntegerAttr(selectionControl);
   auto selectionOp = builder.create<spirv::SelectionOp>(location, control);
   selectionOp.addMergeBlock();
 
   return selectionOp;
 }
 
-spirv::LoopOp ControlFlowStructurizer::createLoopOp() {
+spirv::LoopOp ControlFlowStructurizer::createLoopOp(uint32_t loopControl) {
   // Create a builder and set the insertion point to the beginning of the
   // merge block so that the newly created LoopOp will be inserted there.
   OpBuilder builder(&mergeBlock->front());
 
-  // TODO: handle loop control properly
-  auto loopOp = builder.create<spirv::LoopOp>(location);
+  auto control = builder.getI32IntegerAttr(loopControl);
+  auto loopOp = builder.create<spirv::LoopOp>(location, control);
   loopOp.addEntryAndMergeBlock();
 
   return loopOp;
 }
 
 void ControlFlowStructurizer::collectBlocksInConstruct() {
   assert(constructBlocks.empty() && "expected empty constructBlocks");
 
   // Put the header block in the work list first.
   constructBlocks.insert(headerBlock);
 
   // For each item in the work list, add its successors excluding the merge
   // block.
   for (unsigned i = 0; i < constructBlocks.size(); ++i) {
     for (auto *successor : constructBlocks[i]->getSuccessors())
       if (successor != mergeBlock)
         constructBlocks.insert(successor);
   }
 }
 
 LogicalResult ControlFlowStructurizer::structurizeImpl() {
   Operation *op = nullptr;
   bool isLoop = continueBlock != nullptr;
   if (isLoop) {
-    if (auto loopOp = createLoopOp())
+    if (auto loopOp = createLoopOp(control))
       op = loopOp.getOperation();
   } else {
-    if (auto selectionOp = createSelectionOp())
+    if (auto selectionOp = createSelectionOp(control))
       op = selectionOp.getOperation();
   }
   if (!op)
     return failure();
   Region &body = op->getRegion(0);
 
   BlockAndValueMapping mapper;
   // All references to the old merge block should be directed to the
if (it != blockMergeInfo.end()) {
       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, newMerge, newContinue);
+      blockMergeInfo.try_emplace(newHeader, loc, it->second.control, 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()) {
     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(mergeInfo.loc,
-                                                    blockMergeInfo, headerBlock,
-                                                    mergeBlock, continueBlock)))
+    if (failed(ControlFlowStructurizer::structurize(
+            mergeInfo.loc, mergeInfo.control, blockMergeInfo, headerBlock,
       return failure();
   }
 
   LLVM_DEBUG(llvm::dbgs() << "[cf] completed structurizing control flow\n");
   return success();
 }
 
 //===----------------------------------------------------------------------===//

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

LogicalResult Serializer::processSelectionOp(spirv::SelectionOp selectionOp) {
   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: properly support selection control here
     encodeInstructionInto(
         functionBody, spirv::Opcode::OpSelectionMerge,
-        {mergeID, static_cast<uint32_t>(spirv::SelectionControl::None)});
+        {mergeID, static_cast<uint32_t>(selectionOp.selection_control())});
   };
   // 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.selection op. If there are ops ahead of the spv.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)))
LogicalResult Serializer::processLoopOp(spirv::LoopOp loopOp) {
   // 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: properly support loop control here
     encodeInstructionInto(
         functionBody, spirv::Opcode::OpLoopMerge,
-        {mergeID, continueID, static_cast<uint32_t>(spirv::LoopControl::None)});
+        {mergeID, continueID, static_cast<uint32_t>(loopOp.loop_control())});
   };
   if (failed(processBlock(headerBlock, /*omitLabel=*/false, emitLoopMerge)))
     return failure();
 
   // Process all blocks with a depth-first visitor starting from the header
   // block. The loop header block, loop continue block, and loop merge block are
   // skipped by this visitor and handled later in this function.
   if (failed(visitInPrettyBlockOrder(

mlir/test/Dialect/SPIRV/Serialization/loop.mlir

spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   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>
 
 // CHECK:        spv.Branch ^bb1
 // CHECK-NEXT: ^bb1:
-// CHECK-NEXT:   spv.loop
-    spv.loop {
+// CHECK-NEXT:   spv.loop control(Unroll)
+    spv.loop control(Unroll) {
 // CHECK-NEXT:     spv.Branch ^bb1
       spv.Branch ^header
 
 // CHECK-NEXT:   ^bb1:
     ^header:
 // CHECK-NEXT:     spv.Load
       %ival0 = spv.Load "Function" %ivar : i32
 // CHECK-NEXT:     spv.SLessThan
       %icmp = spv.SLessThan %ival0, %count : i32
 // CHECK-NEXT:     spv.BranchConditional %{{.*}}, ^bb2, ^bb5
       spv.BranchConditional %icmp, ^body, ^merge
 
 // CHECK-NEXT:   ^bb2:
     ^body:
 // CHECK-NEXT:     spv.constant 0
 // CHECK-NEXT:     spv.Store
       spv.Store "Function" %jvar, %zero : i32
 // CHECK-NEXT:     spv.Branch ^bb3
 // CHECK-NEXT:   ^bb3:
-// CHECK-NEXT:     spv.loop
-      spv.loop {
+// CHECK-NEXT:     spv.loop control(DontUnroll)
+      spv.loop control(DontUnroll) {
 // CHECK-NEXT:       spv.Branch ^bb1
         spv.Branch ^header
 
 // CHECK-NEXT:     ^bb1:
       ^header:
 // CHECK-NEXT:       spv.Load
         %jval0 = spv.Load "Function" %jvar : i32
 // CHECK-NEXT:       spv.SLessThan

mlir/test/Dialect/SPIRV/Serialization/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], []> {
   spv.func @selection(%cond: i1) -> () "None" {
 // CHECK:        spv.Branch ^bb1
 // CHECK-NEXT: ^bb1:
     %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.selection {
+// CHECK-NEXT:   spv.selection control(Flatten)
 // CHECK-NEXT:     spv.constant 0
 // CHECK-NEXT:     spv.Variable
-    spv.selection {
+    spv.selection control(Flatten) {
 // CHECK-NEXT: spv.BranchConditional %{{.*}} [5, 10], ^bb1, ^bb2
       spv.BranchConditional %cond [5, 10], ^then, ^else
 
 // CHECK-NEXT:   ^bb1:
     ^then:
 // CHECK-NEXT:     spv.constant 1
 // CHECK-NEXT:     spv.Store
       spv.Store "Function" %var, %one : i32

mlir/test/Dialect/SPIRV/control-flow-ops.mlir

func @empty_region() -> () {
   // CHECK: spv.loop
   spv.loop {
   }
   return
 }
 
 // -----
 
+// CHECK-LABEL: @loop_with_control
+func @loop_with_control() -> () {
+  // CHECK: spv.loop control(Unroll)
+  spv.loop control(Unroll) {
+  }
+  return
+}
+
+// -----
+
 func @wrong_merge_block() -> () {
   // expected-error @+1 {{last block must be the merge block with only one 'spv._merge' op}}
   spv.loop {
     spv.Return
   }
   return
 }
 
func @empty_region() -> () {
   // CHECK: spv.selection
   spv.selection {
   }
   return
 }
 
 // -----
 
+// CHECK-LABEL: @selection_with_control
+func @selection_with_control() -> () {
+  // CHECK: spv.selection control(Flatten)
+  spv.selection control(Flatten) {
+  }
+  return
+}
+
+// -----
+
 func @wrong_merge_block() -> () {
   // expected-error @+1 {{last block must be the merge block with only one 'spv._merge' op}}
   spv.selection {
     spv.Return
   }
   return
 }