Builders.h

Go to the documentation of this file.

//===- Builders.h - Helpers for constructing MLIR Classes -------*- 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
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_IR_BUILDERS_H
#define MLIR_IR_BUILDERS_H

#include "mlir/IR/OpDefinition.h"

namespace mlir {

class AffineExpr;
class BlockAndValueMapping;
class ModuleOp;
class UnknownLoc;
class FileLineColLoc;
class Type;
class PrimitiveType;
class IntegerType;
class FunctionType;
class MemRefType;
class VectorType;
class RankedTensorType;
class UnrankedTensorType;
class TupleType;
class NoneType;
class BoolAttr;
class IntegerAttr;
class FloatAttr;
class StringAttr;
class TypeAttr;
class ArrayAttr;
class SymbolRefAttr;
class ElementsAttr;
class DenseElementsAttr;
class DenseIntElementsAttr;
class AffineMapAttr;
class AffineMap;
class UnitAttr;

class Builder {
public:
  explicit Builder(MLIRContext *context) : context(context) {}
  explicit Builder(ModuleOp module);

  MLIRContext *getContext() const { return context; }

  Identifier getIdentifier(StringRef str);

  // Locations.
  Location getUnknownLoc();
  Location getFileLineColLoc(Identifier filename, unsigned line,
                             unsigned column);
  Location getFusedLoc(ArrayRef<Location> locs,
                       Attribute metadata = Attribute());

  // Types.
  FloatType getBF16Type();
  FloatType getF16Type();
  FloatType getF32Type();
  FloatType getF64Type();

  IndexType getIndexType();

  IntegerType getI1Type();
  IntegerType getIntegerType(unsigned width);
  FunctionType getFunctionType(ArrayRef<Type> inputs, ArrayRef<Type> results);
  TupleType getTupleType(ArrayRef<Type> elementTypes);
  NoneType getNoneType();

  template <typename Ty, typename... Args> Ty getType(Args... args) {
    return Ty::get(context, args...);
  }

  // Attributes.
  NamedAttribute getNamedAttr(StringRef name, Attribute val);

  UnitAttr getUnitAttr();
  BoolAttr getBoolAttr(bool value);
  DictionaryAttr getDictionaryAttr(ArrayRef<NamedAttribute> value);
  IntegerAttr getIntegerAttr(Type type, int64_t value);
  IntegerAttr getIntegerAttr(Type type, const APInt &value);
  FloatAttr getFloatAttr(Type type, double value);
  FloatAttr getFloatAttr(Type type, const APFloat &value);
  StringAttr getStringAttr(StringRef bytes);
  ArrayAttr getArrayAttr(ArrayRef<Attribute> value);
  FlatSymbolRefAttr getSymbolRefAttr(Operation *value);
  FlatSymbolRefAttr getSymbolRefAttr(StringRef value);
  SymbolRefAttr getSymbolRefAttr(StringRef value,
                                 ArrayRef<FlatSymbolRefAttr> nestedReferences);

  // Returns a 0-valued attribute of the given `type`. This function only
  // supports boolean, integer, and 16-/32-/64-bit float types, and vector or
  // ranked tensor of them. Returns null attribute otherwise.
  Attribute getZeroAttr(Type type);

  // Convenience methods for fixed types.
  FloatAttr getF16FloatAttr(float value);
  FloatAttr getF32FloatAttr(float value);
  FloatAttr getF64FloatAttr(double value);

  IntegerAttr getI8IntegerAttr(int8_t value);
  IntegerAttr getI16IntegerAttr(int16_t value);
  IntegerAttr getI32IntegerAttr(int32_t value);
  IntegerAttr getI64IntegerAttr(int64_t value);

  DenseIntElementsAttr getI32VectorAttr(ArrayRef<int32_t> values);

  ArrayAttr getAffineMapArrayAttr(ArrayRef<AffineMap> values);
  ArrayAttr getI32ArrayAttr(ArrayRef<int32_t> values);
  ArrayAttr getI64ArrayAttr(ArrayRef<int64_t> values);
  ArrayAttr getIndexArrayAttr(ArrayRef<int64_t> values);
  ArrayAttr getF32ArrayAttr(ArrayRef<float> values);
  ArrayAttr getF64ArrayAttr(ArrayRef<double> values);
  ArrayAttr getStrArrayAttr(ArrayRef<StringRef> values);

  // Affine expressions and affine maps.
  AffineExpr getAffineDimExpr(unsigned position);
  AffineExpr getAffineSymbolExpr(unsigned position);
  AffineExpr getAffineConstantExpr(int64_t constant);

  // Special cases of affine maps and integer sets
  AffineMap getEmptyAffineMap();
  AffineMap getConstantAffineMap(int64_t val);
  // One dimension id identity map: (i) -> (i).
  AffineMap getDimIdentityMap();
  // Multi-dimensional identity map: (d0, d1, d2) -> (d0, d1, d2).
  AffineMap getMultiDimIdentityMap(unsigned rank);
  // One symbol identity map: ()[s] -> (s).
  AffineMap getSymbolIdentityMap();

  AffineMap getSingleDimShiftAffineMap(int64_t shift);

