[mlir] spirv: Add scf.while spirv conversion

[Hardcode84]

Summary

• It works similar to scf.for coversion, but convert condition and yield ops as part of scf.whille pattern so it don't need to maintain external state

[antiagainst]

Awesome! Yes this is a major missing pattern. I was thinking to add it by myself they other day. Thanks for handling it! Sorry I've quite a few nits, given this is control flow. Mostly around doc though. :)

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

345

Nit: I'd prefer to just name them as beforeRegion and afterRegion for clarity. It's not quite worth saving those extra keystrokes esp. nowadays autocompletion is so good. :)

351

Please add some comment so later it's easier to understand/modify the logic here:

// Move the while before block as the initial loop header block.

354

// Move the while after block as the initial loop body block.

357

// Jump from the loop entry block to the loop header block.

362

Could you explain why we maintain the translation of result values inside this pattern by itself instead of using the common utilities? Something like

// For other SCF ops, the scf.yield op yields the value for the whole SCF op. So we use the scf.yield op as the anchor to create/load/store SPIR-V local variables. But for the scf.while op, the scf.yield op yields a value for the before region, which may not matching the whole op's result. Instead, the scf.condition op returns values matching the whole op's results. So we need to create/load/store variables according to that.

370

// Create local variables before the scf.while op.

375

// Load the final result values after the scf.while op.

379

// Store the current iteration's result value.

388

// Convert the scf.yield op to a branch back to the header block.

mlir/test/Conversion/SCFToSPIRV/while.mlir

39

Could you add an example where the before and after regions take different types of args? You can see the third example in the doc: https://mlir.llvm.org/docs/Dialects/SCFDialect/#scfwhile-mlirscfwhileop.

[Hardcode84]

Side question: is there any particular reason why ForOp and IfOp patterns handle yieldOp via separate pattern? It seems to complicate things as it need to maintain external state (also, I am not 100% sure my handling of yieldOp and condOp is correct wrt dialect conversion).

Also, can we remove existing ForOp lowering and use scf-for-to-while pass instead (so less code to maintain) after this is merged?

[Hardcode84]

adress review comments

[antiagainst]

In D113007#3103806, @Hardcode84 wrote:

Side question: is there any particular reason why ForOp and IfOp patterns handle yieldOp via separate pattern? It seems to complicate things as it need to maintain external state (also, I am not 100% sure my handling of yieldOp and condOp is correct wrt dialect conversion).

Ah actually good point; sorry I neglected this particular part at the beginning. You are right this will be problematic with dialect conversion. The way it works in dialect conversion is that dialect conversion handles op conversion from top to bottom, from external regions to internal regions. In the process the framework will try to automatically handle type conversion for operands (including using unrealized_conversion_cast if allowed by the type converter) and give access to the type converted operands in patterns via OpAdaptor. This whole process does not materialize immediately for each op; instead it "caches" various conversion steps internally and apply once at the very end if everything turns out to be fine. What we have here will intervene the framework because we handle condition ops and yield ops when still in the while op: that's not expected by the framework. It will cause issues when we need to do special handling for those condition/yield ops too (e.g., if the integer type in condition/yield is not supported by the target environment then we need to emulate it or just abort conversion. With the current impl I think that is not properly handled. I think if you remove the Int64 in the capability list in your test you'll see errors.)

Though the above just describes what a normal dialect conversion flow works. In SCF -> SPIR-V specifically, there are additional trickiness, as commented at L64-L77 in the cpp file. I won't repeat here; you can read there.

In general this kind of tightly-coupled-ops-with-regions are hard and tricky to handle. I think we need to handle the scf.condition and scf.yield in the same way it is handled right now for scf.for ops.

It would be great if you can give it a go to fix. But it's also entirely fine if you'd like to leave it to me given that these are all tricky stuff. So I'll accept this and feel free to land. We can iterate on it.

