//===- Builders.h - Helpers for constructing MLIR Classes -------*- C++ -*-===// // // Part of the LLVM 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" #include "llvm/Support/Compiler.h" #include namespace mlir { class AffineExpr; class IRMapping; class UnknownLoc; class FileLineColLoc; class FileLineColRange; class Type; class IntegerType; class FloatType; class FunctionType; class GraphType; class IndexType; 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; /// This class is a general helper class for creating context-global objects /// like types, attributes, and affine expressions. class Builder { public: explicit Builder(MLIRContext *context) : context(context) {} explicit Builder(Operation *op) : Builder(op->getContext()) {} MLIRContext *getContext() const { return context; } // Locations. Location getUnknownLoc(); Location getFusedLoc(ArrayRef locs, Attribute metadata = Attribute()); // Types. FloatType getF8E8M0Type(); FloatType getF8E4M3FNType(); FloatType getF8E5M2Type(); FloatType getBF16Type(); FloatType getF16Type(); FloatType getTF32Type(); FloatType getF32Type(); FloatType getF64Type(); FloatType getF80Type(); FloatType getF128Type(); IndexType getIndexType(); IntegerType getI1Type(); IntegerType getI2Type(); IntegerType getI4Type(); IntegerType getI8Type(); IntegerType getI16Type(); IntegerType getI32Type(); IntegerType getI64Type(); IntegerType getIntegerType(unsigned width); IntegerType getIntegerType(unsigned width, bool isSigned); FunctionType getFunctionType(TypeRange inputs, TypeRange results); GraphType getGraphType(TypeRange inputs, TypeRange results); TupleType getTupleType(TypeRange elementTypes); NoneType getNoneType(); /// Get or construct an instance of the type `Ty` with provided arguments. template Ty getType(Args &&...args) { return Ty::get(context, std::forward(args)...); } /// Get or construct an instance of the attribute `Attr` with provided /// arguments. template Attr getAttr(Args &&...args) { return Attr::get(context, std::forward(args)...); } // Attributes. NamedAttribute getNamedAttr(StringRef name, Attribute val); UnitAttr getUnitAttr(); BoolAttr getBoolAttr(bool value); DictionaryAttr getDictionaryAttr(ArrayRef 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(const Twine &bytes); ArrayAttr getArrayAttr(ArrayRef value); // 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. TypedAttr getZeroAttr(Type type); // Returns a 1-valued attribute of the given `type`. // Type constraints are the same as `getZeroAttr`. TypedAttr getOneAttr(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); IntegerAttr getIndexAttr(int64_t value); /// Signed and unsigned integer attribute getters. IntegerAttr getSI32IntegerAttr(int32_t value); IntegerAttr getUI32IntegerAttr(uint32_t value); /// Vector-typed DenseIntElementsAttr getters. `values` must not be empty. DenseIntElementsAttr getBoolVectorAttr(ArrayRef values); DenseIntElementsAttr getI32VectorAttr(ArrayRef values); DenseIntElementsAttr getI64VectorAttr(ArrayRef values); DenseIntElementsAttr getIndexVectorAttr(ArrayRef values); DenseFPElementsAttr getF32VectorAttr(ArrayRef values); DenseFPElementsAttr getF64VectorAttr(ArrayRef values); /// Tensor-typed DenseIntElementsAttr getters. `values` can be empty. /// These are generally preferable for representing general lists of integers /// as attributes. DenseIntElementsAttr getI32TensorAttr(ArrayRef values); DenseIntElementsAttr getI64TensorAttr(ArrayRef values); DenseIntElementsAttr getIndexTensorAttr(ArrayRef values); /// Tensor-typed DenseArrayAttr getters. DenseBoolArrayAttr getDenseBoolArrayAttr(ArrayRef values); DenseI8ArrayAttr getDenseI8ArrayAttr(ArrayRef values); DenseI16ArrayAttr getDenseI16ArrayAttr(ArrayRef values); DenseI32ArrayAttr getDenseI32ArrayAttr(ArrayRef values); DenseI64ArrayAttr getDenseI64ArrayAttr(ArrayRef values); DenseF32ArrayAttr getDenseF32ArrayAttr(ArrayRef values); DenseF64ArrayAttr getDenseF64ArrayAttr(ArrayRef values); ArrayAttr getAffineMapArrayAttr(ArrayRef values); ArrayAttr getBoolArrayAttr(ArrayRef values); ArrayAttr getI32ArrayAttr(ArrayRef values); ArrayAttr getI64ArrayAttr(ArrayRef values); ArrayAttr getIndexArrayAttr(ArrayRef values); ArrayAttr getF32ArrayAttr(ArrayRef values); ArrayAttr getF64ArrayAttr(ArrayRef values); ArrayAttr getStrArrayAttr(ArrayRef values); ArrayAttr getTypeArrayAttr(TypeRange 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 /// Returns a zero result affine map with no dimensions or symbols: () -> (). AffineMap getEmptyAffineMap(); /// Returns a single constant result affine map with 0 dimensions and 0 /// symbols. One constant result: () -> (val). 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(); /// Returns a map that shifts its (single) input dimension by 'shift'. /// (d0) -> (d0 + shift) AffineMap getSingleDimShiftAffineMap(int64_t shift); /// Returns an affine map that is a translation (shift) of all result /// expressions in 'map' by 'shift'. /// Eg: input: (d0, d1)[s0] -> (d0, d1 + s0), shift = 2 /// returns: (d0, d1)[s0] -> (d0 + 2, d1 + s0 + 2) AffineMap getShiftedAffineMap(AffineMap map, int64_t shift); protected: MLIRContext *context; }; /// This class helps build Operations. Operations that are created are /// automatically inserted at an insertion point. The builder is copyable. class OpBuilder : public Builder { public: class InsertPoint; struct Listener; /// Create a builder with the given context. explicit OpBuilder(MLIRContext *ctx, Listener *listener = nullptr) : Builder(ctx), listener(listener) {} /// Create a builder and set the insertion point to the start of the region. explicit OpBuilder(Region *region, Listener *listener = nullptr) : OpBuilder(region->getContext(), listener) { if (!region->empty()) setInsertionPointToStart(®ion->front()); } explicit OpBuilder(Region ®ion, Listener *listener = nullptr) : OpBuilder(®ion, listener) {} /// Create a builder and set insertion point to the given operation, which /// will cause subsequent insertions to go right before it. explicit OpBuilder(Operation *op, Listener *listener = nullptr) : OpBuilder(op->getContext(), listener) { setInsertionPoint(op); } OpBuilder(Block *block, Block::iterator insertPoint, Listener *listener = nullptr) : OpBuilder(block->getParent()->getContext(), listener) { setInsertionPoint(block, insertPoint); } /// Create a builder and set the insertion point to before the first operation /// in the block but still inside the block. static OpBuilder atBlockBegin(Block *block, Listener *listener = nullptr) { return OpBuilder(block, block->begin(), listener); } /// Create a builder and set the insertion point to after the last operation /// in the block but still inside the block. static OpBuilder atBlockEnd(Block *block, Listener *listener = nullptr) { return OpBuilder(block, block->end(), listener); } /// Create a builder and set the insertion point to before the block /// terminator. static OpBuilder atBlockTerminator(Block *block, Listener *listener = nullptr) { auto *terminator = block->getTerminator(); assert(terminator != nullptr && "the block has no terminator"); return OpBuilder(block, Block::iterator(terminator), listener); } //===--------------------------------------------------------------------===// // Listeners //===--------------------------------------------------------------------===// /// Base class for listeners. struct ListenerBase { /// The kind of listener. enum class Kind { /// OpBuilder::Listener or user-derived class. OpBuilderListener = 0, /// RewriterBase::Listener or user-derived class. RewriterBaseListener = 1 }; Kind getKind() const { return kind; } protected: ListenerBase(Kind kind) : kind(kind) {} private: const Kind kind; }; /// This class represents a listener that may be used to hook into various /// actions within an OpBuilder. struct Listener : public ListenerBase { Listener() : ListenerBase(ListenerBase::Kind::OpBuilderListener) {} virtual ~Listener() = default; /// Notify the listener that the specified operation was inserted. /// /// * If the operation was moved, then `previous` is the previous location /// of the op. /// * If the operation was unlinked before it was inserted, then `previous` /// is empty. /// /// Note: Creating an (unlinked) op does not trigger this notification. virtual void notifyOperationInserted(Operation *op, InsertPoint previous) {} /// Notify the listener that the specified block was inserted. /// /// * If the block was moved, then `previous` and `previousIt` are the /// previous location of the block. /// * If the block was unlinked before it was inserted, then `previous` /// is "nullptr". /// /// Note: Creating an (unlinked) block does not trigger this notification. virtual void notifyBlockInserted(Block *block, Region *previous, Region::iterator previousIt) {} protected: Listener(Kind kind) : ListenerBase(kind) {} }; /// Sets the listener of this builder to the one provided. void setListener(Listener *newListener) { listener = newListener; } /// Returns the current listener of this builder, or nullptr if this builder /// doesn't have a listener. Listener *getListener() const { return listener; } //===--------------------------------------------------------------------===// // Insertion Point Management //===--------------------------------------------------------------------===// /// This class represents a saved insertion point. class InsertPoint { public: /// Creates a new insertion point which doesn't point to anything. InsertPoint() = default; /// Creates a new insertion point at the given location. InsertPoint(Block *insertBlock, Block::iterator insertPt) : block(insertBlock), point(insertPt) {} /// Returns true if this insert point is set. bool isSet() const { return (block != nullptr); } Block *getBlock() const { return block; } Block::iterator getPoint() const { return point; } private: Block *block = nullptr; Block::iterator point; }; /// RAII guard to reset the insertion point of the builder when destroyed. class InsertionGuard { public: InsertionGuard(OpBuilder &builder) : builder(&builder), ip(builder.saveInsertionPoint()) {} ~InsertionGuard() { if (builder) builder->restoreInsertionPoint(ip); } InsertionGuard(const InsertionGuard &) = delete; InsertionGuard &operator=(const InsertionGuard &) = delete; /// Implement the move constructor to clear the builder field of `other`. /// That way it does not restore the insertion point upon destruction as /// that should be done exclusively by the just constructed InsertionGuard. InsertionGuard(InsertionGuard &&other) noexcept : builder(other.builder), ip(other.ip) { other.builder = nullptr; } InsertionGuard &operator=(InsertionGuard &&other) = delete; private: OpBuilder *builder; OpBuilder::InsertPoint ip; }; /// Reset the insertion point to no location. Creating an operation without a /// set insertion point is an error, but this can still be useful when the /// current insertion point a builder refers to is being removed. void clearInsertionPoint() { this->block = nullptr; insertPoint = Block::iterator(); } /// Return a saved insertion point. InsertPoint saveInsertionPoint() const { return InsertPoint(getInsertionBlock(), getInsertionPoint()); } /// Restore the insert point to a previously saved point. void restoreInsertionPoint(InsertPoint ip) { if (ip.isSet()) setInsertionPoint(ip.getBlock(), ip.getPoint()); else clearInsertionPoint(); } /// Set the insertion point to the specified location. 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; } /// Sets the insertion point to the specified operation, which will cause /// subsequent insertions to go right before it. void setInsertionPoint(Operation *op) { setInsertionPoint(op->getBlock(), Block::iterator(op)); } /// Sets the insertion point to the node after the specified operation, which /// will cause subsequent insertions to go right after it. void setInsertionPointAfter(Operation *op) { setInsertionPoint(op->getBlock(), ++Block::iterator(op)); } /// Sets the insertion point to the node after the specified value. If value /// has a defining operation, sets the insertion point to the node after such /// defining operation. This will cause subsequent insertions to go right /// after it. Otherwise, value is a BlockArgument. Sets the insertion point to /// the start of its block. void setInsertionPointAfterValue(Value val) { if (Operation *op = val.getDefiningOp()) { setInsertionPointAfter(op); } else { auto blockArg = llvm::cast(val); setInsertionPointToStart(blockArg.getOwner()); } } /// Sets the insertion point to the start of the specified block. void setInsertionPointToStart(Block *block) { setInsertionPoint(block, block->begin()); } /// Sets the insertion point to the end of the specified block. void setInsertionPointToEnd(Block *block) { setInsertionPoint(block, block->end()); } /// Return the block the current insertion point belongs to. Note that the /// insertion point is not necessarily the end of the block. Block *getInsertionBlock() const { return block; } /// Returns the current insertion point of the builder. Block::iterator getInsertionPoint() const { return insertPoint; } /// Returns the current block of the builder. Block *getBlock() const { return block; } //===--------------------------------------------------------------------===// // Block Creation //===--------------------------------------------------------------------===// /// Add new block with 'argTypes' arguments and set the insertion point to the /// end of it. The block is inserted at the provided insertion point of /// 'parent'. `locs` contains the locations of the inserted arguments, and /// should match the size of `argTypes`. Block *createBlock(Region *parent, Region::iterator insertPt = {}, TypeRange argTypes = {}, ArrayRef locs = {}); /// Add new block with 'argTypes' arguments and set the insertion point to the /// end of it. The block is placed before 'insertBefore'. `locs` contains the /// locations of the inserted arguments, and should match the size of /// `argTypes`. Block *createBlock(Block *insertBefore, TypeRange argTypes = {}, ArrayRef locs = {}); //===--------------------------------------------------------------------===// // Operation Creation //===--------------------------------------------------------------------===// /// Insert the given operation at the current insertion point and return it. Operation *insert(Operation *op); /// Creates an operation given the fields represented as an OperationState. Operation *create(const OperationState &state); /// Creates an operation with the given fields. Operation *create(Location loc, StringAttr opName, ValueRange operands, TypeRange types = {}, ArrayRef attributes = {}, BlockRange successors = {}, MutableArrayRef> regions = {}); private: /// Helper for sanity checking preconditions for create* methods below. template RegisteredOperationName getCheckRegisteredInfo(MLIRContext *ctx) { std::optional opName = RegisteredOperationName::lookup(TypeID::get(), ctx); if (LLVM_UNLIKELY(!opName)) { llvm::report_fatal_error( "Building op `" + OpT::getOperationName() + "` but it isn't known in this MLIRContext: the dialect may not " "be loaded or this operation hasn't been added by the dialect. See " "also https://mlir.llvm.org/getting_started/Faq/" "#registered-loaded-dependent-whats-up-with-dialects-management"); } return *opName; } public: /// Create an operation of specific op type at the current insertion point. template [[deprecated("Use OpTy::create instead")]] OpTy create(Location location, Args &&...args) { OperationState state(location, getCheckRegisteredInfo(location.getContext())); OpTy::build(*this, state, std::forward(args)...); auto *op = create(state); auto result = dyn_cast(op); assert(result && "builder didn't return the right type"); return result; } /// Create an operation of specific op type at the current insertion point, /// and immediately try to fold it. This functions populates 'results' with /// the results of the operation. /// /// Note: This performs opportunistic eager folding during IR construction. /// The folders are designed to operate efficiently on canonical IR, which /// this API does not enforce. Complete folding isn't only expected in the /// context of canonicalization which intertwine folders with pattern /// rewrites until fixed-point. template void createOrFold(SmallVectorImpl &results, Location location, Args &&...args) { // Create the operation without using 'create' as we want to control when // the listener is notified. OperationState state(location, getCheckRegisteredInfo(location.getContext())); OpTy::build(*this, state, std::forward(args)...); Operation *op = Operation::create(state); if (block) block->getOperations().insert(insertPoint, op); // Attempt to fold the operation. if (succeeded(tryFold(op, results)) && !results.empty()) { // Erase the operation, if the fold removed the need for this operation. // Note: The fold already populated the results in this case. op->erase(); return; } ResultRange opResults = op->getResults(); results.assign(opResults.begin(), opResults.end()); if (block && listener) listener->notifyOperationInserted(op, /*previous=*/{}); } /// Overload to create or fold a single result operation. template std::enable_if_t(), Value> createOrFold(Location location, Args &&...args) { SmallVector results; createOrFold(results, location, std::forward(args)...); return results.front(); } /// Overload to create or fold a zero result operation. template std::enable_if_t(), OpTy> createOrFold(Location location, Args &&...args) { auto op = OpTy::create(*this, location, std::forward(args)...); SmallVector unused; (void)tryFold(op.getOperation(), unused); // Folding cannot remove a zero-result operation, so for convenience we // continue to return it. return op; } /// Attempts to fold the given operation and places new results within /// `results`. Returns success if the operation was folded, failure otherwise. /// If the fold was in-place, `results` will not be filled. Optionally, newly /// materialized constant operations can be returned to the caller. /// /// Note: This function does not erase the operation on a successful fold. LogicalResult tryFold(Operation *op, SmallVectorImpl &results, SmallVectorImpl *materializedConstants = nullptr); /// Creates a deep copy of the specified operation, remapping any operands /// that use values outside of the operation using the map that is provided /// ( leaving them alone if no entry is present). Replaces references to /// cloned sub-operations to the corresponding operation that is copied, /// and adds those mappings to the map. Operation *clone(Operation &op, IRMapping &mapper); Operation *clone(Operation &op); /// Creates a deep copy of this operation but keep the operation regions /// empty. Operands are remapped using `mapper` (if present), and `mapper` is /// updated to contain the results. Operation *cloneWithoutRegions(Operation &op, IRMapping &mapper) { return insert(op.cloneWithoutRegions(mapper)); } Operation *cloneWithoutRegions(Operation &op) { return insert(op.cloneWithoutRegions()); } template OpT cloneWithoutRegions(OpT op) { return cast(cloneWithoutRegions(*op.getOperation())); } /// Clone the blocks that belong to "region" before the given position in /// another region "parent". The two regions must be different. The caller is /// responsible for creating or updating the operation transferring flow of /// control to the region and passing it the correct block arguments. void cloneRegionBefore(Region ®ion, Region &parent, Region::iterator before, IRMapping &mapping); void cloneRegionBefore(Region ®ion, Region &parent, Region::iterator before); void cloneRegionBefore(Region ®ion, Block *before); protected: /// The optional listener for events of this builder. Listener *listener; private: /// The current block this builder is inserting into. Block *block = nullptr; /// The insertion point within the block that this builder is inserting /// before. Block::iterator insertPoint; }; /// ImplicitLocOpBuilder maintains a 'current location', allowing use of the /// create<> method without specifying the location. It is otherwise the same /// as OpBuilder. class ImplicitLocOpBuilder : public mlir::OpBuilder { public: /// OpBuilder has a bunch of convenience constructors - we support them all /// with the additional Location. template ImplicitLocOpBuilder(Location loc, T &&...operands) : OpBuilder(std::forward(operands)...), curLoc(loc) {} /// Create a builder and set the insertion point to before the first operation /// in the block but still inside the block. static ImplicitLocOpBuilder atBlockBegin(Location loc, Block *block, Listener *listener = nullptr) { return ImplicitLocOpBuilder(loc, block, block->begin(), listener); } /// Create a builder and set the insertion point to after the last operation /// in the block but still inside the block. static ImplicitLocOpBuilder atBlockEnd(Location loc, Block *block, Listener *listener = nullptr) { return ImplicitLocOpBuilder(loc, block, block->end(), listener); } /// Create a builder and set the insertion point to before the block /// terminator. static ImplicitLocOpBuilder atBlockTerminator(Location loc, Block *block, Listener *listener = nullptr) { auto *terminator = block->getTerminator(); assert(terminator != nullptr && "the block has no terminator"); return ImplicitLocOpBuilder(loc, block, Block::iterator(terminator), listener); } /// Accessors for the implied location. Location getLoc() const { return curLoc; } void setLoc(Location loc) { curLoc = loc; } // We allow clients to use the explicit-loc version of create as well. using OpBuilder::create; using OpBuilder::createOrFold; /// Create an operation of specific op type at the current insertion point and /// location. template OpTy create(Args &&...args) { return OpTy::create(*this, curLoc, std::forward(args)...); } /// Create an operation of specific op type at the current insertion point, /// and immediately try to fold it. This functions populates 'results' with /// the results after folding the operation. template void createOrFold(llvm::SmallVectorImpl &results, Args &&...args) { OpBuilder::createOrFold(results, curLoc, std::forward(args)...); } /// Overload to create or fold a single result operation. template std::enable_if_t(), Value> createOrFold(Args &&...args) { return OpBuilder::createOrFold(curLoc, std::forward(args)...); } /// Overload to create or fold a zero result operation. template std::enable_if_t(), OpTy> createOrFold(Args &&...args) { return OpBuilder::createOrFold(curLoc, std::forward(args)...); } /// This builder can also be used to emit diagnostics to the current location. mlir::InFlightDiagnostic emitError(const llvm::Twine &message = llvm::Twine()) { return mlir::emitError(curLoc, message); } mlir::InFlightDiagnostic emitWarning(const llvm::Twine &message = llvm::Twine()) { return mlir::emitWarning(curLoc, message); } mlir::InFlightDiagnostic emitRemark(const llvm::Twine &message = llvm::Twine()) { return mlir::emitRemark(curLoc, message); } private: Location curLoc; }; } // namespace mlir #endif