mlir/lib/Conversion/FuncToSPIRV/FuncToSPIRV.cpp

   1 //===- FuncToSPIRV.cpp - Func 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 Func dialect to SPIR-V dialect.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #include "mlir/Conversion/FuncToSPIRV/FuncToSPIRV.h"
  14 #include "../SPIRVCommon/Pattern.h"
  15 #include "mlir/Dialect/Func/IR/FuncOps.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 "llvm/Support/Debug.h"
  22
  23 #define DEBUG_TYPE "func-to-spirv-pattern"
  24
  25 using namespace mlir;
  26
  27 //===----------------------------------------------------------------------===//
  28 // Operation conversion
  29 //===----------------------------------------------------------------------===//
  30
  31 // Note that DRR cannot be used for the patterns in this file: we may need to
  32 // convert type along the way, which requires ConversionPattern. DRR generates
  33 // normal RewritePattern.
  34
  35 namespace {
  36
  37 /// Converts func.return to spirv.Return.
  38 class ReturnOpPattern final : public OpConversionPattern<func::ReturnOp> {
  39 public:
  40   using OpConversionPattern<func::ReturnOp>::OpConversionPattern;
  41
  42   LogicalResult
  43   matchAndRewrite(func::ReturnOp returnOp, OpAdaptor adaptor,
  44                   ConversionPatternRewriter &rewriter) const override {
  45     if (returnOp.getNumOperands() > 1)
  46       return failure();
  47
  48     if (returnOp.getNumOperands() == 1) {
  49       rewriter.replaceOpWithNewOp<spirv::ReturnValueOp>(
  50           returnOp, adaptor.getOperands()[0]);
  51     } else {
  52       rewriter.replaceOpWithNewOp<spirv::ReturnOp>(returnOp);
  53     }
  54     return success();
  55   }
  56 };
  57
  58 /// Converts func.call to spirv.FunctionCall.
  59 class CallOpPattern final : public OpConversionPattern<func::CallOp> {
  60 public:
  61   using OpConversionPattern<func::CallOp>::OpConversionPattern;
  62
  63   LogicalResult
  64   matchAndRewrite(func::CallOp callOp, OpAdaptor adaptor,
  65                   ConversionPatternRewriter &rewriter) const override {
  66     // multiple results func was not converted to spirv.func
  67     if (callOp.getNumResults() > 1)
  68       return failure();
  69     if (callOp.getNumResults() == 1) {
  70       auto resultType =
  71           getTypeConverter()->convertType(callOp.getResult(0).getType());
  72       if (!resultType)
  73         return failure();
  74       rewriter.replaceOpWithNewOp<spirv::FunctionCallOp>(
  75           callOp, resultType, adaptor.getOperands(), callOp->getAttrs());
  76     } else {
  77       rewriter.replaceOpWithNewOp<spirv::FunctionCallOp>(
  78           callOp, TypeRange(), adaptor.getOperands(), callOp->getAttrs());
  79     }
  80     return success();
  81   }
  82 };
  83
  84 } // namespace
  85
  86 //===----------------------------------------------------------------------===//
  87 // Pattern population
  88 //===----------------------------------------------------------------------===//
  89
  90 void mlir::populateFuncToSPIRVPatterns(const SPIRVTypeConverter &typeConverter,
  91                                        RewritePatternSet &patterns) {
  92   MLIRContext *context = patterns.getContext();
  93
  94   patterns.add<ReturnOpPattern, CallOpPattern>(typeConverter, context);
  95 }