FoldUtils.h

Go to the documentation of this file.

//===- FoldUtils.h - Operation Fold Utilities -------------------*- 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 header file declares various operation folding utilities. These
// utilities are intended to be used by passes to unify and simply their logic.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_TRANSFORMS_FOLDUTILS_H
#define MLIR_TRANSFORMS_FOLDUTILS_H

#include "mlir/IR/Builders.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/DialectInterface.h"

namespace mlir {
class Operation;
class Value;

//===--------------------------------------------------------------------===//
// Operation Folding Interface
//===--------------------------------------------------------------------===//

class OpFolderDialectInterface
    : public DialectInterface::Base<OpFolderDialectInterface> {
public:
  OpFolderDialectInterface(Dialect *dialect) : Base(dialect) {}

  virtual bool shouldMaterializeInto(Region *region) const { return false; }
};

//===--------------------------------------------------------------------===//
// OperationFolder
//===--------------------------------------------------------------------===//

class OperationFolder {
public:
  OperationFolder(MLIRContext *ctx) : interfaces(ctx) {}

  LogicalResult
  tryToFold(Operation *op,
            function_ref<void(Operation *)> processGeneratedConstants = nullptr,
            function_ref<void(Operation *)> preReplaceAction = nullptr);

  void notifyRemoval(Operation *op);

  template <typename OpTy, typename... Args>
  void create(OpBuilder &builder, SmallVectorImpl<Value> &results,
              Location location, Args &&... args) {
    Operation *op = builder.create<OpTy>(location, std::forward<Args>(args)...);
    if (failed(tryToFold(op, results)))
      results.assign(op->result_begin(), op->result_end());
    else if (op->getNumResults() != 0)
      op->erase();
  }

  template <typename OpTy, typename... Args>
  typename std::enable_if<OpTy::template hasTrait<OpTrait::OneResult>(),
                          Value>::type
  create(OpBuilder &builder, Location location, Args &&... args) {
    SmallVector<Value, 1> results;
    create<OpTy>(builder, results, location, std::forward<Args>(args)...);
    return results.front();
  }

  template <typename OpTy, typename... Args>
  typename std::enable_if<OpTy::template hasTrait<OpTrait::ZeroResult>(),
                          OpTy>::type
  create(OpBuilder &builder, Location location, Args &&... args) {
    auto op = builder.create<OpTy>(location, std::forward<Args>(args)...);
    SmallVector<Value, 0> unused;
    (void)tryToFold(op.getOperation(), unused);

    // Folding cannot remove a zero-result operation, so for convenience we
    // continue to return it.
    return op;
  }

private:
  using ConstantMap =
      DenseMap<std::tuple<Dialect *, Attribute, Type>, Operation *>;

  LogicalResult tryToFold(
      Operation *op, SmallVectorImpl<Value> &results,
      function_ref<void(Operation *)> processGeneratedConstants = nullptr);

  Operation *tryGetOrCreateConstant(ConstantMap &uniquedConstants,
                                    Dialect *dialect, OpBuilder &builder,
                                    Attribute value, Type type, Location loc);

  DenseMap<Region *, ConstantMap> foldScopes;

  DenseMap<Operation *, SmallVector<Dialect *, 2>> referencedDialects;

  DialectInterfaceCollection<OpFolderDialectInterface> interfaces;
};

} // end namespace mlir

#endif // MLIR_TRANSFORMS_FOLDUTILS_H