This is a proposal, not completely fleshed out on every aspect (PDL and DRR for example), but worth discussing since this is touching on a very core aspect of our modeling.
I tentatively plan to present this in the Open Meeting next Thurday, but I may have to push it out one week.
Background
Operations extensibility has been centered around the use of _Attribute_s since the beginning: this is the only mechanism to attach “data” to an operation. The operation only keeps a reference to the data: the Attribute stores the data uniquely in the MLIRContext. When an Attribute is created, the data is hashed and checked against a dedicated storage in the context if it is already present in which case the existing data is used, otherwise the data is copied into the storage. As such all data is owned by the MLIRContext and only gets destroyed when the context is deleted. One advantage of this is that a simple pointer comparison is enough to compare two attributes for equality, which is advantageous for large data.
A new kind of Attribute was recently introduced around the concept of “Resource”. This is still a work in progress and the documentation does not exist yet, here is the dialect interface. The idea is to decouple 1) the storage of the data from the IR (the data isn’t printed to the MLIR file) and 2) the lifetime of the data from the MLIRContext: the resource does not copy the data into the context and as such we can create such an Attribute from already existing large blob of data in memory. The data can also be shared across multiple MLIRContext. This is convenient when referring to large data stored in files (like ML weights) that can be mmap and referred to by the IR.
An Operation actually holds a single Attribute: a DictionaryAttr. All interactions through accessors are wrappers around this DictionaryAttr. For example, let’s take arith.cmpi:
def Arith_CmpIOp
: Arith_CompareOpOfAnyRank<"cmpi",
[DeclareOpInterfaceMethods<InferIntRangeInterface>]> {
let summary = "integer comparison operation";
let arguments = (ins Arith_CmpIPredicateAttr:$predicate,
SignlessIntegerLikeOfAnyRank:$lhs,
SignlessIntegerLikeOfAnyRank:$rhs);
Which appears in the textual assembly this way:
// Custom form of scalar "signed less than" comparison.
%x = arith.cmpi slt, %lhs, %rhs : i32
// Generic form of the same operation.
%x = "arith.cmpi"(%lhs, %rhs) {predicate = 2 : i64} : (i32, i32) -> i1
// Custom form of vector equality comparison.
%x = arith.cmpi eq, %lhs, %rhs : vector<4xi64>
// Generic form of the same operation.
%x = "arith.cmpi"(%lhs, %rhs) {predicate = 0 : i64}
: (vector<4xi64>, vector<4xi64>) -> vector<4xi1>
The defined Attribute here models the predicate for the comparison, accessors are generated:
::mlir::arith::CmpIPredicate getPredicate();
void setPredicate(::mlir::arith::CmpIPredicate attrValue);
The arith::CmpIOp::fold() method shows examples of this usage: getPredicate() and setPredicate(). The underlying implementation works by looking up the Attribute in the Dictionary by name.
Here is how the lookup is optimized:
-
Getting the key to use in the dictionary (a StringAttr), while avoiding
getPredicate()having to query the MLIRContext and rehash the"predicate"string key.An
Operationhas anOperationNameattached to it: this is the structure defined in theMLIRContextwhen anOperationis registered. It holds the traits, interfaces, etc. specific to an Operation. It also contains an array ofStringAttr: these are the names of the attributes used as keys in the Dictionary, stored in a predefined order (for this operation there is a single entry matching the"predicate"StringAttr). This allows to directly retrieve the key at the price of 3 pointers dereferencing. -
The Dictionary stores the attributes in an array sorted by keys, as long as there are less than 16 entries in the dictionary, we use a linear scan comparing just the pointer value of the StringAttr key. Otherwise we fallback to a custom binary search based on more expensive StringRef comparisons.
Now for the setter implementation:
-
The provided enum value has to be converted to an Attribute:
mlir::arith::CmpIPredicateAttr::get(context, value)which will lookup the storage for this attribute in the MLIRContext, hash the value and try to find an already existing Attribute or create a new unique one. -
The
StringAttrkey for the Dictionary entry is looked up the same way as for the lookup above. -
Since
Attributes are immutable, theDictionaryAttrattached to the operation must be entirely rebuilt. We’ll first copy the entire dictionary into aNamedAttrList(basically aSmallVectorofNamedAttribute). TheNamedAttrListis then modified in place to add or modify the attribute. This involves looking up the entry if it exists (using the same process as the lookup above) and modifying it in place, or inserting a new entry in the vector in the sorted position. -
Finally the
NamedAttrListis converted to aDictionaryAttrby looking up the storage in the MLIRContext, hashing the entire array of attributes, and returning an existing Attribute or inserting the array into the storage and returning the new DictionaryAttr.
These accessors are quite involved, this is why it is preferable to cache the result of the getter as much as possible:
…
… (..., op.getPredicate(), …);
…
… (..., op.getPredicate(), …);
would be better written:
auto pred = op.getPredicate();
…
… (..., pred, …);
…
… (..., pred, …);
Also, changing multiple attributes using the setter is very inefficient. With an operation that has two i64 integer attributes axis1 and axis2, the sequence:
op.setAxis1(0);
op.setAxis2(1);
will go through the process described above twice and build two new dictionaries that will “leak” permanently in the MLIRContext. The most efficient way to update multiple attributes involves avoiding the convenient accessors:
int64_t newAxis1 = 0, newAxis2 = 1;
// Copy the dictionary into a vector of attributes.
NamedAttrList attrs(op.getAttrDictionary());
// Mutate the vector of attributes in-place
// Using a string key for example.
attrs.set("axis1", IntegerAttr::get(IntegerType::get(ctx, 64), newAxis1);
// It is more efficient to use precomputed StringAttr keys using accessors like
// getAxis2AttrName() instead of a string key.
attrs.set(op.getAxis2AttrName(), IntegerAttr::get(IntegerType::get(ctx, 64), newAxis2);
// Build a new DictionaryAttr in the context.
DictionaryAttr dict = attrs.getDictionary(ctx);
// Update the operation in-place by swapping-in the new Dictionary.
op.setAttrs(dict);
This is clearly a lot of boilerplate, and it’s not surprising that the convenience accessors are pervasively used everywhere despite the cost.
Introducing Properties
While attributes are a very flexible and extensible mechanism to attach data to an Operation, the price to pay (both in terms of runtime cost associated with the accessors and the lifetime tied to the MLIRContext) does not make it necessarily a good tradeoff for most common small data attached to operations like the predicate on arith.cmpi example above.
This proposal isn’t completely fleshed out and comes with a few possibilities moving forward, each with their own tradeoffs. I will try to expose this in details, so that we can evaluate our options during the next open meeting.
I will first introduce a new concept I call “properties”: the idea is to allow operations to carry some data without involving any attribute:
-
The data is allocated inline with the Operation and can be used as an alternative to attributes to store data that is specific to an operation.
-
An
Attributecan also be stored inside the properties storage if desired, but any other kind of data can be present as well. -
This offers a way to store and mutate data without uniquing in the Context contrary to attributes.
-
The lifetime of the data is tied to the Operation.
D141742 is the implementation for this new feature and the OpPropertiesTest.cpp has an example where a struct with a std::vector<> as well as a std::shared_ptr<const std::string> attached to an operation and mutated in-place.
In terms of ODS, the integration is functional, but there is room for further improvements. Right now it’ll look like this for example:
// Op with a properties struct defined inline.
def TestOpWithProperties : TEST_Op<"with_properties"> {
let assemblyFormat = "prop-attr attr-dict";
let properties = (ins
Property<"int64_t">:$a,
StrAttr:$b, // Attributes can directly be used here.
ArrayProperty<"int64_t", 4>:$array // example of an array
);
}
All the boiler plate is generated automatically, and the properties can be accessed through a new accessor named getProperties(). In the example above it’ll return a mutable reference to a struct looking like:
struct Properties {
int64_t a;
StringAttr b;
int array[4];
}
The generic syntax for operation is modified to print the properties as an Attribute in between < > after the operand list (this implies that unregistered operations will have an extra Attribute storage available). The TestOpWithProperties above will be printed generically as:
"test.with_properties"() <{a = 32 : i64,
array = array<i64: 1, 2, 3, 4>,
b = "foo"}> : () -> ()
Of course custom printers and parsers have full freedom as usual. The ODS declarative assembly does not support yet referring to individual properties members. Mutating the properties for this Operation in an optimal way is trivial and does not involve any gotchas:
auto opWithProp = dyn_cast<TestOpWithProperties>(op.get());
// Get a mutable reference to the properties for this operation and modify it
// in place one member at a time.
TestOpWithProperties::Properties &prop = opWithProp.getProperties();
prop.a += 42;
prop.b = StringAttr::get(ctx, "some label");
prop.array[2] += prop.a;
Right now, build methods generated through ODS haven’t been updated to take Properties as input, however the properties can be updated post-creation in a similar way as the mutation above:
auto opWithProp = builder.create<TestOpWithProperties>(loc);
// Get a mutable reference to the properties for this operation and modify it
// in place one member at a time.
TestOpWithProperties::Properties &prop = opWithProp.getProperties();
prop.a = 42;
prop.b = StringAttr::get(ctx, "some label");
prop.setArray({1,2,3,4);
We will update ODS generators to emit build methods that will support:
auto opWithProp = builder.create<TestOpWithProperties>(loc, /*a=*/42;
/*b=*/StringAttr::get(ctx, "some label"),
/*array=*/{1,2,3,4});
Not Discardable
One aspect of the DictionaryAttr attached to an operation is that it mixes inherent and discardable attributes. As defined in LangRef:
_ The attribute entries are considered to be of two different kinds based on whether their dictionary key has a dialect prefix:_
- inherent attributes are inherent to the definition of an operation’s semantics. The operation itself is expected to verify the consistency of these attributes. An example is the predicate attribute of the arith.cmpi op. These attributes must have names that do not start with a dialect prefix.
- discardable attributes have semantics defined externally to the operation itself, but must be compatible with the operations’s semantics. These attributes must have names that start with a dialect prefix. The dialect indicated by the dialect prefix is expected to verify these attributes. An example is the gpu.container_module attribute.
On the other hand, the Properties data-structure is non-discardable and not designed to be introspectable opaquely (that is the client needs to know the Operation class to access the properties, possible through an OpInterface). It is an alternative solution to inherent attributes.
When to use Properties
Properties are almost always like an appropriate replacement for inherent attributes. In particular, Attribute can still be used inside properties, allowing to make every possible tradeoff within the Properties framework. Ultimately, defining a DictionaryAttr as Properties for an Operation would be strictly equivalent to the existing storage, but the current Attribute storage will always stay relevant for discardable attributes anyway.
When to not use Properties
- When you intend to expose an attribute opaquely through the dictionary. However it is likely that we may prefer to rely on OpInterface to provide access to the attribute. For example for handling FastMathFlags, we likely should have a FastMathOpInterface wrapping the access to the flags instead of using name lookups in the dictionary.
- When you need to model something that will be exposed through an AttributeInterface (you can’t attach interfaces to properties right now). Although, the Attribute can be stored inside the Properties and exposed as such.
- PDL introspection is not supported right now for something not stored in the properties as an attribute. DRR support hasn’t been explored either.
Implementation Notes
Here is a walkthrough of the example in mlir/unittests/IR/OpPropertiesTest.cpp:
// A c++ struct with 4 members.
struct TestProperties {
int a = -1;
float b = -1.;
std::vector<int64_t> array = {-33};
// A shared_ptr to a const object is safe: it is equivalent to a value-based
// member. A non-const shared_ptr would be unsafe: after cloning an Operation,
// any mutation of the pointed value from the original would apply to the clone
// as well.
// Here the label will be deallocated when the last operation
// referring to it is destroyed. However there is no builtin pool-allocation:
// this is offloaded to the client.
std::shared_ptr<const std::string> label;
};
Before being able to use this as properties on an operation, three functions must be defined:
// Compute a hash for the structure: this is needed for
// computing OperationEquivalence, think about CSE.
llvm::hash_code computeHash(const TestProperties &prop);
// Convert the structure to an attribute: this is used when printing
// an operation in the generic form.
Attribute getPropertiesAsAttribute(MLIRContext *ctx,
const TestProperties &prop);
// Convert the structure from an attribute: this is used when
// parsing an operation from the generic form.
LogicalResult setPropertiesFromAttribute(TestProperties &prop,
Attribute attr,
InFlightDiagnostic *diagnostic);
Any operation can very easily register this structure through a simple type alias declaration:
/// A custom operation for the purpose of showcasing how to use "properties".
class OpWithProperties : public Op<OpWithProperties> {
public:
// Begin boilerplate
…
// End boilerplate
// This alias is the only definition needed for enabling "properties" for this
// operation.
using Properties = TestProperties;
};
The properties is then directly accessible, here is an example involving accessing and mutating the data attached to an operation:
auto opWithProp = dyn_cast<OpWithProperties>(op.get());
// Get a mutable reference to the properties for this operation and modify it
// in place one member at a time.
TestProperties &prop = opWithProp.getProperties();
prop.a += 42;
prop.b = 42.;
prop.array.push_back(42); // std::vector privately attached to this operation.
prop.label = std::make_shared<std::string>("foo bar");
The getProperties() accessor is cheap: it retrieves directly a member of the operation by computing an offset, this is no different from getting any other member of an Operation (result, operands, type, …). From there, everything is purely native C++: there is nothing hidden and nothing specific to MLIR involved (in particular the MLIRContext is not involved at all).
While by default, the parsing and printing will use the getPropertiesAsAttribute()/setPropertiesFromAttribute() and rely on existing attribute printing/parsing format, 2 extra functions can also be used to customize the format:
static void customPrintProperties(OpAsmPrinter &p,
PropertiesWithCustomPrint &prop);
static ParseResult customParseProperties(OpAsmParser &parser,
PropertiesWithCustomPrint &prop);
See the TestDialect.cpp file in the revision above for an example.
In terms of ODS, the integration is functional, but there is room for further improvements. See some examples in TestOps.td, including this one for example:
// Op with a properties struct defined inline.
def TestOpWithProperties : TEST_Op<"with_properties"> {
let assemblyFormat = "prop-attr attr-dict";
let properties = (ins
Property<"int64_t">:$a,
StrAttr:$b, // Attributes can directly be used here.
ArrayProperty<"int64_t", 4>:$array // example of an array
);
}
TableGen will generate the following C++:
class TestOpWithProperties : public Op<TestOpWithProperties, /*traits*/… > {
public:
// Begin boilerplate
…
// End boilerplate
struct Properties {
using aTy = int64_t;
using bTy = ::mlir::StringAttr;
using arrayTy = int64_t[4];
aTy a;
bTy b;
arrayTy array;
int64_t getA() {
auto &storage = this->a;
return storage;
}
void setA(const int64_t &value) {
auto &storage = this->a;
storage = value;
}
::mlir::StringAttr getB() {
auto &storage = this->b;
return storage;
}
void setB(const ::mlir::StringAttr &value) {
auto &storage = this->b;
storage = value;
}
::llvm::ArrayRef<int64_t> getArray() {
auto &storage = this->array;
return storage;
}
void setArray(const ::llvm::ArrayRef<int64_t> &value) {
auto &storage = this->array;
::llvm::copy(value, storage);
}
};
…
The setPropertiesFromAttr, getPropertiesAsAttr, and computePropertiesHash are also automatically generated, dispatching to every single individual member. In this case no C++ code needs to be written by hand by the user, the boilerplate is entirely generated.
Performance Downsides & Considerations
Registered Operations that don’t use Properties don’t pay any price for the existence of this feature in the codebase. Unregistered operations (including in registered dialects) will have a single new attribute member added (we don’t know if they will use or not some property storage).
While the advantages of using Properties over Attribute should seem fairly straightforward by now, there are some downsides when using Properties:
- Memory footprint may increase:
Operationallocations get larger than before. While “properties” should stay small there is a tradeoff between “pooling” the allocation in the context (and leaking memory there “forever”) and inlining it within the operation. Note that the proposed scheme would allow a Dialect to implement some intermediate scheme: constant data could be pooled in the context but using reference counting to free the data when all operations referring to them are deleted. - Comparison is no longer “single pointer”: checking that two operations have the same Properties requires calling the Properties comparison operator, while with Attributes comparing the DictionaryAttr pointer is enough.
- As seen on
OperationName, more hooks are introduced to interact with Operation adding some extra runtime cost to come operations:- When creating an operation, we initialize the properties by calling its default constructor (through an indirect call) before calling the assignment operator.
- When cloning an operation, we call the assignment operator and copy the properties.
- When deleting an operation, we call the properties destructor.
- OperationEquivalence (called by CSE for example) will hash the properties (through an indirect call).
Breaking the Model Further: opt-in all inherent attributes into Properties
The proposal above is designed with backward compatibility in mind: landing this as-is wouldn’t break any existing code out there immediately. This allows for a gradual adoption of this feature. However there is an argument that we’re breaking some fundamental assumptions of the underlying MLIR model: trying to piecewise clone or propagate information from one operation to another will be broken. Existing MLIR code that manipulates and transforms operations opaquely this way could become subtly broken when a dialect starts to adopt properties. A more immediately breaking but more consistent approach would be:
- Stop storing the inherent attributes (so the ones declared in ODS) in the DictionaryAttr. Instead store them as individual members of the properties. The DictionaryAttr would be exclusively reserved for Discardable attributes.
- Maintain Operation::getAttr(StringRef) behavior and so backward compatibility by first checking if the name is a registered inherent attribute, in which case return it from the properties, otherwise return it from the discardable attributes DictionaryAttr.
- Rename Operation::getAttrs() to Operation::getDiscardableAttrs(), and Operation::setAttrs(DictionaryAttr) to Operation::setDiscardableAttrs(DictionaryAttr), so that any uses of the former APIs would be a build breakage and would have to be updated, allowing to audit and update to account for properties. This is the only breakage that downstream clients would incur.
This also presents with the benefit that every existing use of inherent attributes in ODS would immediately be turned into properties, maximizing the benefit of the feature. Although a proper migration to properties involves changing the ODS definition to declare members as “properties” so that they aren’t stored as attributes any more (for example storing the arith.cmpi comparator as a C++ enum inline instead of an IntegerAttr).
We can also very well get there in steps gradually, adopting properties as opt-in first before deprecating and removing the storage of inherent attributes in the DictionaryAttr (through an opt-in at the dialect level for example).
