Add custom 'tags' to the memory pool

[igchor]

Add custom 'tags' to the memory pool

Rationale

In UR, there are a few use-cases that we can optimize by allowing the user to associated a 'tag' with the memory pool (and to later, query the 'tag'):

1. Calls to zeMemGetAllocProperties (quite expensive) in L0 adapter: https://github.com/oneapi-src/unified-runtime/blob/main/source/adapters/level_zero/memory.cpp#L51
2. Quite convoluted logic for finding ur_usm_pool_handle_t for a given ptr: https://github.com/oneapi-src/unified-runtime/blob/d61cf6680c7f620bd5092e2bb22524106fbdef63/source/adapters/level_zero/usm.cpp#L637

Description

Add one extra void *tag argument to umfPoolCreate and allow querying it later. Tag can point to a user-provided structure that user is responsible for freeing.

In UR we can create a custom structure struct pool_tag { ur_memory_type type; ur_usm_pool_handle_t urPool; } and store a pointer to that structure as a tag for each umf pool. This will make calls to zeMemGetAllocProperties unnecessary and logic for finding ur_usm_pool_handle_t will be just ((pool_tag *) umfPoolGetTag(umfPoolByPtr(ptr)))-> urPool

API Changes

umf_result_t umfPoolCreate(const umf_memory_pool_ops_t *ops,
                           umf_memory_provider_handle_t provider, void *params,
                           void *tag, // <- NEW PARAMETER
                           umf_pool_create_flags_t flags,
                           umf_memory_pool_handle_t *hPool);
                           
                           
void *umfPoolGetTag(umf_memory_pool_handle_t hPool);                           

Implementation details

Store provided tag in the pool structure.

Meta

[vinser52]

I agree with the problem statement, but we need to discuss API design (as I mentioned in the previous issue).
In the example you provided the memory type (e.g. HOST_PINNED, SHARED, DEVICE) should be a "regular" memory property provided by the API we are going to expose to get allocation properties.

[bratpiorka]

I partially agree with @vinser52 - in the first case we could use either use (not existing yet) API to query memory properties or just use umfPoolByPtr() to get the pool - the type of memory is provided by the user during the pool creation.

So the real problem here is to optimize the 2nd problem - get the UR pool from the UMF pool and this "tag" mechanism should help.

[igchor]

the type of memory is provided by the user during the pool creation.

That's true but there is no way to get this information back from the pool. You can call umfPoolGetMemoryProvider() but this will return umf_memory_provider_handle_t which is a pointer to UMF structure, not to the structure that user passed during provider creation. So event for this case, we need a new API (either 'tag' or some new API for querying memory properties).

[igchor]

the type of memory is provided by the user during the pool creation.

That's true but there is no way to get this information back from the pool. You can call umfPoolGetMemoryProvider() but this will return umf_memory_provider_handle_t which is a pointer to UMF structure, not to the structure that user passed during provider creation. So event for this case, we need a new API (either 'tag' or some new API for querying memory properties).

Also, the application can use some 'cutom' memory type. For example, Unified Runtime has a concept of SharedReadOnly (read only on the device) memory - is this something that we want to define in UMF?

[vinser52]

Also, the application can use some 'cutom' memory type. For example, Unified Runtime has a concept of SharedReadOnly (read only on the device) memory - is this something that we want to define in UMF?

My main point was that we need to design the proper API. The problem is clear. As a next step, I am going to capture the use cases we have so far. After that we can brainstorm on the right API.

In the example you provided the required information is already somehow stored inside UMF structures (pool or provider) for example SharedReadOnly flag is a memory provider config parameter - user should not be required to duplicate it as a custom Tag.

[bratpiorka]

Agree that we could have a brainstorm about querying properties, but tags could be really useful to get e.g. UR or other user defined objects associated with UMF pools and this API is really simple - I think we could implement this

[igchor]

Also, the application can use some 'cutom' memory type. For example, Unified Runtime has a concept of SharedReadOnly (read only on the device) memory - is this something that we want to define in UMF?

My main point was that we need to design the proper API. The problem is clear. As a next step, I am going to capture the use cases we have so far. After that we can brainstorm on the right API.

In the example you provided the required information is already somehow stored inside UMF structures (pool or provider) for example SharedReadOnly flag is a memory provider config parameter - user should not be required to duplicate it as a custom Tag.

Sure, I'm fine with having different APIs to solve 1. and 2. In this issue I just proposed umfPoolGetTag because it's quite simple and generic enough to cover those 2 use cases (and I don't have a better idea how to cover the second use case).

[vinser52]

Ahh, ok, I am not against umfPoolGetTag API. I think it is orthogonal to memory properties API.

But here are my questions/feedback:

1. Should we introduce a dedicated API to assign Tag to the pool instead of adding an additional parameter to the umfPoolCreate API? My first concern is that tag parameter is not mandatory and I suppose it will be NULL in most cases. the second concern is backward compatibility: I understand that our current UMF version is 0.9 but anyway we already have customers and umfPoolCreate is our main entry point. Extending it would brake customer codes.
2. Should we add support of Tags for memory providers as well.
3. Naming... I know it is always hard to choose the right name. Is tag is the best name? Any other ideas? I thought about metadata but it also might be confusing.

[igchor]

Ahh, ok, I am not against umfPoolGetTag API. I think it is orthogonal to memory properties API.

But here are my questions/feedback:
...

1. Yes, we can add a dedicated function. Should we make set/get thread safe? I don't have a use case for that, so we don't have to for now.
2. This is less useful I think. The only way to get a provider from ptr is through umfPoolByPtr(so the user can just store whatever he wants in the pool's tag). If you have a provider handle from somewhere else and need to query it, you can probably just change the code to pass a struct that contains required info and the handle (using umfProviderGetTag would require the function to know the tag type anyway).
3. I don't have better ideas. We could name it something like 'userData' but I prefer 'tag'.

[vinser52]

Another question what if tag is tried to be assigned multiple time to the same pool. E.g. one client creates pool and assign tag and than another client get the poo via umfPoolByPtr() and tries to assign another tag.

[igchor]

Another question what if tag is tried to be assigned multiple time to the same pool. E.g. one client creates pool and assign tag and than another client get the poo via umfPoolByPtr() and tries to assign another tag.

That's why I think passing the tag in create would be a bit better but for the set function I think we can allow it - as long as we describe what are the guarantees (e.g. no thread-safety, etc.).

[vinser52]

With the set function it is not clear what to do in case some tag is already set. Since tag is allocated by a client we cannot free memory by ourselves.

Regarding thread-safety, it is not clear to me why we cannot support it?

[igchor]

With the set function it is not clear what to do in case some tag is already set. Since tag is allocated by a client we cannot free memory by ourselves.

Regarding thread-safety, it is not clear to me why we cannot support it?

Well, that's true. User should be responsible for freeing it. It may very well be the case that 'tag' is just a magic value (i.e. uint64_t id casted to void* in which case no freeing is necessary).

There's no reason we cannot support it, I was just saying that I don't see a use case for it right now. If we do, then the set function should probably return the old value of the tag so that the user can release it if necesary.