[MLIR][SPIRVToLLVM] Conversion pattern for loop op

[georgemitenkov]

Summary

This patch introduces a conversion of spv.loop to LLVM dialect.
Similarly to spv.selection, op's control attributes are not mapped to LLVM
yet and therefore the conversion fails if the loop control is not None. Also,
all blocks within the loop should be reachable in order for conversion to
succeed.

[mravishankar]

Nice!

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

736

This seems out of place here. Better to add it as a canonicalization or folding pattern for spirv::LoopOp

745

Block *currentBlock = loopOp.getBlock() ?

Also the convention of auto is to use it only when its clear from the context what the type is or if it is too verbose.

747

auto position = Block::iterator(loopOp) ?

Note that here auto is OK.

Btw, if you are changing the insertion point just the get the Operation after loop, you can just say loopOp->getNextNode()

[rriddle]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

736

Please do not use .size(), it is O(N). Use loopOp.body().empty() instead.

[georgemitenkov]

Addressed comments

[georgemitenkov]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

745

Not sure loopOp.getBlock() works in this case? I think you meant rewriter.getBlock() to get the current block of the builder

[mravishankar]

THanks for the update. Minor edit comments here. Sorry didnt post these earlierl

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

722

I think some more comments about the loop structure in LLVM dialect would be helpful here.

745

Right. I meant rewriter.getBlock() . Thanks for interpreting it correctly!

mlir/test/Conversion/SPIRVToLLVM/control-flow-ops-to-llvm.mlir

88

Instead of splitting the CHECK: across the entire IR snippet, its better and more readable to write it as a sequence of CHECK instructions. Also do not hard code basic block IDs. For example, the below could be written as

     CHECK:   llvm.br ^[[BB1:.*]]
     CHECK: ^[[BB1]]:
     CHECK:   %[[COND:.*]] = ...
CHECK-NEXT:   llvm.cond_br %[[COND]], ^[[BB2:.*]], ^[[BB4:.*]]
     CHECK: ^[{BB2]]:
     CHECK:   llvm.br ^[[BB3:.*]]
     CHECK: ^[[BB3]]:
     CHECK:   llvm.br ^[{BB1]]
     CHECK: ^[[BB4]]
     CHECK:   llvm.br ^[[BB5:.*]]
     CHECK: ^[[BB5]]:
CHECK-NEXT:   llvm.return

[georgemitenkov]

Addressed comments:

• Put check statements in one block
• Added a diagram of the loop structure in LLVM

[georgemitenkov]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

722

Would the diagram be sufficient, or do you think more details should be added?

[mravishankar]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

723

Wow! THis is super nice! Thanks

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

public:
   matchAndRewrite(spirv::ReturnValueOp returnValueOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(returnValueOp, ArrayRef<Type>(),
                                                 operands);
     return success();
   }
 };
 
