Using operation properties for better lifetime control in MLIR

MLIR
1

Hello,

In the project I am working on, we would like to have more control over the data that is being used by the operations across the compilation pipeline. At the moment, this data is stored as mlir::DenseElementsAttr / mlir::DenseResourceElementsAttr and additionally wrapped into our own “fancy” attribute. The oversimplified view of the wrapper is:

def MyAttr : Some_Attr<...> {
  let parameters = (ins
     "mlir::ElementsAttr":$theData
  );
}

The problem in such a workflow is that we do not have any control over the lifetime of the objects here (namely, the mlir::Dense*ElementsAttr field) and the problem arises when these attributes are no longer used by any of the operations.

My current “working” idea would be to convert our “fancy” wrapper attribute into a property. This way, I could introduce some manual lifetime management around mlir::DenseResourceElementsAttr at least (since it works on top of a blob resource manager). Due to “property semantics” I would be able to have my wrapper’s destructor call every time an operation is destroyed (i.e. when doing replaceOp and alike).

In principle, this should work fine (far from perfect but it is a start). The catch, however, is that our fancy attribute MyAttr participates in folding. For a random operation, the code like this exists:

// returns "modified" attr
MyAttr doSomethingSpecial(MyAttr attr);

mlir::OpFoldResult MySpecialOp::fold(FoldAdaptor adaptor) {
    auto operands = adaptor.getOperands(); // returns ArrayRef<mlir::Attribute>?
    if (auto attr = mlir::dyn_cast_or_null<MyAttr>(operands[0])) { // oops: cannot cast mlir::Attribute to MyAttr property
        return doSomethingSpecial(attr); // oops: cannot convert MyAttr property to mlir::OpFoldResult
    }
    return nullptr;
}

Now, the two problems here are 1) casting; 2) conversion to fold result. In both of these cases, it seems like the intended behavior is to use attributes. This practically means that I should perform attribute-to-property and property-to-attribute conversion, which I believe is not cheap in my case.

My questions would be:

  • Generally, am I stuck with an XY problem? Maybe there’s a better way to do what I want (that is, have an ability to deallocate mlir::ElementsAttr that goes “out of scope”)?
  • Is there a good example somewhere that shows the path from operation attribute to operation property? (So far I am looking at the tests in MLIR and tablegen’s boilerplate C++ but it’s not much)
  • Given that with properties I would have a simple C++ class, what I would like is to avoid back-and-forth attribute conversions at all costs (in real wrapper type I have an array of things, mlir::Types, boolean fields, etc. etc. - “serializing” is a pain). Is there a way somehow to maintain the “good old” folding logic, yet rely on properties instead? (I am fairly new to the whole folding mechanism so maybe I’m missing something obvious!)

Thank you for the help (and perserverance reading all this!) in advance,
Andrei

2

One way around it is to not use fold for constant folding, but to implement the same thing as a canonicalization pattern: there you’re not bound by the OpFoldResult interface that is mandated by the fold() API.

Another way around this may be to have a shared_ptr/weak_ptr setup: your properties should always be setup as a shared_ptr owned by the Operations, but you can always create an attribute from a weak_ptr to the same object.
Of course this attribute would be meant as an ephemeral object: not to be added as a Attribute on an operation for example.

3

to implement the same thing as a canonicalization pattern

I think this should work. We do rely on .createOrFold() functionality in quite a lot of passes though, but I imagine running a canonicalizer pass after any pass that relied on folding logic before could work out?

Another way around this may be to have a shared_ptr/weak_ptr setup

I guess, what you mean is to create an attribute that holds a pointer? (Not sure how to directly store the std::weak_ptr though).

In theory, this would work as well since we have our own materialization (materializeConstant). So this “weak ptr” would only live during the call to createOrFold<FooOp>().

4

For completeness: We also do need to revisit our folding API. Its long ongoing discussion, but it isn’t a simple change.

2 Likes