Also, can we remove existing ForOp lowering and use scf-for-to-while pass instead (so less code to maintain) after this is merged?

Let's take it step by step. The existing scf.for conversion is quite load bearing. We first need to have the scf.while op well supported and then see how to migrate over.

[Hardcode84]

Can it be solved by splitting SCF conversion into 2 steps:

1. convert types of scf.for, while, yield and cond using DialectConversion, without converting operations itself
2. Convert them using simple RewritePattern + greedyRewriter (all types were already successfully converted at this stage)

?

[rriddle]

In D113007#3125393, @antiagainst wrote:

In D113007#3103806, @Hardcode84 wrote:

Side question: is there any particular reason why ForOp and IfOp patterns handle yieldOp via separate pattern? It seems to complicate things as it need to maintain external state (also, I am not 100% sure my handling of yieldOp and condOp is correct wrt dialect conversion).

Ah actually good point; sorry I neglected this particular part at the beginning. You are right this will be problematic with dialect conversion. The way it works in dialect conversion is that dialect conversion handles op conversion from top to bottom, from external regions to internal regions. In the process the framework will try to automatically handle type conversion for operands (including using unrealized_conversion_cast if allowed by the type converter) and give access to the type converted operands in patterns via OpAdaptor. This whole process does not materialize immediately for each op; instead it "caches" various conversion steps internally and apply once at the very end if everything turns out to be fine. What we have here will intervene the framework because we handle condition ops and yield ops when still in the while op: that's not expected by the framework.

That is fine, you should be able to use getRemappedValues on the yield operations operands, which would insert temporary casts as necessary. getRemappedValues should be able to properly handle unconverted values given that D111620 landed.

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

64–77

Is this comment still true, this pattern is passing the type converter into the conversion pattern.

[Hardcode84]

use getRemappedValue

[antiagainst]

In D113007#3125674, @rriddle wrote:

In D113007#3125393, @antiagainst wrote:

In D113007#3103806, @Hardcode84 wrote:

Side question: is there any particular reason why ForOp and IfOp patterns handle yieldOp via separate pattern? It seems to complicate things as it need to maintain external state (also, I am not 100% sure my handling of yieldOp and condOp is correct wrt dialect conversion).

Ah actually good point; sorry I neglected this particular part at the beginning. You are right this will be problematic with dialect conversion. The way it works in dialect conversion is that dialect conversion handles op conversion from top to bottom, from external regions to internal regions. In the process the framework will try to automatically handle type conversion for operands (including using unrealized_conversion_cast if allowed by the type converter) and give access to the type converted operands in patterns via OpAdaptor. This whole process does not materialize immediately for each op; instead it "caches" various conversion steps internally and apply once at the very end if everything turns out to be fine. What we have here will intervene the framework because we handle condition ops and yield ops when still in the while op: that's not expected by the framework.

That is fine, you should be able to use getRemappedValues on the yield operations operands, which would insert temporary casts as necessary. getRemappedValues should be able to properly handle unconverted values given that D111620 landed.

Oh wow, I didn't notice that. Seems relatively new. As always, River come to the rescue!

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

64–77

Probably not true anymore given that I haven't seen anything broken since you changed this. I'll verify and clean up. Thanks a ton for fixing this, River!

[antiagainst]

