[mlir][spirv] Add support for lowering scf.for scf/if with return values

[ThomasRaoux]

Summary

This allow lowering to support scf.for and scf.if with results. As right now spv region operations don't have return value the results are demoted to Function memory. We create one allocation per result right before the region and store the yield values in it. Then we can load back the value from allocation to be able to use the result.

[antiagainst]

Cool, thanks!

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

45

This comment needs to be updated. It would be nice to give a high-level description of what we are doing here regarding yield handling.

58

llvm::to_vector<4>(operands) will work here?

142

s/IfOp/scf ops/

143

There is no spv.if. :)

143

s/doesn't have result/doesn't yield value/

144

s/ForOp/scf op/

146

s/Replace/replace/

160

I think you can directly pass in spirv::StorageClass::Function here. We should have a build method for that.

mlir/test/Conversion/GPUToSPIRV/if.mlir

134

Let's add a TODO here to mention this is missing proper VariablePointer capability.

[ThomasRaoux]

I realized my previous patch breaks seriazliation/deserialization since those require merge block to only have merge op in it. I rebased my change on top of D82914 and use the scf context to keep track of the VariableOp created during lowering of the control flow region. This allows adding store when lowering yield operation instead and solve the problem with merge block.
Sorry for the mess this cause in the review process.

[ThomasRaoux]

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

58

This still needs to be merge with br.getBlockArguments() so I'm not sure how using to_vector helps?

160

I couldn't get it to work. Maybe I'm missing something, I couldn't find any example in the code either.

[rriddle]

mlir/include/mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h

26 ↗

