1 //===- ControlFlowToSPIRV.cpp - ControlFlow to SPIR-V Patterns ------------===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This file implements patterns to convert standard dialect to SPIR-V dialect. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #include "mlir/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRV.h" 14 #include "../SPIRVCommon/Pattern.h" 15 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" 16 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h" 17 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h" 18 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h" 19 #include "mlir/Dialect/SPIRV/Utils/LayoutUtils.h" 20 #include "mlir/IR/AffineMap.h" 21 #include "mlir/IR/PatternMatch.h" 22 #include "mlir/Transforms/DialectConversion.h" 23 #include "llvm/Support/Debug.h" 24 #include "llvm/Support/FormatVariadic.h" 25 26 #define DEBUG_TYPE "cf-to-spirv-pattern" 27 28 using namespace mlir; 29 30 /// Legailze target block arguments. 31 static LogicalResult legalizeBlockArguments(Block &block, Operation *op, 32 PatternRewriter &rewriter, 33 const TypeConverter &converter) { 34 auto builder = OpBuilder::atBlockBegin(&block); 35 for (unsigned i = 0; i < block.getNumArguments(); ++i) { 36 BlockArgument arg = block.getArgument(i); 37 if (converter.isLegal(arg.getType())) 38 continue; 39 Type ty = arg.getType(); 40 Type newTy = converter.convertType(ty); 41 if (!newTy) { 42 return rewriter.notifyMatchFailure( 43 op, llvm::formatv("failed to legalize type for argument {0})", arg)); 44 } 45 unsigned argNum = arg.getArgNumber(); 46 Location loc = arg.getLoc(); 47 Value newArg = block.insertArgument(argNum, newTy, loc); 48 Value convertedValue = converter.materializeSourceConversion( 49 builder, op->getLoc(), ty, newArg); 50 if (!convertedValue) { 51 return rewriter.notifyMatchFailure( 52 op, llvm::formatv("failed to cast new argument {0} to type {1})", 53 newArg, ty)); 54 } 55 arg.replaceAllUsesWith(convertedValue); 56 block.eraseArgument(argNum + 1); 57 } 58 return success(); 59 } 60 61 //===----------------------------------------------------------------------===// 62 // Operation conversion 63 //===----------------------------------------------------------------------===// 64 65 namespace { 66 /// Converts cf.br to spirv.Branch. 67 struct BranchOpPattern final : OpConversionPattern<cf::BranchOp> { 68 using OpConversionPattern::OpConversionPattern; 69 70 LogicalResult 71 matchAndRewrite(cf::BranchOp op, OpAdaptor adaptor, 72 ConversionPatternRewriter &rewriter) const override { 73 if (failed(legalizeBlockArguments(*op.getDest(), op, rewriter, 74 *getTypeConverter()))) 75 return failure(); 76 77 rewriter.replaceOpWithNewOp<spirv::BranchOp>(op, op.getDest(), 78 adaptor.getDestOperands()); 79 return success(); 80 } 81 }; 82 83 /// Converts cf.cond_br to spirv.BranchConditional. 84 struct CondBranchOpPattern final : OpConversionPattern<cf::CondBranchOp> { 85 using OpConversionPattern::OpConversionPattern; 86 87 LogicalResult 88 matchAndRewrite(cf::CondBranchOp op, OpAdaptor adaptor, 89 ConversionPatternRewriter &rewriter) const override { 90 if (failed(legalizeBlockArguments(*op.getTrueDest(), op, rewriter, 91 *getTypeConverter()))) 92 return failure(); 93 94 if (failed(legalizeBlockArguments(*op.getFalseDest(), op, rewriter, 95 *getTypeConverter()))) 96 return failure(); 97 98 rewriter.replaceOpWithNewOp<spirv::BranchConditionalOp>( 99 op, adaptor.getCondition(), op.getTrueDest(), 100 adaptor.getTrueDestOperands(), op.getFalseDest(), 101 adaptor.getFalseDestOperands()); 102 return success(); 103 } 104 }; 105 } // namespace 106 107 //===----------------------------------------------------------------------===// 108 // Pattern population 109 //===----------------------------------------------------------------------===// 110 111 void mlir::cf::populateControlFlowToSPIRVPatterns( 112 const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns) { 113 MLIRContext *context = patterns.getContext(); 114 115 patterns.add<BranchOpPattern, CondBranchOpPattern>(typeConverter, context); 116 } 117