DialectConversion.h

Go to the documentation of this file.

//===- DialectConversion.h - MLIR dialect conversion pass -------*- C++ -*-===//
//
// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file declares a generic pass for converting between MLIR dialects.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_TRANSFORMS_DIALECTCONVERSION_H_
#define MLIR_TRANSFORMS_DIALECTCONVERSION_H_

#include "mlir/IR/PatternMatch.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/StringMap.h"

namespace mlir {

// Forward declarations.
class Block;
class ConversionPatternRewriter;
class FuncOp;
class MLIRContext;
class Operation;
class Type;
class Value;

//===----------------------------------------------------------------------===//
// Type Conversion
//===----------------------------------------------------------------------===//

class TypeConverter {
public:
  virtual ~TypeConverter() = default;

  class SignatureConversion {
  public:
    SignatureConversion(unsigned numOrigInputs)
        : remappedInputs(numOrigInputs) {}

    struct InputMapping {
      size_t inputNo, size;
      Value replacementValue;
    };

    ArrayRef<Type> getConvertedTypes() const { return argTypes; }

    Optional<InputMapping> getInputMapping(unsigned input) const {
      return remappedInputs[input];
    }

    //===------------------------------------------------------------------===//
    // Conversion Hooks
    //===------------------------------------------------------------------===//

    void addInputs(unsigned origInputNo, ArrayRef<Type> types);

    void addInputs(ArrayRef<Type> types);

    void remapInput(unsigned origInputNo, unsigned newInputNo,
                    unsigned newInputCount = 1);

    void remapInput(unsigned origInputNo, Value replacement);

  private:
    SmallVector<Optional<InputMapping>, 4> remappedInputs;

    SmallVector<Type, 4> argTypes;
  };

  virtual LogicalResult convertType(Type t, SmallVectorImpl<Type> &results);

  virtual Type convertType(Type t) { return t; }

  LogicalResult convertTypes(ArrayRef<Type> types,
                             SmallVectorImpl<Type> &results);

  bool isLegal(Type type);

  bool isSignatureLegal(FunctionType funcType);

  virtual LogicalResult convertSignatureArg(unsigned inputNo, Type type,
                                            SignatureConversion &result);

  Optional<SignatureConversion> convertBlockSignature(Block *block);

  virtual Operation *materializeConversion(PatternRewriter &rewriter,
                                           Type resultType,
                                           ArrayRef<Value> inputs,
                                           Location loc) {
    llvm_unreachable("expected 'materializeConversion' to be overridden");
  }
};

//===----------------------------------------------------------------------===//
// Conversion Patterns
//===----------------------------------------------------------------------===//

class ConversionPattern : public RewritePattern {
public:
  ConversionPattern(StringRef rootName, PatternBenefit benefit,
                    MLIRContext *ctx)
      : RewritePattern(rootName, benefit, ctx) {}

  virtual void rewrite(Operation *op, ArrayRef<Value> operands,
                       ConversionPatternRewriter &rewriter) const {
    llvm_unreachable("unimplemented rewrite");
  }

  virtual void rewrite(Operation *op, ArrayRef<Value> properOperands,
                       ArrayRef<Block *> destinations,
                       ArrayRef<ArrayRef<Value>> operands,
                       ConversionPatternRewriter &rewriter) const {
    llvm_unreachable("unimplemented rewrite for terminators");
  }

  virtual PatternMatchResult
  matchAndRewrite(Operation *op, ArrayRef<Value> properOperands,
                  ArrayRef<Block *> destinations,
                  ArrayRef<ArrayRef<Value>> operands,
                  ConversionPatternRewriter &rewriter) const {
    if (!match(op))
      return matchFailure();
    rewrite(op, properOperands, destinations, operands, rewriter);
    return matchSuccess();
  }

  virtual PatternMatchResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const {
    if (!match(op))
      return matchFailure();
    rewrite(op, operands, rewriter);
    return matchSuccess();
  }

  PatternMatchResult matchAndRewrite(Operation *op,
                                     PatternRewriter &rewriter) const final;

private:
  using RewritePattern::rewrite;
};

template <typename SourceOp>
struct OpConversionPattern : public ConversionPattern {
  OpConversionPattern(MLIRContext *context, PatternBenefit benefit = 1)
      : ConversionPattern(SourceOp::getOperationName(), benefit, context) {}

