#include "mlir/IR/Block.h"
#include "mlir/IR/OpDefinition.h"
+#include "llvm/ADT/SmallString.h"
namespace mlir {
class BlockAndValueMapping;
llvm::iterator_range<args_iterator> getArguments() {
return {args_begin(), args_end()};
}
+
+ //===--------------------------------------------------------------------===//
+ // Argument Attributes
+ //===--------------------------------------------------------------------===//
+
+ /// FuncOp allows for attaching attributes to each of the respective function
+ /// arguments. These argument attributes are stored as DictionaryAttrs in the
+ /// main operation attribute dictionary. The name of these entries is `arg`
+ /// followed by the index of the argument. These argument attribute
+ /// dictionaries are optional, and will generally only exist if they are
+ /// non-empty.
+
+ /// Return all of the attributes for the argument at 'index'.
+ ArrayRef<NamedAttribute> getArgAttrs(unsigned index) {
+ auto argDict = getArgAttrDict(index);
+ return argDict ? argDict.getValue() : llvm::None;
+ }
+
+ /// Return all argument attributes of this function.
+ void getAllArgAttrs(SmallVectorImpl<DictionaryAttr> &result) {
+ for (unsigned i = 0, e = getNumArguments(); i != e; ++i)
+ result.emplace_back(getArgAttrDict(i));
+ }
+
+ /// Return the specified attribute, if present, for the argument at 'index',
+ /// null otherwise.
+ Attribute getArgAttr(unsigned index, Identifier name) {
+ auto argDict = getArgAttrDict(index);
+ return argDict ? argDict.get(name) : nullptr;
+ }
+ Attribute getArgAttr(unsigned index, StringRef name) {
+ auto argDict = getArgAttrDict(index);
+ return argDict ? argDict.get(name) : nullptr;
+ }
+
+ template <typename AttrClass>
+ AttrClass getArgAttrOfType(unsigned index, Identifier name) {
+ return getArgAttr(index, name).dyn_cast_or_null<AttrClass>();
+ }
+ template <typename AttrClass>
+ AttrClass getArgAttrOfType(unsigned index, StringRef name) {
+ return getArgAttr(index, name).dyn_cast_or_null<AttrClass>();
+ }
+
+ /// Set the attributes held by the argument at 'index'.
+ void setArgAttrs(unsigned index, ArrayRef<NamedAttribute> attributes);
+ void setArgAttrs(unsigned index, NamedAttributeList attributes);
+ void setAllArgAttrs(ArrayRef<NamedAttributeList> attributes) {
+ assert(attributes.size() == getNumArguments());
+ for (unsigned i = 0, e = attributes.size(); i != e; ++i)
+ setArgAttrs(i, attributes[i]);
+ }
+
+ /// If the an attribute exists with the specified name, change it to the new
+ /// value. Otherwise, add a new attribute with the specified name/value.
+ void setArgAttr(unsigned index, Identifier name, Attribute value);
+ void setArgAttr(unsigned index, StringRef name, Attribute value) {
+ setArgAttr(index, Identifier::get(name, getContext()), value);
+ }
+
+ /// Remove the attribute 'name' from the argument at 'index'.
+ NamedAttributeList::RemoveResult removeArgAttr(unsigned index,
+ Identifier name);
+
+private:
+ /// Returns the attribute entry name for the set of argument attributes at
+ /// index 'arg'.
+ static StringRef getArgAttrName(unsigned arg, SmallVectorImpl<char> &out);
+
+ /// Returns the dictionary attribute corresponding to the argument at 'index'.
+ /// If there are no argument attributes at 'index', a null attribute is
+ /// returned.
+ DictionaryAttr getArgAttrDict(unsigned index) {
+ assert(index < getNumArguments() && "invalid argument number");
+ SmallString<8> nameOut;
+ return getAttrOfType<DictionaryAttr>(getArgAttrName(index, nameOut));
+ }
};
} // end namespace mlir
// Add the argument type.
argTypes.push_back(argumentType);
- // TODO(riverriddle) Parse argument attributes.
- // Parse the attribute dict.
- // SmallVector<NamedAttribute, 2> attrs;
- // if (parser->parseOptionalAttributeDict(attrs))
- // return failure();
- // argAttrs.push_back(attrs);
+ // Parse any argument attributes.
+ SmallVector<NamedAttribute, 2> attrs;
+ if (parser->parseOptionalAttributeDict(attrs))
+ return failure();
+ argAttrs.push_back(attrs);
return success();
};
if (parser->parseOptionalAttributeDict(result->attributes))
return failure();
- // TODO(riverriddle) Parse argument attributes.
// Add the attributes to the function arguments.
- // for (unsigned i = 0, e = type.getNumInputs(); i != e; ++i)
- // if (!argAttrs[i].empty())
- // result->addAttribute(("arg" + Twine(i)).str(),
- // builder.getDictionaryAttr(argAttrs[i]));
+ SmallString<8> argAttrName;
+ for (unsigned i = 0, e = type.getNumInputs(); i != e; ++i)
+ if (!argAttrs[i].empty())
+ result->addAttribute(getArgAttrName(i, argAttrName),
+ builder.getDictionaryAttr(argAttrs[i]));
// Parse the optional function body.
auto *body = result->addRegion();
*p << ": ";
}
+ // Print the type followed by any argument attributes.
p->printType(fnType.getInput(i));
-
- // TODO(riverriddle) Print argument attributes.
- // Print the attributes for this argument.
- // p->printOptionalAttrDict(op.getArgAttrs(i));
+ p->printOptionalAttrDict(op.getArgAttrs(i));
}
*p << ')';
printFunctionSignature(p, *this);
// Print out function attributes, if present.
- auto attrs = getAttrs();
+ SmallVector<StringRef, 2> ignoredAttrs = {"name", "type"};
- // We must have more attributes than <name, type>.
- constexpr unsigned kNumHiddenAttrs = 2;
- if (attrs.size() > kNumHiddenAttrs) {
+ // Ignore any argument attributes.
+ std::vector<SmallString<8>> argAttrStorage;
+ SmallString<8> argAttrName;
+ for (unsigned i = 0, e = getNumArguments(); i != e; ++i)
+ if (getAttr(getArgAttrName(i, argAttrName)))
+ argAttrStorage.emplace_back(argAttrName);
+ ignoredAttrs.append(argAttrStorage.begin(), argAttrStorage.end());
+
+ auto attrs = getAttrs();
+ if (attrs.size() > ignoredAttrs.size()) {
*p << "\n attributes ";
- p->printOptionalAttrDict(attrs, {"name", "type"});
+ p->printOptionalAttrDict(attrs, ignoredAttrs);
}
// Print the body if this is not an external function.
}
*p << '\n';
}
+
+//===----------------------------------------------------------------------===//
+// Function Argument Attribute.
+//===----------------------------------------------------------------------===//
+
+/// Set the attributes held by the argument at 'index'.
+void FuncOp::setArgAttrs(unsigned index, ArrayRef<NamedAttribute> attributes) {
+ assert(index < getNumArguments() && "invalid argument number");
+ SmallString<8> nameOut;
+ getArgAttrName(index, nameOut);
+
+ if (attributes.empty())
+ return (void)removeAttr(nameOut);
+ setAttr(nameOut, DictionaryAttr::get(attributes, getContext()));
+}
+
+void FuncOp::setArgAttrs(unsigned index, NamedAttributeList attributes) {
+ assert(index < getNumArguments() && "invalid argument number");
+ SmallString<8> nameOut;
+ if (auto newAttr = attributes.getDictionary())
+ return setAttr(getArgAttrName(index, nameOut), newAttr);
+ removeAttr(getArgAttrName(index, nameOut));
+}
+
+/// If the an attribute exists with the specified name, change it to the new
+/// value. Otherwise, add a new attribute with the specified name/value.
+void FuncOp::setArgAttr(unsigned index, Identifier name, Attribute value) {
+ auto curAttr = getArgAttrDict(index);
+ NamedAttributeList attrList(curAttr);
+ attrList.set(name, value);
+
+ // If the attribute changed, then set the new arg attribute list.
+ if (curAttr != attrList.getDictionary())
+ setArgAttrs(index, attrList);
+}
+
+/// Remove the attribute 'name' from the argument at 'index'.
+NamedAttributeList::RemoveResult FuncOp::removeArgAttr(unsigned index,
+ Identifier name) {
+ // Build an attribute list and remove the attribute at 'name'.
+ NamedAttributeList attrList(getArgAttrDict(index));
+ auto result = attrList.remove(name);
+
+ // If the attribute was removed, then update the argument dictionary.
+ if (result == NamedAttributeList::RemoveResult::Removed)
+ setArgAttrs(index, attrList);
+ return result;
+}
+
+/// Returns the attribute entry name for the set of argument attributes at index
+/// 'arg'.
+StringRef FuncOp::getArgAttrName(unsigned arg, SmallVectorImpl<char> &out) {
+ out.clear();
+ return ("arg" + Twine(arg)).toStringRef(out);
+}
projects / platform / upstream / llvm.git / commitdiff