Take ownership of tensor data ptr

C++
1

Hey!

I currently have most of my application abstracted into “modules”. My current problem is that when I get a at::Tensor tensor from the neural network output in C++, I want to be able to take ownership of the tensor.data_ptr() using my own internal C++ class and delete the interface at::Tensor. Is this possible?

2

Hi,

This is not possible as of today because you wouldn’t know how to free that memory.
Is there a specific reason why you can keep the Tensor object around to keep the buffer alive?

2 Likes
3

This is an interesting point. I’ve been using global variables or copies as workaround.

4

I architected my codebase so that modules don’t know the data types of other modules it may reference. The only structs shared between modules are my own internal structs.

So I have a struct called ManagedMemory which handles memory management. So what I usually do is instantiate a ManagedMemory(dataSize, dataPtr, memoryType=cuda/cpu) inside the ModulePytorch from the inference output and then send that through pubsub to another module. ManagedMemory is now the struct which owns the data.

Just a thought which may work (tell me if this is doable):

Maybe an alternative would be to up the refCount artificially by 1 and then pass in a custom function into my ManagedMemory which will decrease the refCount when it’s not being used anymore.
ManagedMemory(dataSize, dataPtr, memoryType=cuda/cpu, [&]() { // Decrease Tensor Refcount })
Then that means the data is still managed by Module PyTorch but will have a floating pointer until the refCount falls to 0.

I just don’t like the idea of doing this. I’m more of favor of telling tensors to release ownership to the pointers.

5

I just don’t like the idea of doing this. I’m more of favor of telling tensors to release ownership to the pointers.

The problem is not with the ownership but with how will you know how to free the memory?
We use different allocators depending on what you do (regular cpu memory, regular gpu memory, shared cpu memory, paged cpu memory,…). And you won’t be able to call just free or delete on the pointer to clean it properly!

Do you think it would be possible for your ManagedMemory to have an optional Tensor field that will make sure it keep the dataPtr alive? And when the ManagedMemory is freed, the Tensor will go away and free the memory.

6

Yeah, but if we know the type of memory, the size, and have the pointer, should I be able to just call a release() and it will automatically give up ownership of whatever memory type it was allocated with? It will be up to me as the developer to now do the proper de-allocations elsewhere depending on what the tensor memory type was. A little unsafe but I guess that’s what happens when you go against the grain.

But yeah, I can add some more code to my ManagedMemory to do that and make the tensor live until the duration of the ManagedMemory. Not doing the optional Tensor field method (because the point of ManagedMemory was to have basic dependencies and nothing additional) but some other way.

Thanks for helping me out albanD, it’s really appreciated.

7

What you describe is a bit like DLPack (DLManagedTensor on the C side), maybe that can help you achieve what you need.
In particular DLManagedTensor* toDLPack(const Tensor& src);, from ATen/DLConvertor.h (mind the o at the end).

Best regards

Thomas

8

Still couldn’t figure out how to take ownership of the pointer since they store the Tensor internally, but I did what I mentioned above and modified by ManagedMemory to have a custom deallocation func.

{
        at::Tensor a = at::randn({50, 50, 50});
        int tensorDataSize = std::accumulate(a.sizes().begin(), a.sizes().end(), 1, std::multiplies<int>()) * sizeof(float);
        DLManagedTensor * dlPackTensor = at::toDLPack(a);
        ManagedMemory managedMemory((char *) dlPackTensor->dl_tensor.data, tensorDataSize, CPU, true);
        managedMemory.deallocationFunc = [&]() { dlPackTensor->deleter(dlPackTensor); };
}