//===----------------------------------------------------------------------===//
#include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.h"
#include "mlir/Dialect/GPU/GPUDialect.h"
+#include "mlir/Dialect/LoopOps/LoopOps.h"
#include "mlir/Dialect/SPIRV/SPIRVDialect.h"
#include "mlir/Dialect/SPIRV/SPIRVOps.h"
#include "mlir/Pass/Pass.h"
namespace {
+/// Pattern to convert a loop::ForOp within kernel functions into spirv::LoopOp.
+class ForOpConversion final : public SPIRVOpLowering<loop::ForOp> {
+public:
+ using SPIRVOpLowering<loop::ForOp>::SPIRVOpLowering;
+
+ PatternMatchResult
+ matchAndRewrite(Operation *op, ArrayRef<Value *> operands,
+ ConversionPatternRewriter &rewriter) const override;
+};
+
/// Pattern lowering GPU block/thread size/id to loading SPIR-V invocation
/// builin variables.
template <typename OpTy, spirv::BuiltIn builtin>
matchAndRewrite(Operation *op, ArrayRef<Value *> operands,
ConversionPatternRewriter &rewriter) const override;
};
+
} // namespace
+PatternMatchResult
+ForOpConversion::matchAndRewrite(Operation *op, ArrayRef<Value *> operands,
+ ConversionPatternRewriter &rewriter) const {
+ // loop::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 forOp = cast<loop::ForOp>(op);
+ loop::ForOpOperandAdaptor forOperands(operands);
+ auto loc = op->getLoc();
+ auto loopControl = rewriter.getI32IntegerAttr(
+ static_cast<uint32_t>(spirv::LoopControl::None));
+ auto loopOp = rewriter.create<spirv::LoopOp>(loc, loopControl);
+ 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);
+
+ // Create the new induction variable to use.
+ BlockArgument *newIndVar =
+ header->addArgument(forOperands.lowerBound()->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);
+ rewriter.applySignatureConversion(&forOp.getOperation()->getRegion(0),
+ 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));
+
+ // Branch into it from the entry.
+ rewriter.setInsertionPointToEnd(&(loopOp.body().front()));
+ rewriter.create<spirv::BranchOp>(loc, header, forOperands.lowerBound());
+
+ // 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());
+ rewriter.create<spirv::BranchConditionalOp>(
+ loc, cmpOp, body, ArrayRef<Value *>(), mergeBlock, ArrayRef<Value *>());
+
+ // 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);
+ return matchSuccess();
+}
+
template <typename OpTy, spirv::BuiltIn builtin>
PatternMatchResult LaunchConfigConversion<OpTy, builtin>::matchAndRewrite(
Operation *op, ArrayRef<Value *> operands,
SPIRVTypeConverter typeConverter(&basicTypeConverter);
OwningRewritePatternList patterns;
patterns.insert<
- KernelFnConversion,
+ ForOpConversion, KernelFnConversion,
LaunchConfigConversion<gpu::BlockDimOp, spirv::BuiltIn::WorkgroupSize>,
LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
--- /dev/null
+// RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s
+
+module attributes {gpu.container_module} {
+ func @loop(%arg0 : memref<10xf32>, %arg1 : memref<10xf32>) {
+ %c0 = constant 1 : index
+ "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0, %arg0, %arg1) { kernel = "loop_kernel", kernel_module = @kernels} : (index, index, index, index, index, index, memref<10xf32>, memref<10xf32>) -> ()
+ return
+ }
+
+ module @kernels attributes {gpu.kernel_module} {
+ func @loop_kernel(%arg2 : memref<10xf32>, %arg3 : memref<10xf32>)
+ attributes {gpu.kernel} {
+ // 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: spv.loop {
+ // CHECK-NEXT: spv.Branch [[HEADER:\^.*]]([[LB]] : i32)
+ // CHECK: [[HEADER]]([[INDVAR:%.*]]: i32):
+ // CHECK: [[CMP:%.*]] = spv.SLessThan [[INDVAR]], [[UB]] : i32
+ // CHECK: spv.BranchConditional [[CMP]], [[BODY:\^.*]], [[MERGE:\^.*]]
+ // CHECK: [[BODY]]:
+ // CHECK: spv.AccessChain {{%.*}}{{\[}}[[INDVAR]]{{\]}} : {{.*}}
+ // CHECK: spv.AccessChain {{%.*}}{{\[}}[[INDVAR]]{{\]}} : {{.*}}
+ // CHECK: [[INCREMENT:%.*]] = spv.IAdd [[INDVAR]], [[STEP]] : i32
+ // CHECK: spv.Branch [[HEADER]]([[INCREMENT]] : i32)
+ // CHECK: [[MERGE]]
+ // CHECK: spv._merge
+ // CHECK: }
+ loop.for %arg4 = %lb to %ub step %step {
+ %1 = load %arg2[%arg4] : memref<10xf32>
+ store %1, %arg3[%arg4] : memref<10xf32>
+ }
+ return
+ }
+ }
+}
\ No newline at end of file