  void rewrite(Operation *op, ArrayRef<Value> operands,
               ConversionPatternRewriter &rewriter) const final {
    rewrite(cast<SourceOp>(op), operands, rewriter);
  }
  void rewrite(Operation *op, ArrayRef<Value> properOperands,
               ArrayRef<Block *> destinations,
               ArrayRef<ArrayRef<Value>> operands,
               ConversionPatternRewriter &rewriter) const final {
    rewrite(cast<SourceOp>(op), properOperands, destinations, operands,
            rewriter);
  }
  PatternMatchResult
  matchAndRewrite(Operation *op, ArrayRef<Value> properOperands,
                  ArrayRef<Block *> destinations,
                  ArrayRef<ArrayRef<Value>> operands,
                  ConversionPatternRewriter &rewriter) const final {
    return matchAndRewrite(cast<SourceOp>(op), properOperands, destinations,
                           operands, rewriter);
  }
  PatternMatchResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const final {
    return matchAndRewrite(cast<SourceOp>(op), operands, rewriter);
  }

  // TODO(b/142763075): Use OperandAdaptor when it supports access to unnamed
  // operands.

  virtual void rewrite(SourceOp op, ArrayRef<Value> operands,
                       ConversionPatternRewriter &rewriter) const {
    llvm_unreachable("must override matchAndRewrite or a rewrite method");
  }

  virtual void rewrite(SourceOp op, ArrayRef<Value> properOperands,
                       ArrayRef<Block *> destinations,
                       ArrayRef<ArrayRef<Value>> operands,
                       ConversionPatternRewriter &rewriter) const {
    llvm_unreachable("unimplemented rewrite for terminators");
  }

  virtual PatternMatchResult
  matchAndRewrite(SourceOp op, ArrayRef<Value> properOperands,
                  ArrayRef<Block *> destinations,
                  ArrayRef<ArrayRef<Value>> operands,
                  ConversionPatternRewriter &rewriter) const {
    if (!match(op))
      return matchFailure();
    rewrite(op, properOperands, destinations, operands, rewriter);
    return matchSuccess();
  }

  virtual PatternMatchResult
  matchAndRewrite(SourceOp op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const {
    if (!match(op))
      return matchFailure();
    rewrite(op, operands, rewriter);
    return matchSuccess();
  }

private:
  using ConversionPattern::matchAndRewrite;
};

void populateFuncOpTypeConversionPattern(OwningRewritePatternList &patterns,
                                         MLIRContext *ctx,
                                         TypeConverter &converter);

//===----------------------------------------------------------------------===//
// Conversion PatternRewriter
//===----------------------------------------------------------------------===//

namespace detail {
struct ConversionPatternRewriterImpl;
} // end namespace detail

class ConversionPatternRewriter final : public PatternRewriter {
public:
  ConversionPatternRewriter(MLIRContext *ctx, TypeConverter *converter);
  ~ConversionPatternRewriter() override;

  Block *
  applySignatureConversion(Region *region,
                           TypeConverter::SignatureConversion &conversion);

  void replaceUsesOfBlockArgument(BlockArgument from, Value to);

  Value getRemappedValue(Value key);

  //===--------------------------------------------------------------------===//
  // PatternRewriter Hooks
  //===--------------------------------------------------------------------===//

  void replaceOp(Operation *op, ValueRange newValues,
                 ValueRange valuesToRemoveIfDead) override;
  using PatternRewriter::replaceOp;

  void eraseOp(Operation *op) override;

  Block *splitBlock(Block *block, Block::iterator before) override;

  void mergeBlocks(Block *source, Block *dest, ValueRange argValues) override;

  void inlineRegionBefore(Region &region, Region &parent,
                          Region::iterator before) override;
  using PatternRewriter::inlineRegionBefore;

  void cloneRegionBefore(Region &region, Region &parent,
                         Region::iterator before,
                         BlockAndValueMapping &mapping) override;
  using PatternRewriter::cloneRegionBefore;

  Operation *insert(Operation *op) override;

  void startRootUpdate(Operation *op) override;

  void finalizeRootUpdate(Operation *op) override;

