forked from OSchip/llvm-project
445 lines
15 KiB
C++
445 lines
15 KiB
C++
//===- Builders.cpp - Helpers for constructing MLIR Classes ---------------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "mlir/IR/Dialect.h"
|
|
#include "mlir/IR/IntegerSet.h"
|
|
#include "mlir/IR/Matchers.h"
|
|
#include "mlir/IR/Module.h"
|
|
#include "mlir/IR/StandardTypes.h"
|
|
#include "llvm/Support/raw_ostream.h"
|
|
using namespace mlir;
|
|
|
|
Builder::Builder(ModuleOp module) : context(module.getContext()) {}
|
|
|
|
Identifier Builder::getIdentifier(StringRef str) {
|
|
return Identifier::get(str, context);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Locations.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
Location Builder::getUnknownLoc() { return UnknownLoc::get(context); }
|
|
|
|
Location Builder::getFileLineColLoc(Identifier filename, unsigned line,
|
|
unsigned column) {
|
|
return FileLineColLoc::get(filename, line, column, context);
|
|
}
|
|
|
|
Location Builder::getFusedLoc(ArrayRef<Location> locs, Attribute metadata) {
|
|
return FusedLoc::get(locs, metadata, context);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Types.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
FloatType Builder::getBF16Type() { return FloatType::getBF16(context); }
|
|
|
|
FloatType Builder::getF16Type() { return FloatType::getF16(context); }
|
|
|
|
FloatType Builder::getF32Type() { return FloatType::getF32(context); }
|
|
|
|
FloatType Builder::getF64Type() { return FloatType::getF64(context); }
|
|
|
|
IndexType Builder::getIndexType() { return IndexType::get(context); }
|
|
|
|
IntegerType Builder::getI1Type() { return IntegerType::get(1, context); }
|
|
|
|
IntegerType Builder::getIntegerType(unsigned width) {
|
|
return IntegerType::get(width, context);
|
|
}
|
|
|
|
IntegerType Builder::getIntegerType(unsigned width, bool isSigned) {
|
|
return IntegerType::get(
|
|
width, isSigned ? IntegerType::Signed : IntegerType::Unsigned, context);
|
|
}
|
|
|
|
FunctionType Builder::getFunctionType(ArrayRef<Type> inputs,
|
|
ArrayRef<Type> results) {
|
|
return FunctionType::get(inputs, results, context);
|
|
}
|
|
|
|
TupleType Builder::getTupleType(ArrayRef<Type> elementTypes) {
|
|
return TupleType::get(elementTypes, context);
|
|
}
|
|
|
|
NoneType Builder::getNoneType() { return NoneType::get(context); }
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Attributes.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
NamedAttribute Builder::getNamedAttr(StringRef name, Attribute val) {
|
|
return NamedAttribute(getIdentifier(name), val);
|
|
}
|
|
|
|
UnitAttr Builder::getUnitAttr() { return UnitAttr::get(context); }
|
|
|
|
BoolAttr Builder::getBoolAttr(bool value) {
|
|
return BoolAttr::get(value, context);
|
|
}
|
|
|
|
DictionaryAttr Builder::getDictionaryAttr(ArrayRef<NamedAttribute> value) {
|
|
return DictionaryAttr::get(value, context);
|
|
}
|
|
|
|
IntegerAttr Builder::getIndexAttr(int64_t value) {
|
|
return IntegerAttr::get(getIndexType(), APInt(64, value));
|
|
}
|
|
|
|
IntegerAttr Builder::getI64IntegerAttr(int64_t value) {
|
|
return IntegerAttr::get(getIntegerType(64), APInt(64, value));
|
|
}
|
|
|
|
DenseIntElementsAttr Builder::getI32VectorAttr(ArrayRef<int32_t> values) {
|
|
return DenseIntElementsAttr::get(
|
|
VectorType::get(static_cast<int64_t>(values.size()), getIntegerType(32)),
|
|
values);
|
|
}
|
|
|
|
DenseIntElementsAttr Builder::getI64VectorAttr(ArrayRef<int64_t> values) {
|
|
return DenseIntElementsAttr::get(
|
|
VectorType::get(static_cast<int64_t>(values.size()), getIntegerType(64)),
|
|
values);
|
|
}
|
|
|
|
DenseIntElementsAttr Builder::getI32TensorAttr(ArrayRef<int32_t> values) {
|
|
return DenseIntElementsAttr::get(
|
|
RankedTensorType::get(static_cast<int64_t>(values.size()),
|
|
getIntegerType(32)),
|
|
values);
|
|
}
|
|
|
|
DenseIntElementsAttr Builder::getI64TensorAttr(ArrayRef<int64_t> values) {
|
|
return DenseIntElementsAttr::get(
|
|
RankedTensorType::get(static_cast<int64_t>(values.size()),
|
|
getIntegerType(64)),
|
|
values);
|
|
}
|
|
|
|
IntegerAttr Builder::getI32IntegerAttr(int32_t value) {
|
|
return IntegerAttr::get(getIntegerType(32), APInt(32, value));
|
|
}
|
|
|
|
IntegerAttr Builder::getSI32IntegerAttr(int32_t value) {
|
|
return IntegerAttr::get(getIntegerType(32, /*isSigned=*/true),
|
|
APInt(32, value, /*isSigned=*/true));
|
|
}
|
|
|
|
IntegerAttr Builder::getUI32IntegerAttr(uint32_t value) {
|
|
return IntegerAttr::get(getIntegerType(32, /*isSigned=*/false),
|
|
APInt(32, (uint64_t)value, /*isSigned=*/false));
|
|
}
|
|
|
|
IntegerAttr Builder::getI16IntegerAttr(int16_t value) {
|
|
return IntegerAttr::get(getIntegerType(16), APInt(16, value));
|
|
}
|
|
|
|
IntegerAttr Builder::getI8IntegerAttr(int8_t value) {
|
|
return IntegerAttr::get(getIntegerType(8), APInt(8, value));
|
|
}
|
|
|
|
IntegerAttr Builder::getIntegerAttr(Type type, int64_t value) {
|
|
if (type.isIndex())
|
|
return IntegerAttr::get(type, APInt(64, value));
|
|
return IntegerAttr::get(
|
|
type, APInt(type.getIntOrFloatBitWidth(), value, type.isSignedInteger()));
|
|
}
|
|
|
|
IntegerAttr Builder::getIntegerAttr(Type type, const APInt &value) {
|
|
return IntegerAttr::get(type, value);
|
|
}
|
|
|
|
FloatAttr Builder::getF64FloatAttr(double value) {
|
|
return FloatAttr::get(getF64Type(), APFloat(value));
|
|
}
|
|
|
|
FloatAttr Builder::getF32FloatAttr(float value) {
|
|
return FloatAttr::get(getF32Type(), APFloat(value));
|
|
}
|
|
|
|
FloatAttr Builder::getF16FloatAttr(float value) {
|
|
return FloatAttr::get(getF16Type(), value);
|
|
}
|
|
|
|
FloatAttr Builder::getFloatAttr(Type type, double value) {
|
|
return FloatAttr::get(type, value);
|
|
}
|
|
|
|
FloatAttr Builder::getFloatAttr(Type type, const APFloat &value) {
|
|
return FloatAttr::get(type, value);
|
|
}
|
|
|
|
StringAttr Builder::getStringAttr(StringRef bytes) {
|
|
return StringAttr::get(bytes, context);
|
|
}
|
|
|
|
ArrayAttr Builder::getArrayAttr(ArrayRef<Attribute> value) {
|
|
return ArrayAttr::get(value, context);
|
|
}
|
|
|
|
FlatSymbolRefAttr Builder::getSymbolRefAttr(Operation *value) {
|
|
auto symName =
|
|
value->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
|
|
assert(symName && "value does not have a valid symbol name");
|
|
return getSymbolRefAttr(symName.getValue());
|
|
}
|
|
FlatSymbolRefAttr Builder::getSymbolRefAttr(StringRef value) {
|
|
return SymbolRefAttr::get(value, getContext());
|
|
}
|
|
SymbolRefAttr
|
|
Builder::getSymbolRefAttr(StringRef value,
|
|
ArrayRef<FlatSymbolRefAttr> nestedReferences) {
|
|
return SymbolRefAttr::get(value, nestedReferences, getContext());
|
|
}
|
|
|
|
ArrayAttr Builder::getI32ArrayAttr(ArrayRef<int32_t> values) {
|
|
auto attrs = llvm::to_vector<8>(llvm::map_range(
|
|
values, [this](int32_t v) -> Attribute { return getI32IntegerAttr(v); }));
|
|
return getArrayAttr(attrs);
|
|
}
|
|
|
|
ArrayAttr Builder::getI64ArrayAttr(ArrayRef<int64_t> values) {
|
|
auto attrs = llvm::to_vector<8>(llvm::map_range(
|
|
values, [this](int64_t v) -> Attribute { return getI64IntegerAttr(v); }));
|
|
return getArrayAttr(attrs);
|
|
}
|
|
|
|
ArrayAttr Builder::getIndexArrayAttr(ArrayRef<int64_t> values) {
|
|
auto attrs = llvm::to_vector<8>(
|
|
llvm::map_range(values, [this](int64_t v) -> Attribute {
|
|
return getIntegerAttr(IndexType::get(getContext()), v);
|
|
}));
|
|
return getArrayAttr(attrs);
|
|
}
|
|
|
|
ArrayAttr Builder::getF32ArrayAttr(ArrayRef<float> values) {
|
|
auto attrs = llvm::to_vector<8>(llvm::map_range(
|
|
values, [this](float v) -> Attribute { return getF32FloatAttr(v); }));
|
|
return getArrayAttr(attrs);
|
|
}
|
|
|
|
ArrayAttr Builder::getF64ArrayAttr(ArrayRef<double> values) {
|
|
auto attrs = llvm::to_vector<8>(llvm::map_range(
|
|
values, [this](double v) -> Attribute { return getF64FloatAttr(v); }));
|
|
return getArrayAttr(attrs);
|
|
}
|
|
|
|
ArrayAttr Builder::getStrArrayAttr(ArrayRef<StringRef> values) {
|
|
auto attrs = llvm::to_vector<8>(llvm::map_range(
|
|
values, [this](StringRef v) -> Attribute { return getStringAttr(v); }));
|
|
return getArrayAttr(attrs);
|
|
}
|
|
|
|
ArrayAttr Builder::getAffineMapArrayAttr(ArrayRef<AffineMap> values) {
|
|
auto attrs = llvm::to_vector<8>(llvm::map_range(
|
|
values, [](AffineMap v) -> Attribute { return AffineMapAttr::get(v); }));
|
|
return getArrayAttr(attrs);
|
|
}
|
|
|
|
Attribute Builder::getZeroAttr(Type type) {
|
|
switch (type.getKind()) {
|
|
case StandardTypes::BF16:
|
|
case StandardTypes::F16:
|
|
case StandardTypes::F32:
|
|
case StandardTypes::F64:
|
|
return getFloatAttr(type, 0.0);
|
|
case StandardTypes::Integer: {
|
|
auto width = type.cast<IntegerType>().getWidth();
|
|
if (width == 1)
|
|
return getBoolAttr(false);
|
|
return getIntegerAttr(type, APInt(width, 0));
|
|
}
|
|
case StandardTypes::Vector:
|
|
case StandardTypes::RankedTensor: {
|
|
auto vtType = type.cast<ShapedType>();
|
|
auto element = getZeroAttr(vtType.getElementType());
|
|
if (!element)
|
|
return {};
|
|
return DenseElementsAttr::get(vtType, element);
|
|
}
|
|
default:
|
|
break;
|
|
}
|
|
return {};
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Affine Expressions, Affine Maps, and Integer Sets.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
AffineExpr Builder::getAffineDimExpr(unsigned position) {
|
|
return mlir::getAffineDimExpr(position, context);
|
|
}
|
|
|
|
AffineExpr Builder::getAffineSymbolExpr(unsigned position) {
|
|
return mlir::getAffineSymbolExpr(position, context);
|
|
}
|
|
|
|
AffineExpr Builder::getAffineConstantExpr(int64_t constant) {
|
|
return mlir::getAffineConstantExpr(constant, context);
|
|
}
|
|
|
|
AffineMap Builder::getEmptyAffineMap() { return AffineMap::get(context); }
|
|
|
|
AffineMap Builder::getConstantAffineMap(int64_t val) {
|
|
return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/0,
|
|
getAffineConstantExpr(val));
|
|
}
|
|
|
|
AffineMap Builder::getDimIdentityMap() {
|
|
return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, getAffineDimExpr(0));
|
|
}
|
|
|
|
AffineMap Builder::getMultiDimIdentityMap(unsigned rank) {
|
|
SmallVector<AffineExpr, 4> dimExprs;
|
|
dimExprs.reserve(rank);
|
|
for (unsigned i = 0; i < rank; ++i)
|
|
dimExprs.push_back(getAffineDimExpr(i));
|
|
return AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/0, dimExprs,
|
|
context);
|
|
}
|
|
|
|
AffineMap Builder::getSymbolIdentityMap() {
|
|
return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/1,
|
|
getAffineSymbolExpr(0));
|
|
}
|
|
|
|
AffineMap Builder::getSingleDimShiftAffineMap(int64_t shift) {
|
|
// expr = d0 + shift.
|
|
auto expr = getAffineDimExpr(0) + shift;
|
|
return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
|
|
}
|
|
|
|
AffineMap Builder::getShiftedAffineMap(AffineMap map, int64_t shift) {
|
|
SmallVector<AffineExpr, 4> shiftedResults;
|
|
shiftedResults.reserve(map.getNumResults());
|
|
for (auto resultExpr : map.getResults())
|
|
shiftedResults.push_back(resultExpr + shift);
|
|
return AffineMap::get(map.getNumDims(), map.getNumSymbols(), shiftedResults,
|
|
context);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// OpBuilder
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
OpBuilder::Listener::~Listener() {}
|
|
|
|
/// Insert the given operation at the current insertion point and return it.
|
|
Operation *OpBuilder::insert(Operation *op) {
|
|
if (block)
|
|
block->getOperations().insert(insertPoint, op);
|
|
|
|
if (listener)
|
|
listener->notifyOperationInserted(op);
|
|
return op;
|
|
}
|
|
|
|
/// 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'.
|
|
Block *OpBuilder::createBlock(Region *parent, Region::iterator insertPt,
|
|
TypeRange argTypes) {
|
|
assert(parent && "expected valid parent region");
|
|
if (insertPt == Region::iterator())
|
|
insertPt = parent->end();
|
|
|
|
Block *b = new Block();
|
|
b->addArguments(argTypes);
|
|
parent->getBlocks().insert(insertPt, b);
|
|
setInsertionPointToEnd(b);
|
|
|
|
if (listener)
|
|
listener->notifyBlockCreated(b);
|
|
return b;
|
|
}
|
|
|
|
/// Add new block with 'argTypes' arguments and set the insertion point to the
|
|
/// end of it. The block is placed before 'insertBefore'.
|
|
Block *OpBuilder::createBlock(Block *insertBefore, TypeRange argTypes) {
|
|
assert(insertBefore && "expected valid insertion block");
|
|
return createBlock(insertBefore->getParent(), Region::iterator(insertBefore),
|
|
argTypes);
|
|
}
|
|
|
|
/// Create an operation given the fields represented as an OperationState.
|
|
Operation *OpBuilder::createOperation(const OperationState &state) {
|
|
return insert(Operation::create(state));
|
|
}
|
|
|
|
/// Attempts to fold the given operation and places new results within
|
|
/// 'results'. Returns success if the operation was folded, failure otherwise.
|
|
/// Note: This function does not erase the operation on a successful fold.
|
|
LogicalResult OpBuilder::tryFold(Operation *op,
|
|
SmallVectorImpl<Value> &results) {
|
|
results.reserve(op->getNumResults());
|
|
auto cleanupFailure = [&] {
|
|
results.assign(op->result_begin(), op->result_end());
|
|
return failure();
|
|
};
|
|
|
|
// If this operation is already a constant, there is nothing to do.
|
|
if (matchPattern(op, m_Constant()))
|
|
return cleanupFailure();
|
|
|
|
// Check to see if any operands to the operation is constant and whether
|
|
// the operation knows how to constant fold itself.
|
|
SmallVector<Attribute, 4> constOperands(op->getNumOperands());
|
|
for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i)
|
|
matchPattern(op->getOperand(i), m_Constant(&constOperands[i]));
|
|
|
|
// Try to fold the operation.
|
|
SmallVector<OpFoldResult, 4> foldResults;
|
|
if (failed(op->fold(constOperands, foldResults)) || foldResults.empty())
|
|
return cleanupFailure();
|
|
|
|
// A temporary builder used for creating constants during folding.
|
|
OpBuilder cstBuilder(context);
|
|
SmallVector<Operation *, 1> generatedConstants;
|
|
|
|
// Populate the results with the folded results.
|
|
Dialect *dialect = op->getDialect();
|
|
for (auto &it : llvm::enumerate(foldResults)) {
|
|
// Normal values get pushed back directly.
|
|
if (auto value = it.value().dyn_cast<Value>()) {
|
|
results.push_back(value);
|
|
continue;
|
|
}
|
|
|
|
// Otherwise, try to materialize a constant operation.
|
|
if (!dialect)
|
|
return cleanupFailure();
|
|
|
|
// Ask the dialect to materialize a constant operation for this value.
|
|
Attribute attr = it.value().get<Attribute>();
|
|
auto *constOp = dialect->materializeConstant(
|
|
cstBuilder, attr, op->getResult(it.index()).getType(), op->getLoc());
|
|
if (!constOp) {
|
|
// Erase any generated constants.
|
|
for (Operation *cst : generatedConstants)
|
|
cst->erase();
|
|
return cleanupFailure();
|
|
}
|
|
assert(matchPattern(constOp, m_Constant()));
|
|
|
|
generatedConstants.push_back(constOp);
|
|
results.push_back(constOp->getResult(0));
|
|
}
|
|
|
|
// If we were successful, insert any generated constants.
|
|
for (Operation *cst : generatedConstants)
|
|
insert(cst);
|
|
|
|
return success();
|
|
}
|