Can I have custom data owned (and freed) by an Operation?

Hi,

Is there a way to have an operation own some custom data (similar to an Attribute), and have the data destroyed when the operation is erased.

I would like to have on my dialect a ConstantOp with a huge buffer (usually hold millions of elements), and have this buffer dealocated when the op is erased.

To give more context, I am working on implementing my own math Dialect on tensors which manipulate huge constants, and supports folding.

I have implemented constants this way:

def Toy_ConstantOp : Toy_Op<"constant", [NoSideEffect, ConstantLike]> {
  let arguments = (ins ElementsAttr:$value);
  let results = (outs Toy_Tensor:$result);
  let hasFolder = 1;
}

When performing folding (For Example folding an Add op with 2 constants operands into a new constant), I ran into 2 problems:

  • Iterating over DenseElementsAttr for huge tensors is slow
  • Once the constants op for the operands are optimized away, I don’t need the data anymore. But it’s only deallocated once the Context is destroyed. That’s hundreds (or even thousands) of Mbs of memory wasted.

I was able to solve the first problem by using my own custom attribute class:

let arguments = (ins ToyRawBufferAttr:$value);

Inside ToyRawBufferAttr, I am manually allocating a huge linear buffer, which is a lot faster to manipulate compared to DenseElementsAttr / APFloat.

But again the memory is only deallocated when the Context is destroyed.

Looking at the source, I have seen the Operation has it’s own ~Operation() destructor being called, but I couldn’t figure out if there is a way to extend / override this for my own operations.

Thanks.

No this isn’t possible.

We also have this problem (everyone does :wink: ), and the current direction is the recent addition of “resources” to MLIR: this allows to externalize the storage of the data and decouple their lifetime from the Context.
This does not provide directly an automatic deallocation when unused, but it makes it possible to “garbage collect” these easily by walking the IR and marking only the resources used.

I was looking for doc on this, but it seems like its not there yet (@River707 ?), here is a relevant revision for example: âš™ D130021 [mlir] Add a generic DialectResourceBlobManager to simplify resource blob management

We should implement some more integration and example to achieve exactly what you’re asking for here, this is commonly needed.

1 Like

See [RFC] Adding a generalized “config” section to MLIR files for resource discussion.

Thanks for your answer, resources has most of what I need, except for automatic deallocation. But walking over the IR to deallocate unused resources feels like a good compromise.

1 Like

Hi, I’d like to give some updates about this.
I used resources to implement constants, and a garbage collector to free unused memory. And it does the job alright, but I had lots of issues along the road:

  • Multiple Modules share the same context, but when creating my custom ResourceAttribute, it is bound to a context, not a module, so during GC, I need to have access to all the modules and walk through all of them.
  • I need to use locks or other synchronization primitives to read / write the data inside my ResourcesManager since it can be accessed by multiple threads.
  • Calling GC after every pass is not enough. On Some experiments I go from 1.7GB to 3.4GB memory peak because of unused constants. I am able to keep the memory peak ~1.7GB by managing the constants with a refcount. But I decrease the refcount inside the folders and OpRewritePatterns, which is definitely the wrong place to do this. But I couldn’t find any other ways.

Are they any new features upstream that would help solve my problems ? If not, I have some ideas and I’d be happy to contribute.

I don’t think of anything upstream that can help with this right now, I’m curious about your ideas though!

Using a GC seems to be a terrible idea, especially for multi-thread uses. It’s slow, requires lock, wastes too much memory, and there is no real need for it. Having a refcount on the resource is a much better solution, when an operation is created, we increment it, and when an operation is erased, we decrement it, and the memory is released when the refcount reaches 0, it should be as simple as this.
The GC is there only because the current design of MLIR doesn’t allow us to update properly such a refcount.

It would be complicated to make MLIR aware of this refcount, especially since a resource is not always linked to an attribute or an operation, and also we want the user to have full control over the resource lifetime.
Instead, a more practical solution would be to create 2 new native Interfaces: OnOperationCreation (called right after an operation is built), and OnOperationDestruction (called when an operation is erased).
The goal is to help the user being able to manage more easily the allocation / destruction of resources thanks to IR changes callbacks.

2 Likes