  void cancelRootUpdate(Operation *op) override;

  detail::ConversionPatternRewriterImpl &getImpl();

private:
  std::unique_ptr<detail::ConversionPatternRewriterImpl> impl;
};

//===----------------------------------------------------------------------===//
// ConversionTarget
//===----------------------------------------------------------------------===//

class ConversionTarget {
public:
  enum class LegalizationAction {
    Legal,

    Dynamic,

    Illegal,
  };

  struct LegalOpDetails {
    bool isRecursivelyLegal = false;
  };

  using DynamicLegalityCallbackFn = std::function<bool(Operation *)>;

  ConversionTarget(MLIRContext &ctx) : ctx(ctx) {}
  virtual ~ConversionTarget() = default;

  //===--------------------------------------------------------------------===//
  // Legality Registration
  //===--------------------------------------------------------------------===//

  void setOpAction(OperationName op, LegalizationAction action);
  template <typename OpT> void setOpAction(LegalizationAction action) {
    setOpAction(OperationName(OpT::getOperationName(), &ctx), action);
  }

  template <typename OpT> void addLegalOp() {
    setOpAction<OpT>(LegalizationAction::Legal);
  }
  template <typename OpT, typename OpT2, typename... OpTs> void addLegalOp() {
    addLegalOp<OpT>();
    addLegalOp<OpT2, OpTs...>();
  }

  template <typename OpT> void addDynamicallyLegalOp() {
    setOpAction<OpT>(LegalizationAction::Dynamic);
  }
  template <typename OpT, typename OpT2, typename... OpTs>
  void addDynamicallyLegalOp() {
    addDynamicallyLegalOp<OpT>();
    addDynamicallyLegalOp<OpT2, OpTs...>();
  }

  template <typename OpT>
  void addDynamicallyLegalOp(const DynamicLegalityCallbackFn &callback) {
    OperationName opName(OpT::getOperationName(), &ctx);
    setOpAction(opName, LegalizationAction::Dynamic);
    setLegalityCallback(opName, callback);
  }
  template <typename OpT, typename OpT2, typename... OpTs>
  void addDynamicallyLegalOp(const DynamicLegalityCallbackFn &callback) {
    addDynamicallyLegalOp<OpT>(callback);
    addDynamicallyLegalOp<OpT2, OpTs...>(callback);
  }
  template <typename OpT, class Callable>
  typename std::enable_if<!is_invocable<Callable, Operation *>::value>::type
  addDynamicallyLegalOp(Callable &&callback) {
    addDynamicallyLegalOp<OpT>(
        [=](Operation *op) { return callback(cast<OpT>(op)); });
  }

  template <typename OpT> void addIllegalOp() {
    setOpAction<OpT>(LegalizationAction::Illegal);
  }
  template <typename OpT, typename OpT2, typename... OpTs> void addIllegalOp() {
    addIllegalOp<OpT>();
    addIllegalOp<OpT2, OpTs...>();
  }

  template <typename OpT>
  void markOpRecursivelyLegal(const DynamicLegalityCallbackFn &callback = {}) {
    OperationName opName(OpT::getOperationName(), &ctx);
    markOpRecursivelyLegal(opName, callback);
  }
  template <typename OpT, typename OpT2, typename... OpTs>
  void markOpRecursivelyLegal(const DynamicLegalityCallbackFn &callback = {}) {
    markOpRecursivelyLegal<OpT>(callback);
    markOpRecursivelyLegal<OpT2, OpTs...>(callback);
  }
  template <typename OpT, class Callable>
  typename std::enable_if<!is_invocable<Callable, Operation *>::value>::type
  markOpRecursivelyLegal(Callable &&callback) {
    markOpRecursivelyLegal<OpT>(
        [=](Operation *op) { return callback(cast<OpT>(op)); });
  }

  void setDialectAction(ArrayRef<StringRef> dialectNames,
                        LegalizationAction action);

  template <typename... Names>
  void addLegalDialect(StringRef name, Names... names) {
    SmallVector<StringRef, 2> dialectNames({name, names...});
    setDialectAction(dialectNames, LegalizationAction::Legal);
  }
  template <typename... Args> void addLegalDialect() {
    SmallVector<StringRef, 2> dialectNames({Args::getDialectNamespace()...});
    setDialectAction(dialectNames, LegalizationAction::Legal);
  }