+/// Converts `spv.loop` to LLVM dialect. All blocks within selection should be
+/// reachable for conversion to succeed.
+/// The structure of the loop in LLVM dialect will be the following:
+///
+///      +------------------------------------+
+///      | <code before spv.loop>             |
+///      | llvm.br ^header                    |
+///      +------------------------------------+
+///                           |
+///   +----------------+      |
+///   |                |      |
+///   |                V      V
+///   |  +------------------------------------+
+///   |  | ^header:                           |
+///   |  |   <header code>                    |
+///   |  |   llvm.cond_br %cond, ^body, ^exit |
+///   |  +------------------------------------+
+///   |                    |
+///   |                    |----------------------+
+///   |                    |                      |
+///   |                    V                      |
+///   |  +------------------------------------+   |
+///   |  | ^body:                             |   |
+///   |  |   <body code>                      |   |
+///   |  |   llvm.br ^continue                |   |
+///   |  +------------------------------------+   |
+///   |                    |                      |
+///   |                    V                      |
+///   |  +------------------------------------+   |
+///   |  | ^continue:                         |   |
+///   |  |   <continue code>                  |   |
+///   |  |   llvm.br ^header                  |   |
+///   |  +------------------------------------+   |
+///   |               |                           |
+///   +---------------+    +----------------------+
+///                        |
+///                        V
+///      +------------------------------------+
+///      | ^exit:                             |
+///      |   llvm.br ^remaining               |
+///      +------------------------------------+
+///                        |
+///                        V
+///      +------------------------------------+
+///      | ^remaining:                        |
+///      |   <code after spv.loop>            |
+///      +------------------------------------+
+///
+class LoopPattern : public SPIRVToLLVMConversion<spirv::LoopOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::LoopOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::LoopOp loopOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    // There is no support of loop control at the moment.
+    if (loopOp.loop_control() != spirv::LoopControl::None)
+      return failure();
+
+    Location loc = loopOp.getLoc();
+
+    // Split the current block after `spv.loop`. The remaing ops will be used in
+    // `endBlock`.
+    Block *currentBlock = rewriter.getBlock();
+    auto position = Block::iterator(loopOp);
+    Block *endBlock = rewriter.splitBlock(currentBlock, position);
+
+    // Remove entry block and create a branch in the current block going to the
+    // header block.
+    Block *entryBlock = loopOp.getEntryBlock();
+    assert(entryBlock->getOperations().size() == 1);
+    auto brOp = dyn_cast<spirv::BranchOp>(entryBlock->getOperations().front());
+    if (!brOp)
+      return failure();
+    Block *headerBlock = loopOp.getHeaderBlock();
+    rewriter.setInsertionPointToEnd(currentBlock);
+    rewriter.create<LLVM::BrOp>(loc, brOp.getBlockArguments(), headerBlock);
+    rewriter.eraseBlock(entryBlock);
+
+    // Branch from merge block to end block.
+    Block *mergeBlock = loopOp.getMergeBlock();
+    Operation *terminator = mergeBlock->getTerminator();
+    ValueRange terminatorOperands = terminator->getOperands();
+    rewriter.setInsertionPointToEnd(mergeBlock);
+    rewriter.create<LLVM::BrOp>(loc, terminatorOperands, endBlock);
+
+    rewriter.inlineRegionBefore(loopOp.body(), endBlock);
+    rewriter.replaceOp(loopOp, endBlock->getArguments());
+    return success();
+  }
+};
+
 class MergePattern : public SPIRVToLLVMConversion<spirv::MergeOp> {
 public:
   using SPIRVToLLVMConversion<spirv::MergeOp>::SPIRVToLLVMConversion;
 
   LogicalResult
   matchAndRewrite(spirv::MergeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.eraseOp(op);
patterns.insert<
       IComparePattern<spirv::UGreaterThanEqualOp, LLVM::ICmpPredicate::uge>,
       IComparePattern<spirv::ULessThanEqualOp, LLVM::ICmpPredicate::ule>,
       IComparePattern<spirv::ULessThanOp, LLVM::ICmpPredicate::ult>,
 
       // Constant op
       ConstantScalarAndVectorPattern,
 
       // Control Flow ops
-      BranchConversionPattern, BranchConditionalConversionPattern,
+      BranchConversionPattern, BranchConditionalConversionPattern, LoopPattern,
       SelectionPattern, MergePattern,
 
       // Function Call op
       FunctionCallPattern,
 
       // GLSL extended instruction set ops
       DirectConversionPattern<spirv::GLSLCeilOp, LLVM::FCeilOp>,
       DirectConversionPattern<spirv::GLSLCosOp, LLVM::CosOp>,

mlir/test/Conversion/SPIRVToLLVM/control-flow-ops-to-llvm.mlir

spv.module Logical GLSL450 {
   ^false:
     spv.Return
   }
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
+// spv.loop
+//===----------------------------------------------------------------------===//
+
+spv.module Logical GLSL450 {
+  // CHECK-LABEL: @infinite_loop
+  spv.func @infinite_loop(%count : i32) -> () "None" {
+    // CHECK:   llvm.br ^[[BB1:.*]]
+    // CHECK: ^[[BB1]]:
+    // CHECK:   %[[COND:.*]] = llvm.mlir.constant(true) : !llvm.i1
+    // CHECK:   llvm.cond_br %[[COND]], ^[[BB2:.*]], ^[[BB4:.*]]
+    // CHECK: ^[[BB2]]:
+    // CHECK:   llvm.br ^[[BB3:.*]]
+    // CHECK: ^[[BB3]]:
+    // CHECK:   llvm.br ^[[BB1:.*]]
+    // CHECK: ^[[BB4]]:
+    // CHECK:   llvm.br ^[[BB5:.*]]
+    // CHECK: ^[[BB5]]:
+    // CHECK:   llvm.return
+    spv.loop {
+      spv.Branch ^header
+    ^header:
+      %cond = spv.constant true
+      spv.BranchConditional %cond, ^body, ^merge
+    ^body:
+      // Do nothing
+      spv.Branch ^continue
+    ^continue:
+      // Do nothing
+      spv.Branch ^header
+    ^merge:
+      spv._merge
+    }
+    spv.Return
+  }
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // spv.selection
 //===----------------------------------------------------------------------===//
 
 spv.module Logical GLSL450 {
   spv.func @selection_empty() -> () "None" {
     // CHECK: llvm.return
     spv.selection {
     }