Thanks for adding this again, @Hardcode84!

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

 public:
   SCFToSPIRVPattern<OpTy>(MLIRContext *context, SPIRVTypeConverter &converter,
                           ScfToSPIRVContextImpl *scfToSPIRVContext)
       : OpConversionPattern<OpTy>::OpConversionPattern(converter, context),
         scfToSPIRVContext(scfToSPIRVContext), typeConverter(converter) {}
 
 protected:
   ScfToSPIRVContextImpl *scfToSPIRVContext;
   // FIXME: We explicitly keep a reference of the type converter here instead of
   // passing it to OpConversionPattern during construction. This effectively
   // bypasses the conversion framework's automation on type conversion. This is
   // needed right now because the conversion framework will unconditionally
   // legalize all types used by SCF ops upon discovering them, for example, the
   // types of loop carried values. We use SPIR-V variables for those loop
   // carried values. Depending on the available capabilities, the SPIR-V
   // variable can be different, for example, cooperative matrix or normal
   // variable. We'd like to detach the conversion of the loop carried values
   // from the SCF ops (which is mainly a region). So we need to "mark" types
   // used by SCF ops as legal, if to use the conversion framework for type
   // conversion. There isn't a straightforward way to do that yet, as when
   // converting types, ops aren't taken into consideration. Therefore, we just
   // bypass the framework's type conversion for now.
   SPIRVTypeConverter &typeConverter;
 };
 
 /// Pattern to convert a scf::ForOp within kernel functions into spirv::LoopOp.
 class ForOpConversion final : public SCFToSPIRVPattern<scf::ForOp> {
 public:
   using SCFToSPIRVPattern<scf::ForOp>::SCFToSPIRVPattern;
 

 class TerminatorOpConversion final : public SCFToSPIRVPattern<scf::YieldOp> {
 public:
   using SCFToSPIRVPattern<scf::YieldOp>::SCFToSPIRVPattern;
 
   LogicalResult
   matchAndRewrite(scf::YieldOp terminatorOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
+
+class WhileOpConversion final : public SCFToSPIRVPattern<scf::WhileOp> {
+public:
+  using SCFToSPIRVPattern<scf::WhileOp>::SCFToSPIRVPattern;
+
+  LogicalResult
+  matchAndRewrite(scf::WhileOp forOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
 } // namespace
 
 /// Helper function to replaces SCF op outputs with SPIR-V variable loads.
 /// We create VariableOp to handle the results value of the control flow region.
 /// spv.mlir.loop/spv.mlir.selection currently don't yield value. Right after
 /// the loop we load the value from the allocation and use it as the SCF op
 /// result.
 template <typename ScfOp, typename OpTy>
for (Type convertedType : returnTypes) {
     allocas.push_back(alloc);
     rewriter.setInsertionPointAfter(newOp);
     Value loadResult = rewriter.create<spirv::LoadOp>(loc, alloc);
     resultValue.push_back(loadResult);
   }
   rewriter.replaceOp(scfOp, resultValue);
 }
 
+static Region::iterator getBlockIt(Region &region, unsigned index) {
+  return std::next(region.begin(), index);
+}
+
 //===----------------------------------------------------------------------===//
 // scf::ForOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ForOpConversion::matchAndRewrite(scf::ForOp forOp, OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const {
   // scf::ForOp can be lowered to the structured control flow represented by
   // spirv::LoopOp by making the continue block of the spirv::LoopOp the loop
   // latch and the merge block the exit block. The resulting spirv::LoopOp has a
   // single back edge from the continue to header block, and a single exit from
   // header to merge.
   auto loc = forOp.getLoc();
   auto loopOp = rewriter.create<spirv::LoopOp>(loc, spirv::LoopControl::None);
   loopOp.addEntryAndMergeBlock();
 
   OpBuilder::InsertionGuard guard(rewriter);
   // Create the block for the header.
   auto *header = new Block();
   // Insert the header.
-  loopOp.body().getBlocks().insert(std::next(loopOp.body().begin(), 1), header);
+  loopOp.body().getBlocks().insert(getBlockIt(loopOp.body(), 1), header);
 
   // Create the new induction variable to use.
   BlockArgument newIndVar = header->addArgument(adaptor.lowerBound().getType());
   for (Value arg : adaptor.initArgs())
     header->addArgument(arg.getType());
   Block *body = forOp.getBody();
 
   // Apply signature conversion to the body of the forOp. It has a single block,
   // with argument which is the induction variable. That has to be replaced with
   // the new induction variable.
   TypeConverter::SignatureConversion signatureConverter(
       body->getNumArguments());
   signatureConverter.remapInput(0, newIndVar);
   for (unsigned i = 1, e = body->getNumArguments(); i < e; i++)
     signatureConverter.remapInput(i, header->getArgument(i));
   body = rewriter.applySignatureConversion(&forOp.getLoopBody(),
                                            signatureConverter);
 
   // Move the blocks from the forOp into the loopOp. This is the body of the
   // loopOp.
   rewriter.inlineRegionBefore(forOp->getRegion(0), loopOp.body(),
-                              std::next(loopOp.body().begin(), 2));
+                              getBlockIt(loopOp.body(), 2));
 
   SmallVector<Value, 8> args(1, adaptor.lowerBound());
   args.append(adaptor.initArgs().begin(), adaptor.initArgs().end());
   // Branch into it from the entry.
   rewriter.setInsertionPointToEnd(&(loopOp.body().front()));
   rewriter.create<spirv::BranchOp>(loc, header, args);
 
   // Generate the rest of the loop header.

 LogicalResult TerminatorOpConversion::matchAndRewrite(
     scf::YieldOp terminatorOp, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   ValueRange operands = adaptor.getOperands();
 
   // If the region is return values, store each value into the associated
   // VariableOp created during lowering of the parent region.
   if (!operands.empty()) {
-    auto loc = terminatorOp.getLoc();
     auto &allocas = scfToSPIRVContext->outputVars[terminatorOp->getParentOp()];
-    assert(allocas.size() == operands.size());
+    if (allocas.size() != operands.size())
+      return failure();
+
+    auto loc = terminatorOp.getLoc();
     for (unsigned i = 0, e = operands.size(); i < e; i++)
       rewriter.create<spirv::StoreOp>(loc, allocas[i], operands[i]);
     if (isa<spirv::LoopOp>(terminatorOp->getParentOp())) {
       // For loops we also need to update the branch jumping back to the header.
       auto br =
           cast<spirv::BranchOp>(rewriter.getInsertionBlock()->getTerminator());
       SmallVector<Value, 8> args(br.getBlockArguments());
       args.append(operands.begin(), operands.end());
       rewriter.setInsertionPoint(br);
       rewriter.create<spirv::BranchOp>(terminatorOp.getLoc(), br.getTarget(),
                                        args);
       rewriter.eraseOp(br);
     }
   }
   rewriter.eraseOp(terminatorOp);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// scf::WhileOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+WhileOpConversion::matchAndRewrite(scf::WhileOp whileOp, OpAdaptor adaptor,
+                                   ConversionPatternRewriter &rewriter) const {
+  auto loc = whileOp.getLoc();
+  auto loopOp = rewriter.create<spirv::LoopOp>(loc, spirv::LoopControl::None);
+  loopOp.addEntryAndMergeBlock();
+
+  OpBuilder::InsertionGuard guard(rewriter);
+
+  Region &beforeRegion = whileOp.before();
+  Region &afterRegion = whileOp.after();
+
+  Block &entryBlock = *loopOp.getEntryBlock();
+  Block &beforeBlock = beforeRegion.front();
+  Block &afterBlock = afterRegion.front();
+  Block &mergeBlock = *loopOp.getMergeBlock();
+
+  auto cond = cast<scf::ConditionOp>(beforeBlock.getTerminator());
+  SmallVector<Value> condArgs;
+  if (failed(rewriter.getRemappedValues(cond.args(), condArgs)))
+    return failure();
+
+  Value conditionVal = rewriter.getRemappedValue(cond.condition());
+  if (!conditionVal)
+    return failure();
+
+  auto yield = cast<scf::YieldOp>(afterBlock.getTerminator());
+  SmallVector<Value> yieldArgs;
+  if (failed(rewriter.getRemappedValues(yield.results(), yieldArgs)))
+    return failure();
+
+  // Move the while before block as the initial loop header block.
+  rewriter.inlineRegionBefore(beforeRegion, loopOp.body(),
+                              getBlockIt(loopOp.body(), 1));
+
+  // Move the while after block as the initial loop body block.
+  rewriter.inlineRegionBefore(afterRegion, loopOp.body(),
+                              getBlockIt(loopOp.body(), 2));
+
+  // Jump from the loop entry block to the loop header block.
+  rewriter.setInsertionPointToEnd(&entryBlock);
+  rewriter.create<spirv::BranchOp>(loc, &beforeBlock, adaptor.inits());
+
+  auto condLoc = cond.getLoc();
+
+  SmallVector<Value> resultValues(condArgs.size());
+
+  // For other SCF ops, the scf.yield op yields the value for the whole SCF op.
+  // So we use the scf.yield op as the anchor to create/load/store SPIR-V local
+  // variables. But for the scf.while op, the scf.yield op yields a value for
+  // the before region, which may not matching the whole op's result. Instead,
+  // the scf.condition op returns values matching the whole op's results. So we
+  // need to create/load/store variables according to that.
+  for (auto it : llvm::enumerate(condArgs)) {
+    auto res = it.value();
+    auto i = it.index();
+    auto pointerType =
+        spirv::PointerType::get(res.getType(), spirv::StorageClass::Function);
+
+    // Create local variables before the scf.while op.
+    rewriter.setInsertionPoint(loopOp);
+    auto alloc = rewriter.create<spirv::VariableOp>(
+        condLoc, pointerType, spirv::StorageClass::Function,
+        /*initializer=*/nullptr);
+
+    // Load the final result values after the scf.while op.
+    rewriter.setInsertionPointAfter(loopOp);
+    auto loadResult = rewriter.create<spirv::LoadOp>(condLoc, alloc);
+    resultValues[i] = loadResult;
+
+    // Store the current iteration's result value.
+    rewriter.setInsertionPointToEnd(&beforeBlock);
+    rewriter.create<spirv::StoreOp>(condLoc, alloc, res);
+  }
+
+  rewriter.setInsertionPointToEnd(&beforeBlock);
+  rewriter.replaceOpWithNewOp<spirv::BranchConditionalOp>(
+      cond, conditionVal, &afterBlock, condArgs, &mergeBlock, llvm::None);
+
+  // Convert the scf.yield op to a branch back to the header block.
+  rewriter.setInsertionPointToEnd(&afterBlock);
+  rewriter.replaceOpWithNewOp<spirv::BranchOp>(yield, &beforeBlock, yieldArgs);
+
+  rewriter.replaceOp(whileOp, resultValues);
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // Hooks
 //===----------------------------------------------------------------------===//
 
 void mlir::populateSCFToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                       ScfToSPIRVContext &scfToSPIRVContext,
                                       RewritePatternSet &patterns) {
-  patterns.add<ForOpConversion, IfOpConversion, TerminatorOpConversion>(
-      patterns.getContext(), typeConverter, scfToSPIRVContext.getImpl());
+  patterns.add<ForOpConversion, IfOpConversion, TerminatorOpConversion,
+               WhileOpConversion>(patterns.getContext(), typeConverter,
+                                  scfToSPIRVContext.getImpl());
 }

mlir/test/Conversion/SCFToSPIRV/while.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -convert-scf-to-spirv %s -o - | FileCheck %s
+
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader, Int64], [SPV_KHR_storage_buffer_storage_class]>, {}>
+} {
+
+// CHECK-LABEL: @while_loop1
+func @while_loop1(%arg0: i32, %arg1: i32) -> i32 {
+  // CHECK-SAME: (%[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
+  // CHECK: %[[INITVAR:.*]] = spv.Constant 2 : i32
+  // CHECK: %[[VAR1:.*]] = spv.Variable : !spv.ptr<i32, Function>
+  // CHECK: spv.mlir.loop {
+  // CHECK:   spv.Branch ^[[HEADER:.*]](%[[ARG1]] : i32)
+  // CHECK: ^[[HEADER]](%[[INDVAR1:.*]]: i32):
+  // CHECK:   %[[CMP:.*]] = spv.SLessThan %[[INDVAR1]], %[[ARG2]] : i32
+  // CHECK:   spv.Store "Function" %[[VAR1]], %[[INDVAR1]] : i32
+  // CHECK:   spv.BranchConditional %[[CMP]], ^[[BODY:.*]](%[[INDVAR1]] : i32), ^[[MERGE:.*]]
+  // CHECK: ^[[BODY]](%[[INDVAR2:.*]]: i32):
+  // CHECK:   %[[UPDATED:.*]] = spv.IMul %[[INDVAR2]], %[[INITVAR]] : i32
+  // CHECK: spv.Branch ^[[HEADER]](%[[UPDATED]] : i32)
+  // CHECK: ^[[MERGE]]:
+  // CHECK:   spv.mlir.merge
+  // CHECK: }
+  %c2_i32 = arith.constant 2 : i32
+  %0 = scf.while (%arg3 = %arg0) : (i32) -> (i32) {
+    %1 = arith.cmpi slt, %arg3, %arg1 : i32
+    scf.condition(%1) %arg3 : i32
+  } do {
+  ^bb0(%arg5: i32):
+    %1 = arith.muli %arg5, %c2_i32 : i32
+    scf.yield %1 : i32
+  }
+  // CHECK: %[[OUT:.*]] = spv.Load "Function" %[[VAR1]] : i32
+  // CHECK: spv.ReturnValue %[[OUT]] : i32
+  return %0 : i32
+}
+
+// -----
+
+// CHECK-LABEL: @while_loop2
+func @while_loop2(%arg0: f32) -> i64 {
+  // CHECK-SAME: (%[[ARG:.*]]: f32)
+  // CHECK: %[[VAR:.*]] = spv.Variable : !spv.ptr<i64, Function>
+  // CHECK: spv.mlir.loop {
+  // CHECK:   spv.Branch ^[[HEADER:.*]](%[[ARG]] : f32)
+  // CHECK: ^[[HEADER]](%[[INDVAR1:.*]]: f32):
+  // CHECK:   %[[SHARED:.*]] = "foo.shared_compute"(%[[INDVAR1]]) : (f32) -> i64
+  // CHECK:   %[[CMP:.*]] = "foo.evaluate_condition"(%[[INDVAR1]], %[[SHARED]]) : (f32, i64) -> i1
+  // CHECK:   spv.Store "Function" %[[VAR]], %[[SHARED]] : i64
+  // CHECK:   spv.BranchConditional %[[CMP]], ^[[BODY:.*]](%[[SHARED]] : i64), ^[[MERGE:.*]]
+  // CHECK: ^[[BODY]](%[[INDVAR2:.*]]: i64):
+  // CHECK:   %[[UPDATED:.*]] = "foo.payload"(%[[INDVAR2]]) : (i64) -> f32
+  // CHECK: spv.Branch ^[[HEADER]](%[[UPDATED]] : f32)
+  // CHECK: ^[[MERGE]]:
+  // CHECK:   spv.mlir.merge
+  // CHECK: }
+  %res = scf.while (%arg1 = %arg0) : (f32) -> i64 {
+    %shared = "foo.shared_compute"(%arg1) : (f32) -> i64
+    %condition = "foo.evaluate_condition"(%arg1, %shared) : (f32, i64) -> i1
+    scf.condition(%condition) %shared : i64
+  } do {
+  ^bb0(%arg2: i64):
+    %res = "foo.payload"(%arg2) : (i64) -> f32
+    scf.yield %res : f32
+  }
+  // CHECK: %[[OUT:.*]] = spv.Load "Function" %[[VAR]] : i64
+  // CHECK: spv.ReturnValue %[[OUT]] : i64
+  return %res : i64
+}
+
+} // end module