  template <typename... Names>
  void addDynamicallyLegalDialect(StringRef name, Names... names) {
    SmallVector<StringRef, 2> dialectNames({name, names...});
    setDialectAction(dialectNames, LegalizationAction::Dynamic);
  }
  template <typename... Args>
  void addDynamicallyLegalDialect(
      Optional<DynamicLegalityCallbackFn> callback = llvm::None) {
    SmallVector<StringRef, 2> dialectNames({Args::getDialectNamespace()...});
    setDialectAction(dialectNames, LegalizationAction::Dynamic);
    if (callback)
      setLegalityCallback(dialectNames, *callback);
  }

  template <typename... Names>
  void addIllegalDialect(StringRef name, Names... names) {
    SmallVector<StringRef, 2> dialectNames({name, names...});
    setDialectAction(dialectNames, LegalizationAction::Illegal);
  }
  template <typename... Args> void addIllegalDialect() {
    SmallVector<StringRef, 2> dialectNames({Args::getDialectNamespace()...});
    setDialectAction(dialectNames, LegalizationAction::Illegal);
  }

  //===--------------------------------------------------------------------===//
  // Legality Querying
  //===--------------------------------------------------------------------===//

  Optional<LegalizationAction> getOpAction(OperationName op) const;

  Optional<LegalOpDetails> isLegal(Operation *op) const;

protected:
  virtual bool isDynamicallyLegal(Operation *op) const {
    llvm_unreachable(
        "targets with custom legalization must override 'isDynamicallyLegal'");
  }

private:
  void setLegalityCallback(OperationName name,
                           const DynamicLegalityCallbackFn &callback);

  void setLegalityCallback(ArrayRef<StringRef> dialects,
                           const DynamicLegalityCallbackFn &callback);

  void markOpRecursivelyLegal(OperationName name,
                              const DynamicLegalityCallbackFn &callback);

  struct LegalizationInfo {
    LegalizationAction action;

    bool isRecursivelyLegal;
  };

  Optional<LegalizationInfo> getOpInfo(OperationName op) const;

  llvm::MapVector<OperationName, LegalizationInfo> legalOperations;

  DenseMap<OperationName, DynamicLegalityCallbackFn> opLegalityFns;

  DenseMap<OperationName, DynamicLegalityCallbackFn> opRecursiveLegalityFns;

  llvm::StringMap<LegalizationAction> legalDialects;

  llvm::StringMap<DynamicLegalityCallbackFn> dialectLegalityFns;

  MLIRContext &ctx;
};

//===----------------------------------------------------------------------===//
// Op Conversion Entry Points
//===----------------------------------------------------------------------===//


LLVM_NODISCARD LogicalResult
applyPartialConversion(ArrayRef<Operation *> ops, ConversionTarget &target,
                       const OwningRewritePatternList &patterns,
                       TypeConverter *converter = nullptr);
LLVM_NODISCARD LogicalResult
applyPartialConversion(Operation *op, ConversionTarget &target,
                       const OwningRewritePatternList &patterns,
                       TypeConverter *converter = nullptr);

LLVM_NODISCARD LogicalResult
applyFullConversion(ArrayRef<Operation *> ops, ConversionTarget &target,
                    const OwningRewritePatternList &patterns,
                    TypeConverter *converter = nullptr);
LLVM_NODISCARD LogicalResult
applyFullConversion(Operation *op, ConversionTarget &target,
                    const OwningRewritePatternList &patterns,
                    TypeConverter *converter = nullptr);

LLVM_NODISCARD LogicalResult applyAnalysisConversion(
    ArrayRef<Operation *> ops, ConversionTarget &target,
    const OwningRewritePatternList &patterns,
    DenseSet<Operation *> &convertedOps, TypeConverter *converter = nullptr);
LLVM_NODISCARD LogicalResult applyAnalysisConversion(
    Operation *op, ConversionTarget &target,
    const OwningRewritePatternList &patterns,
    DenseSet<Operation *> &convertedOps, TypeConverter *converter = nullptr);
} // end namespace mlir

#endif // MLIR_TRANSFORMS_DIALECTCONVERSION_H_