(On Diff #274651)

Can you please comment on what this class is, and its intended purpose?

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

23 ↗

(On Diff #274651)

Please document what this is and what it is used for.

25 ↗

(On Diff #274651)

nit: Move these out of the mlir namespace and into the global namespace for the file.

32 ↗

(On Diff #274651)

nit: ///

75 ↗

(On Diff #274651)

nit: Please use /// for top-level comments.

80 ↗

(On Diff #274651)

Static methods should not be within namespaces. Please split this out of the anonymous namespace. Anonymous namespaces should not be used for functions.

85 ↗

(On Diff #274651)

nit: Please spell out auto here.

88 ↗

(On Diff #274651)

nit: Please spell out auto here.

97 ↗

(On Diff #274651)

nit: Please add a parameter comment here, /*...=*/nullptr

134 ↗

(On Diff #274651)

nit: Spell out auto here.

144 ↗

(On Diff #274651)

nit: Please cache the end iterator of for loops, and drop the int part of unsigned int.

for (unsigned i = 1, e = ...; i != e; ++i)

249 ↗

(On Diff #274651)

nit: Drop trivial braces, and cache the end iterator.

[antiagainst]

Nice! Thanks for this! I just have a few more nits. Please address River's comments.

mlir/include/mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h

26 ↗

(On Diff #274651)

+1

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

58

Oh sorry misread. :)

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

75 ↗

(On Diff #274651)

"Replaces SCF op outputs with SPIR-V variable loads."

78 ↗

(On Diff #274651)

s/srf/SCF/

243 ↗

(On Diff #274651)

If the region returns values, store ...

[antiagainst]

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

160

I think we have a build methods that is:

static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type pointer, ::mlir::spirv::StorageClass storage_class, /*optional*/::mlir::Value initializer);

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

 public:
   using SPIRVOpLowering<scf::IfOp>::SPIRVOpLowering;
 
   LogicalResult
   matchAndRewrite(scf::IfOp IfOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 /// Pattern to erase a scf::YieldOp.
 class TerminatorOpConversion final : public SPIRVOpLowering<scf::YieldOp> {
 public:
   using SPIRVOpLowering<scf::YieldOp>::SPIRVOpLowering;
 
   LogicalResult
   matchAndRewrite(scf::YieldOp terminatorOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
+    // Update the branch to merge block to pass the arguments that will escape
+    // the region.
+    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();
   }
 };
 
 /// Pattern lowering GPU block/thread size/id to loading SPIR-V invocation
 /// builtin variables.
 template <typename SourceOp, spirv::BuiltIn builtin>
public:
 
   LogicalResult
   matchAndRewrite(gpu::ReturnOp returnOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 } // namespace
 
+// Helper function to handle the results of scf operations.
+// We store values returned by the IfOp in an allocation as
+// spv.loop/spv.if/spv.selection currently doesn't have result. Right after the loop we
+// load the value from the allocation and use it as the ForOp result.
+template <typename ScfOp, typename OpTy>
+static void ReplaceSCFOutputValue(ScfOp scfOp, OpTy newOp,
+                                  SPIRVTypeConverter &typeConverter,
+                                  ConversionPatternRewriter &rewriter,
+                                  Block *mergeBlock) {
+
+  auto loc = scfOp.getLoc();
+  SmallVector<Value, 8> resultValue;
+  for (auto result : scfOp.results()) {
+    auto convertedType = typeConverter.convertType(result.getType());
+    auto pointerType =
+        spirv::PointerType::get(convertedType, spirv::StorageClass::Function);
+    rewriter.setInsertionPoint(newOp);
+    auto alloc = rewriter.create<spirv::VariableOp>(
+        loc, pointerType,
+        rewriter.getI32IntegerAttr(
+            static_cast<int32_t>(spirv::StorageClass::Function)),
+        nullptr);
+    rewriter.setInsertionPointAfter(newOp);
+    Value loadResult = rewriter.create<spirv::LoadOp>(loc, alloc);
+    resultValue.push_back(loadResult);
+    rewriter.setInsertionPoint(mergeBlock, mergeBlock->begin());
+    auto arg = mergeBlock->addArgument(convertedType);
+    rewriter.create<spirv::StoreOp>(loc, alloc, arg);
+  }
+  rewriter.replaceOp(scfOp, resultValue);
+}
+
 //===----------------------------------------------------------------------===//
 // scf::ForOp.
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ForOpConversion::matchAndRewrite(scf::ForOp forOp, ArrayRef<Value> operands,
                                  ConversionPatternRewriter &rewriter) const {
   // scf::ForOp can be lowered to the structured control flow represented by
ForOpConversion::matchAndRewrite(scf::ForOp forOp, ArrayRef<Value> operands,
   // Create the block for the header.
   auto header = new Block();
   // Insert the header.
   loopOp.body().getBlocks().insert(std::next(loopOp.body().begin(), 1), header);
 
   // Create the new induction variable to use.
   BlockArgument newIndVar =
       header->addArgument(forOperands.lowerBound().getType());
+  for (auto arg : forOperands.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 int i = 1; i < body->getNumArguments(); i++)
+    signatureConverter.remapInput(i, header->getArgument(i));
   body = rewriter.applySignatureConversion(&forOp.getLoopBody(),
                                            signatureConverter);
 
-  // Delete the loop terminator.
-  rewriter.eraseOp(body->getTerminator());
-
   // Move the blocks from the forOp into the loopOp. This is the body of the
   // loopOp.
   rewriter.inlineRegionBefore(forOp.getOperation()->getRegion(0), loopOp.body(),
                               std::next(loopOp.body().begin(), 2));
 
+  SmallVector<Value, 8> args(1, forOperands.lowerBound());
+  args.append(forOperands.initArgs().begin(), forOperands.initArgs().end());
   // Branch into it from the entry.
   rewriter.setInsertionPointToEnd(&(loopOp.body().front()));
-  rewriter.create<spirv::BranchOp>(loc, header, forOperands.lowerBound());
+  rewriter.create<spirv::BranchOp>(loc, header, args);
 
   // Generate the rest of the loop header.
   rewriter.setInsertionPointToEnd(header);
   auto mergeBlock = loopOp.getMergeBlock();
   auto cmpOp = rewriter.create<spirv::SLessThanOp>(
       loc, rewriter.getI1Type(), newIndVar, forOperands.upperBound());
+  args.clear();
+  // Forward all the arguments to the merge block except the induction variable.
+  for (unsigned i = 1; i < header->getNumArguments(); i++)
+    args.push_back(header->getArgument(i));
+
   rewriter.create<spirv::BranchConditionalOp>(
-      loc, cmpOp, body, ArrayRef<Value>(), mergeBlock, ArrayRef<Value>());
+      loc, cmpOp, body, ArrayRef<Value>(), mergeBlock, args);
 
   // Generate instructions to increment the step of the induction variable and
   // branch to the header.
   Block *continueBlock = loopOp.getContinueBlock();
   rewriter.setInsertionPointToEnd(continueBlock);
 
   // Add the step to the induction variable and branch to the header.
   Value updatedIndVar = rewriter.create<spirv::IAddOp>(
       loc, newIndVar.getType(), newIndVar, forOperands.step());
   rewriter.create<spirv::BranchOp>(loc, header, updatedIndVar);
 
-  rewriter.eraseOp(forOp);
+  ReplaceSCFOutputValue(forOp, loopOp, typeConverter, rewriter, mergeBlock);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // scf::IfOp.
 //===----------------------------------------------------------------------===//
 
 LogicalResult
IfOpConversion::matchAndRewrite(scf::IfOp ifOp, ArrayRef<Value> operands,
   }
 
   // Create a `spv.BranchConditional` operation for selection header block.
   rewriter.setInsertionPointToEnd(selectionHeaderBlock);
   rewriter.create<spirv::BranchConditionalOp>(loc, ifOperands.condition(),
                                               thenBlock, ArrayRef<Value>(),
                                               elseBlock, ArrayRef<Value>());
 
-  rewriter.eraseOp(ifOp);
+  ReplaceSCFOutputValue(ifOp, selectionOp, typeConverter, rewriter, mergeBlock);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Builtins.
 //===----------------------------------------------------------------------===//
 
 static Optional<int32_t> getLaunchConfigIndex(Operation *op) {

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

 static Type getElementType(Type type, Attribute indices, OpAsmParser &parser,
                            llvm::SMLoc loc) {
   auto errorFn = [&](StringRef err) -> InFlightDiagnostic {
     return parser.emitError(loc, err);
   };
   return getElementType(type, indices, errorFn);
 }
 
-/// Returns true if the given `block` only contains one `spv._merge` op.
+/// Returns true if the given `block` ends with `spv._merge` op.
 static inline bool isMergeBlock(Block &block) {
-  return !block.empty() && std::next(block.begin()) == block.end() &&
-         isa<spirv::MergeOp>(block.front());
+  return !block.empty() && isa<spirv::MergeOp>(block.back());
 }
 
 //===----------------------------------------------------------------------===//
 // Common parsers and printers
 //===----------------------------------------------------------------------===//
 
 // Parses an atomic update op. If the update op does not take a value (like
 // AtomicIIncrement) `hasValue` must be false.

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

 Type PointerType::getPointeeType() const { return getImpl()->pointeeType; }
 
 StorageClass PointerType::getStorageClass() const {
   return getImpl()->getStorageClass();
 }
 
 void PointerType::getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                                 Optional<StorageClass> storage) {
-  if (storage)
-    assert(*storage == getStorageClass() && "inconsistent storage class!");
-
   // Use this pointer type's storage class because this pointer indicates we are
   // using the pointee type in that specific storage class.
   getPointeeType().cast<SPIRVType>().getExtensions(extensions,
                                                    getStorageClass());
 
   if (auto scExts = spirv::getExtensions(getStorageClass()))
     extensions.push_back(*scExts);
 }
 
 void PointerType::getCapabilities(
     SPIRVType::CapabilityArrayRefVector &capabilities,
     Optional<StorageClass> storage) {
-  if (storage)
-    assert(*storage == getStorageClass() && "inconsistent storage class!");
-
   // Use this pointer type's storage class because this pointer indicates we are
   // using the pointee type in that specific storage class.
   getPointeeType().cast<SPIRVType>().getCapabilities(capabilities,
                                                      getStorageClass());
 
   if (auto scCaps = spirv::getCapabilities(getStorageClass()))
     capabilities.push_back(*scCaps);
 }

mlir/test/Conversion/GPUToSPIRV/if.mlir

attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
           store %value, %arg4[%j] : memref<10xf32>
         } else {
           %value = load %arg4[%j] : memref<10xf32>
           store %value, %arg3[%j] : memref<10xf32>
         }
       }
       gpu.return
     }
+
+    // CHECK-LABEL: @simple_if_yield
+    gpu.func @simple_if_yield(%arg2 : memref<10xf32>, %arg3 : i1) kernel
+    attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
+      // CHECK: %[[VAR1:.*]] = spv.Variable : !spv.ptr<f32, Function>
+      // CHECK: %[[VAR2:.*]] = spv.Variable : !spv.ptr<f32, Function>
+      // CHECK:       spv.selection {
+      // CHECK-NEXT:    spv.BranchConditional {{%.*}}, [[TRUE:\^.*]], [[FALSE:\^.*]]
+      // CHECK-NEXT:  [[TRUE]]:
+      // CHECK:         %[[RET1TRUE:.*]] = spv.constant 0.000000e+00 : f32
+      // CHECK:         %[[RET2TRUE:.*]] = spv.constant 1.000000e+00 : f32
+      // CHECK:         spv.Branch ^[[MERGE:.*]](%[[RET1TRUE]], %[[RET2TRUE]] : f32, f32)
+      // CHECK-NEXT:  [[FALSE]]:
+      // CHECK:         %[[RET2FALSE:.*]] = spv.constant 2.000000e+00 : f32
+      // CHECK:         %[[RET1FALSE:.*]] = spv.constant 3.000000e+00 : f32
+      // CHECK:         spv.Branch ^[[MERGE]](%[[RET1FALSE]], %[[RET2FALSE]] : f32, f32)
+      // CHECK-NEXT:  ^[[MERGE]](%[[RET1:.*]]: f32, %[[RET2:.*]]: f32):
+      // CHECK-DAG:     spv.Store "Function" %[[VAR1]], %[[RET1]] : f32
+      // CHECK-DAG:     spv.Store "Function" %[[VAR2]], %[[RET2]] : f32
+      // CHECK:         spv._merge
+      // CHECK-NEXT:  }
+      // CHECK-DAG:   %[[OUT1:.*]] = spv.Load "Function" %[[VAR1]] : f32
+      // CHECK-DAG:   %[[OUT2:.*]] = spv.Load "Function" %[[VAR2]] : f32
+      // CHECK:       spv.Store "StorageBuffer" {{%.*}}, %[[OUT1]] : f32
+      // CHECK:       spv.Store "StorageBuffer" {{%.*}}, %[[OUT2]] : f32
+      // CHECK:       spv.Return
+      %0:2 = scf.if %arg3 -> (f32, f32) {
+        %c0 = constant 0.0 : f32
+        %c1 = constant 1.0 : f32
+        scf.yield %c0, %c1 : f32, f32
+      } else {
+        %c0 = constant 2.0 : f32
+        %c1 = constant 3.0 : f32
+        scf.yield %c1, %c0 : f32, f32
+      }
+      %i = constant 0 : index
+      %j = constant 1 : index
+      store %0#0, %arg2[%i] : memref<10xf32>
+      store %0#1, %arg2[%j] : memref<10xf32>
+      gpu.return
+    }
+
+    // CHECK-LABEL: @simple_if_yield_type_change
+    // CHECK:       %[[VAR:.*]] = spv.Variable : !spv.ptr<!spv.ptr<!spv.struct<!spv.array<10 x f32, stride=4> [0]>, StorageBuffer>, Function>
+    // CHECK:       spv.selection {
+    // CHECK-NEXT:    spv.BranchConditional {{%.*}}, [[TRUE:\^.*]], [[FALSE:\^.*]]
+    // CHECK-NEXT:  [[TRUE]]:
+    // CHECK:         spv.Branch ^[[MERGE:.*]]({{%.*}} : !spv.ptr<!spv.struct<!spv.array<10 x f32, stride=4> [0]>, StorageBuffer>)
+    // CHECK-NEXT:  [[FALSE]]:
+    // CHECK:         spv.Branch ^[[MERGE]]({{%.*}} : !spv.ptr<!spv.struct<!spv.array<10 x f32, stride=4> [0]>, StorageBuffer>)
+    // CHECK-NEXT:  ^[[MERGE]](%[[RET:.*]]: !spv.ptr<!spv.struct<!spv.array<10 x f32, stride=4> [0]>, StorageBuffer>):
+    // CHECK:         spv.Store "Function" %[[VAR]], %[[RET]] : !spv.ptr<!spv.struct<!spv.array<10 x f32, stride=4> [0]>, StorageBuffer>
+    // CHECK:         spv._merge
+    // CHECK-NEXT:  }
+    // CHECK:       %[[OUT:.*]] = spv.Load "Function" %[[VAR]] : !spv.ptr<!spv.struct<!spv.array<10 x f32, stride=4> [0]>, StorageBuffer>
+    // CHECK:       %[[ADD:.*]] = spv.AccessChain %[[OUT]][{{%.*}}, {{%.*}}] : !spv.ptr<!spv.struct<!spv.array<10 x f32, stride=4> [0]>, StorageBuffer>
+    // CHECK:       spv.Store "StorageBuffer" %[[ADD]], {{%.*}} : f32
+    // CHECK:       spv.Return
+    gpu.func @simple_if_yield_type_change(%arg2 : memref<10xf32>, %arg3 : memref<10xf32>, %arg4 : i1) kernel
+    attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
+      %i = constant 0 : index
+      %value = constant 0.0 : f32
+
+      %0 = scf.if %arg4 -> (memref<10xf32>) {
+        scf.yield %arg2 : memref<10xf32>
+      } else {
+        scf.yield %arg3 : memref<10xf32>
+      }
+      store %value, %0[%i] : memref<10xf32>
+      gpu.return
+    }
   }
 }

mlir/test/Conversion/GPUToSPIRV/loop.mlir

-// RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s
+// RUN: mlir-opt -convert-gpu-to-spirv -split-input-file %s -o - | FileCheck %s
 
 module attributes {
   gpu.container_module,
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
     {max_compute_workgroup_invocations = 128 : i32,
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
 } {
attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
       // CHECK:        spv._merge
       // CHECK:      }
       scf.for %arg4 = %lb to %ub step %step {
         %1 = load %arg2[%arg4] : memref<10xf32>
         store %1, %arg3[%arg4] : memref<10xf32>
       }
       gpu.return
     }
+
+
+    // CHECK-LABEL: @loop_yield
+    gpu.func @loop_yield(%arg2 : memref<10xf32>, %arg3 : memref<10xf32>) kernel
+    attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
+      // CHECK: %[[LB:.*]] = spv.constant 4 : i32
+      %lb = constant 4 : index
+      // CHECK: %[[UB:.*]] = spv.constant 42 : i32
+      %ub = constant 42 : index
+      // CHECK: %[[STEP:.*]] = spv.constant 2 : i32
+      %step = constant 2 : index
+      // CHECK: %[[INITVAR1:.*]] = spv.constant 0.000000e+00 : f32
+      %s0 = constant 0.0 : f32
+      // CHECK: %[[INITVAR2:.*]] = spv.constant 1.000000e+00 : f32
+      %s1 = constant 1.0 : f32
+      // CHECK: %[[VAR1:.*]] = spv.Variable : !spv.ptr<f32, Function>
+      // CHECK: %[[VAR2:.*]] = spv.Variable : !spv.ptr<f32, Function>
+      // CHECK: spv.loop {
+      // CHECK:   spv.Branch ^[[HEADER:.*]](%[[LB]], %[[INITVAR1]], %[[INITVAR2]] : i32, f32, f32)
+      // CHECK: ^[[HEADER]](%[[INDVAR:.*]]: i32, %[[CARRIED1:.*]]: f32, %[[CARRIED2:.*]]: f32):
+      // CHECK:   %[[CMP:.*]] = spv.SLessThan %[[INDVAR]], %[[UB]] : i32
+      // CHECK:   spv.BranchConditional %[[CMP]], ^[[BODY:.*]], ^[[MERGE:.*]](%[[CARRIED1]], %[[CARRIED2]] : f32, f32)
+      // CHECK: ^[[BODY]]:
+      // CHECK:   %[[UPDATED:.*]] = spv.FAdd %[[CARRIED1]], %[[CARRIED1]] : f32
+      // CHECK:   %[[INCREMENT:.*]] = spv.IAdd %[[INDVAR]], %[[STEP]] : i32
+      // CHECK: spv.Branch ^[[HEADER]](%[[INCREMENT]], %[[UPDATED]], %[[UPDATED]] : i32, f32, f32)
+      // CHECK: ^[[MERGE]](%[[OUT1:.*]]: f32, %[[OUT2:.*]]: f32)
+      // CHECK-DAG:   spv.Store "Function" %[[VAR1]], %[[OUT1]] : f32
+      // CHECK-DAG:   spv.Store "Function" %[[VAR2]], %[[OUT2]] : f32
+      // CHECK:   spv._merge
+      // CHECK: }
+      %result:2 = scf.for %i0 = %lb to %ub step %step iter_args(%si = %s0, %sj = %s1) -> (f32, f32) {
+        %sn = addf %si, %si : f32
+        scf.yield %sn, %sn : f32, f32
+      }
+      // CHECK-DAG: %[[OUT1:.*]] = spv.Load "Function" %[[VAR1]] : f32
+      // CHECK-DAG: %[[OUT2:.*]] = spv.Load "Function" %[[VAR2]] : f32
+      // CHECK: spv.Store "StorageBuffer" {{%.*}}, %[[OUT1]] : f32
+      // CHECK: spv.Store "StorageBuffer" {{%.*}}, %[[OUT2]] : f32
+      store %result#0, %arg3[%lb] : memref<10xf32>
+      store %result#1, %arg3[%ub] : memref<10xf32>
+      gpu.return
+    }
   }
 }