9c9a7e9268
This allows dialect-specific attributes to be attached to func results. (or more specifically, FunctionLike ops).
For example:
```
func @f() -> (i32 {my_dialect.some_attr = 3})
```
This attaches my_dialect.some_attr with value 3 to the first result of func @f.
Another more complex example:
```
func @g() -> (i32, f32 {my_dialect.some_attr = "foo", other_dialect.some_other_attr = [1,2,3]}, i1)
```
Here, the second result has two attributes attached.
PiperOrigin-RevId: 275564165
316 lines
12 KiB
C++
316 lines
12 KiB
C++
//===- FunctionSupport.cpp - Utility types for function-like ops ----------===//
|
|
//
|
|
// Copyright 2019 The MLIR Authors.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
// =============================================================================
|
|
|
|
#include "mlir/IR/FunctionSupport.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/OpImplementation.h"
|
|
|
|
using namespace mlir;
|
|
|
|
static ParseResult
|
|
parseArgumentList(OpAsmParser &parser, bool allowVariadic,
|
|
SmallVectorImpl<Type> &argTypes,
|
|
SmallVectorImpl<OpAsmParser::OperandType> &argNames,
|
|
SmallVectorImpl<SmallVector<NamedAttribute, 2>> &argAttrs,
|
|
bool &isVariadic) {
|
|
if (parser.parseLParen())
|
|
return failure();
|
|
|
|
// The argument list either has to consistently have ssa-id's followed by
|
|
// types, or just be a type list. It isn't ok to sometimes have SSA ID's and
|
|
// sometimes not.
|
|
auto parseArgument = [&]() -> ParseResult {
|
|
llvm::SMLoc loc = parser.getCurrentLocation();
|
|
|
|
// Parse argument name if present.
|
|
OpAsmParser::OperandType argument;
|
|
Type argumentType;
|
|
if (succeeded(parser.parseOptionalRegionArgument(argument)) &&
|
|
!argument.name.empty()) {
|
|
// Reject this if the preceding argument was missing a name.
|
|
if (argNames.empty() && !argTypes.empty())
|
|
return parser.emitError(loc, "expected type instead of SSA identifier");
|
|
argNames.push_back(argument);
|
|
|
|
if (parser.parseColonType(argumentType))
|
|
return failure();
|
|
} else if (allowVariadic && succeeded(parser.parseOptionalEllipsis())) {
|
|
isVariadic = true;
|
|
return success();
|
|
} else if (!argNames.empty()) {
|
|
// Reject this if the preceding argument had a name.
|
|
return parser.emitError(loc, "expected SSA identifier");
|
|
} else if (parser.parseType(argumentType)) {
|
|
return failure();
|
|
}
|
|
|
|
// Add the argument type.
|
|
argTypes.push_back(argumentType);
|
|
|
|
// Parse any argument attributes.
|
|
SmallVector<NamedAttribute, 2> attrs;
|
|
if (parser.parseOptionalAttributeDict(attrs))
|
|
return failure();
|
|
argAttrs.push_back(attrs);
|
|
return success();
|
|
};
|
|
|
|
// Parse the function arguments.
|
|
if (parser.parseOptionalRParen()) {
|
|
do {
|
|
unsigned numTypedArguments = argTypes.size();
|
|
if (parseArgument())
|
|
return failure();
|
|
|
|
llvm::SMLoc loc = parser.getCurrentLocation();
|
|
if (argTypes.size() == numTypedArguments &&
|
|
succeeded(parser.parseOptionalComma()))
|
|
return parser.emitError(
|
|
loc, "variadic arguments must be in the end of the argument list");
|
|
} while (succeeded(parser.parseOptionalComma()));
|
|
parser.parseRParen();
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
/// Parse a function result list.
|
|
///
|
|
/// function-result-list ::= function-result-list-parens
|
|
/// | non-function-type
|
|
/// function-result-list-parens ::= `(` `)`
|
|
/// | `(` function-result-list-no-parens `)`
|
|
/// function-result-list-no-parens ::= function-result (`,` function-result)*
|
|
/// function-result ::= type attribute-dict?
|
|
///
|
|
static ParseResult parseFunctionResultList(
|
|
OpAsmParser &parser, SmallVectorImpl<Type> &resultTypes,
|
|
SmallVectorImpl<SmallVector<NamedAttribute, 2>> &resultAttrs) {
|
|
if (failed(parser.parseOptionalLParen())) {
|
|
// We already know that there is no `(`, so parse a type.
|
|
// Because there is no `(`, it cannot be a function type.
|
|
Type ty;
|
|
if (parser.parseType(ty))
|
|
return failure();
|
|
resultTypes.push_back(ty);
|
|
resultAttrs.emplace_back();
|
|
return success();
|
|
}
|
|
|
|
// Special case for an empty set of parens.
|
|
if (succeeded(parser.parseOptionalRParen()))
|
|
return success();
|
|
|
|
// Parse individual function results.
|
|
do {
|
|
resultTypes.emplace_back();
|
|
resultAttrs.emplace_back();
|
|
if (parser.parseType(resultTypes.back()) ||
|
|
parser.parseOptionalAttributeDict(resultAttrs.back())) {
|
|
return failure();
|
|
}
|
|
} while (succeeded(parser.parseOptionalComma()));
|
|
return parser.parseRParen();
|
|
}
|
|
|
|
/// Parse a function signature, starting with a name and including the
|
|
/// parameter list.
|
|
static ParseResult parseFunctionSignature(
|
|
OpAsmParser &parser, bool allowVariadic,
|
|
SmallVectorImpl<OpAsmParser::OperandType> &argNames,
|
|
SmallVectorImpl<Type> &argTypes,
|
|
SmallVectorImpl<SmallVector<NamedAttribute, 2>> &argAttrs, bool &isVariadic,
|
|
SmallVectorImpl<Type> &resultTypes,
|
|
SmallVectorImpl<SmallVector<NamedAttribute, 2>> &resultAttrs) {
|
|
if (parseArgumentList(parser, allowVariadic, argTypes, argNames, argAttrs,
|
|
isVariadic))
|
|
return failure();
|
|
if (succeeded(parser.parseOptionalArrow()))
|
|
return parseFunctionResultList(parser, resultTypes, resultAttrs);
|
|
return success();
|
|
}
|
|
|
|
/// Parser implementation for function-like operations. Uses `funcTypeBuilder`
|
|
/// to construct the custom function type given lists of input and output types.
|
|
ParseResult
|
|
mlir::impl::parseFunctionLikeOp(OpAsmParser &parser, OperationState &result,
|
|
bool allowVariadic,
|
|
mlir::impl::FuncTypeBuilder funcTypeBuilder) {
|
|
SmallVector<OpAsmParser::OperandType, 4> entryArgs;
|
|
SmallVector<SmallVector<NamedAttribute, 2>, 4> argAttrs;
|
|
SmallVector<SmallVector<NamedAttribute, 2>, 4> resultAttrs;
|
|
SmallVector<Type, 4> argTypes;
|
|
SmallVector<Type, 4> resultTypes;
|
|
auto &builder = parser.getBuilder();
|
|
|
|
// Parse the name as a symbol reference attribute.
|
|
SymbolRefAttr nameAttr;
|
|
if (parser.parseAttribute(nameAttr, ::mlir::SymbolTable::getSymbolAttrName(),
|
|
result.attributes))
|
|
return failure();
|
|
// Convert the parsed function attr into a string attr.
|
|
result.attributes.back().second = builder.getStringAttr(nameAttr.getValue());
|
|
|
|
// Parse the function signature.
|
|
auto signatureLocation = parser.getCurrentLocation();
|
|
bool isVariadic = false;
|
|
if (parseFunctionSignature(parser, allowVariadic, entryArgs, argTypes,
|
|
argAttrs, isVariadic, resultTypes, resultAttrs))
|
|
return failure();
|
|
|
|
std::string errorMessage;
|
|
if (auto type = funcTypeBuilder(builder, argTypes, resultTypes,
|
|
impl::VariadicFlag(isVariadic), errorMessage))
|
|
result.addAttribute(getTypeAttrName(), TypeAttr::get(type));
|
|
else
|
|
return parser.emitError(signatureLocation)
|
|
<< "failed to construct function type"
|
|
<< (errorMessage.empty() ? "" : ": ") << errorMessage;
|
|
|
|
// If function attributes are present, parse them.
|
|
if (succeeded(parser.parseOptionalKeyword("attributes")))
|
|
if (parser.parseOptionalAttributeDict(result.attributes))
|
|
return failure();
|
|
|
|
// Add the attributes to the function arguments.
|
|
SmallString<8> attrNameBuf;
|
|
for (unsigned i = 0, e = argTypes.size(); i != e; ++i)
|
|
if (!argAttrs[i].empty())
|
|
result.addAttribute(getArgAttrName(i, attrNameBuf),
|
|
builder.getDictionaryAttr(argAttrs[i]));
|
|
|
|
// Add the attributes to the function results.
|
|
for (unsigned i = 0, e = resultTypes.size(); i != e; ++i)
|
|
if (!resultAttrs[i].empty())
|
|
result.addAttribute(getResultAttrName(i, attrNameBuf),
|
|
builder.getDictionaryAttr(resultAttrs[i]));
|
|
|
|
// Parse the optional function body.
|
|
auto *body = result.addRegion();
|
|
if (parser.parseOptionalRegion(*body, entryArgs,
|
|
entryArgs.empty() ? llvm::ArrayRef<Type>()
|
|
: argTypes))
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
// Print a function result list.
|
|
static void printFunctionResultList(OpAsmPrinter &p, ArrayRef<Type> types,
|
|
ArrayRef<ArrayRef<NamedAttribute>> attrs) {
|
|
assert(!types.empty() && "Should not be called for empty result list.");
|
|
auto &os = p.getStream();
|
|
bool needsParens =
|
|
types.size() > 1 || types[0].isa<FunctionType>() || !attrs[0].empty();
|
|
if (needsParens)
|
|
os << '(';
|
|
interleaveComma(llvm::zip(types, attrs), os,
|
|
[&](const std::tuple<Type, ArrayRef<NamedAttribute>> &t) {
|
|
p.printType(std::get<0>(t));
|
|
p.printOptionalAttrDict(std::get<1>(t));
|
|
});
|
|
if (needsParens)
|
|
os << ')';
|
|
}
|
|
|
|
/// Print the signature of the function-like operation `op`. Assumes `op` has
|
|
/// the FunctionLike trait and passed the verification.
|
|
static void printSignature(OpAsmPrinter &p, Operation *op,
|
|
ArrayRef<Type> argTypes, bool isVariadic,
|
|
ArrayRef<Type> resultTypes) {
|
|
Region &body = op->getRegion(0);
|
|
bool isExternal = body.empty();
|
|
|
|
p << '(';
|
|
for (unsigned i = 0, e = argTypes.size(); i < e; ++i) {
|
|
if (i > 0)
|
|
p << ", ";
|
|
|
|
if (!isExternal) {
|
|
p.printOperand(body.front().getArgument(i));
|
|
p << ": ";
|
|
}
|
|
|
|
p.printType(argTypes[i]);
|
|
p.printOptionalAttrDict(::mlir::impl::getArgAttrs(op, i));
|
|
}
|
|
|
|
if (isVariadic) {
|
|
if (!argTypes.empty())
|
|
p << ", ";
|
|
p << "...";
|
|
}
|
|
|
|
p << ')';
|
|
|
|
if (!resultTypes.empty()) {
|
|
p.getStream() << " -> ";
|
|
SmallVector<ArrayRef<NamedAttribute>, 4> resultAttrs;
|
|
for (int i = 0, e = resultTypes.size(); i < e; ++i)
|
|
resultAttrs.push_back(::mlir::impl::getResultAttrs(op, i));
|
|
printFunctionResultList(p, resultTypes, resultAttrs);
|
|
}
|
|
}
|
|
|
|
/// Printer implementation for function-like operations. Accepts lists of
|
|
/// argument and result types to use while printing.
|
|
void mlir::impl::printFunctionLikeOp(OpAsmPrinter &p, Operation *op,
|
|
ArrayRef<Type> argTypes, bool isVariadic,
|
|
ArrayRef<Type> resultTypes) {
|
|
// Print the operation and the function name.
|
|
auto funcName =
|
|
op->getAttrOfType<StringAttr>(::mlir::SymbolTable::getSymbolAttrName())
|
|
.getValue();
|
|
p << op->getName() << ' ';
|
|
p.printSymbolName(funcName);
|
|
|
|
// Print the signature.
|
|
printSignature(p, op, argTypes, isVariadic, resultTypes);
|
|
|
|
// Print out function attributes, if present.
|
|
SmallVector<StringRef, 2> ignoredAttrs = {
|
|
::mlir::SymbolTable::getSymbolAttrName(), getTypeAttrName()};
|
|
|
|
SmallString<8> attrNameBuf;
|
|
|
|
// Ignore any argument attributes.
|
|
std::vector<SmallString<8>> argAttrStorage;
|
|
for (unsigned i = 0, e = argTypes.size(); i != e; ++i)
|
|
if (op->getAttr(getArgAttrName(i, attrNameBuf)))
|
|
argAttrStorage.emplace_back(attrNameBuf);
|
|
ignoredAttrs.append(argAttrStorage.begin(), argAttrStorage.end());
|
|
|
|
// Ignore any result attributes.
|
|
std::vector<SmallString<8>> resultAttrStorage;
|
|
for (unsigned i = 0, e = resultTypes.size(); i != e; ++i)
|
|
if (op->getAttr(getResultAttrName(i, attrNameBuf)))
|
|
resultAttrStorage.emplace_back(attrNameBuf);
|
|
ignoredAttrs.append(resultAttrStorage.begin(), resultAttrStorage.end());
|
|
|
|
auto attrs = op->getAttrs();
|
|
if (attrs.size() > ignoredAttrs.size()) {
|
|
p << "\n attributes ";
|
|
p.printOptionalAttrDict(attrs, ignoredAttrs);
|
|
}
|
|
|
|
// Print the body if this is not an external function.
|
|
Region &body = op->getRegion(0);
|
|
if (!body.empty())
|
|
p.printRegion(body, /*printEntryBlockArgs=*/false,
|
|
/*printBlockTerminators=*/true);
|
|
}
|