From: Mahesh Ravishankar Date: Tue, 12 Nov 2019 19:32:54 +0000 (-0800) Subject: Add Conversion to lower loop::ForOp to spirv::LoopOp. X-Git-Tag: llvmorg-11-init~1466^2~356 X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=104af84f4c8788da2d2202cac720fd7aa62b01cd;p=platform%2Fupstream%2Fllvm.git Add Conversion to lower loop::ForOp to spirv::LoopOp. 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. PiperOrigin-RevId: 280015614 --- diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp index 4f73b9b..ece4269 100644 --- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp +++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp @@ -21,6 +21,7 @@ //===----------------------------------------------------------------------===// #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" @@ -29,6 +30,16 @@ using namespace mlir; namespace { +/// Pattern to convert a loop::ForOp within kernel functions into spirv::LoopOp. +class ForOpConversion final : public SPIRVOpLowering { +public: + using SPIRVOpLowering::SPIRVOpLowering; + + PatternMatchResult + matchAndRewrite(Operation *op, ArrayRef operands, + ConversionPatternRewriter &rewriter) const override; +}; + /// Pattern lowering GPU block/thread size/id to loading SPIR-V invocation /// builin variables. template @@ -51,8 +62,79 @@ public: matchAndRewrite(Operation *op, ArrayRef operands, ConversionPatternRewriter &rewriter) const override; }; + } // namespace +PatternMatchResult +ForOpConversion::matchAndRewrite(Operation *op, ArrayRef 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(op); + loop::ForOpOperandAdaptor forOperands(operands); + auto loc = op->getLoc(); + auto loopControl = rewriter.getI32IntegerAttr( + static_cast(spirv::LoopControl::None)); + auto loopOp = rewriter.create(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(loc, header, forOperands.lowerBound()); + + // Generate the rest of the loop header. + rewriter.setInsertionPointToEnd(header); + auto mergeBlock = loopOp.getMergeBlock(); + auto cmpOp = rewriter.create( + loc, rewriter.getI1Type(), newIndVar, forOperands.upperBound()); + rewriter.create( + loc, cmpOp, body, ArrayRef(), mergeBlock, ArrayRef()); + + // 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( + loc, newIndVar->getType(), newIndVar, forOperands.step()); + rewriter.create(loc, header, updatedIndVar); + + rewriter.eraseOp(forOp); + return matchSuccess(); +} + template PatternMatchResult LaunchConfigConversion::matchAndRewrite( Operation *op, ArrayRef operands, @@ -148,7 +230,7 @@ void GPUToSPIRVPass::runOnModule() { SPIRVTypeConverter typeConverter(&basicTypeConverter); OwningRewritePatternList patterns; patterns.insert< - KernelFnConversion, + ForOpConversion, KernelFnConversion, LaunchConfigConversion, LaunchConfigConversion, LaunchConfigConversion, diff --git a/mlir/test/Conversion/GPUToSPIRV/loop.mlir b/mlir/test/Conversion/GPUToSPIRV/loop.mlir new file mode 100644 index 0000000..870f4e3 --- /dev/null +++ b/mlir/test/Conversion/GPUToSPIRV/loop.mlir @@ -0,0 +1,39 @@ +// 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