  AffineMap getShiftedAffineMap(AffineMap map, int64_t shift);

protected:
  MLIRContext *context;
};

class OpBuilder : public Builder {
public:
  explicit OpBuilder(MLIRContext *ctx) : Builder(ctx) {}

  explicit OpBuilder(Region *region) : Builder(region->getContext()) {
    if (!region->empty())
      setInsertionPoint(&region->front(), region->front().begin());
  }
  explicit OpBuilder(Region &region) : OpBuilder(&region) {}

  virtual ~OpBuilder();

  explicit OpBuilder(Operation *op) : Builder(op->getContext()) {
    setInsertionPoint(op);
  }

  explicit OpBuilder(Block *block) : OpBuilder(block, block->end()) {}

  OpBuilder(Block *block, Block::iterator insertPoint)
      : OpBuilder(block->getParent()) {
    setInsertionPoint(block, insertPoint);
  }

  class InsertPoint {
  public:
    InsertPoint() = default;

    InsertPoint(Block *insertBlock, Block::iterator insertPt)
        : block(insertBlock), point(insertPt) {}

    bool isSet() const { return (block != nullptr); }

    Block *getBlock() const { return block; }
    Block::iterator getPoint() const { return point; }

  private:
    Block *block = nullptr;
    Block::iterator point;
  };

  class InsertionGuard {
  public:
    InsertionGuard(OpBuilder &builder)
        : builder(builder), ip(builder.saveInsertionPoint()) {}
    ~InsertionGuard() { builder.restoreInsertionPoint(ip); }

  private:
    OpBuilder &builder;
    OpBuilder::InsertPoint ip;
  };

  void clearInsertionPoint() {
    this->block = nullptr;
    insertPoint = Block::iterator();
  }

  InsertPoint saveInsertionPoint() const {
    return InsertPoint(getInsertionBlock(), getInsertionPoint());
  }

  void restoreInsertionPoint(InsertPoint ip) {
    if (ip.isSet())
      setInsertionPoint(ip.getBlock(), ip.getPoint());
    else
      clearInsertionPoint();
  }

  void setInsertionPoint(Block *block, Block::iterator insertPoint) {
    // TODO: check that insertPoint is in this rather than some other block.
    this->block = block;
    this->insertPoint = insertPoint;
  }

  void setInsertionPoint(Operation *op) {
    setInsertionPoint(op->getBlock(), Block::iterator(op));
  }

  void setInsertionPointAfter(Operation *op) {
    setInsertionPoint(op->getBlock(), ++Block::iterator(op));
  }

  void setInsertionPointToStart(Block *block) {
    setInsertionPoint(block, block->begin());
  }

  void setInsertionPointToEnd(Block *block) {
    setInsertionPoint(block, block->end());
  }

  Block *getInsertionBlock() const { return block; }

  Block::iterator getInsertionPoint() const { return insertPoint; }

  virtual Operation *insert(Operation *op);

  Block *createBlock(Region *parent, Region::iterator insertPt = {});

  Block *createBlock(Block *insertBefore);

  Block *getBlock() const { return block; }

  Operation *createOperation(const OperationState &state);

  template <typename OpTy, typename... Args>
  OpTy create(Location location, Args &&... args) {
    OperationState state(location, OpTy::getOperationName());
    OpTy::build(this, state, std::forward<Args>(args)...);
    auto *op = createOperation(state);
    auto result = dyn_cast<OpTy>(op);
    assert(result && "Builder didn't return the right type");
    return result;
  }

  template <typename OpTy, typename... Args>
  void createOrFold(SmallVectorImpl<Value> &results, Location location,
                    Args &&... args) {
    // Create the operation without using 'createOperation' as we don't want to
    // insert it yet.
    OperationState state(location, OpTy::getOperationName());
    OpTy::build(this, state, std::forward<Args>(args)...);
    Operation *op = Operation::create(state);

    // Fold the operation. If successful destroy it, otherwise insert it.
    if (succeeded(tryFold(op, results)))
      op->destroy();
    else
      insert(op);
  }

  template <typename OpTy, typename... Args>
  typename std::enable_if<OpTy::template hasTrait<OpTrait::OneResult>(),
                          Value>::type
  createOrFold(Location location, Args &&... args) {
    SmallVector<Value, 1> results;
    createOrFold<OpTy>(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
  createOrFold(Location location, Args &&... args) {
    auto op = create<OpTy>(location, std::forward<Args>(args)...);
    SmallVector<Value, 0> unused;
    tryFold(op.getOperation(), unused);

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

  LogicalResult tryFold(Operation *op, SmallVectorImpl<Value> &results);

  Operation *clone(Operation &op, BlockAndValueMapping &mapper) {
    return insert(op.clone(mapper));
  }
  Operation *clone(Operation &op) { return insert(op.clone()); }

  Operation *cloneWithoutRegions(Operation &op, BlockAndValueMapping &mapper) {
    return insert(op.cloneWithoutRegions(mapper));
  }
  Operation *cloneWithoutRegions(Operation &op) {
    return insert(op.cloneWithoutRegions());
  }
  template <typename OpT> OpT cloneWithoutRegions(OpT op) {
    return cast<OpT>(cloneWithoutRegions(*op.getOperation()));
  }

private:
  Block *block = nullptr;
  Block::iterator insertPoint;
};

} // namespace mlir

#endif