Skip to content

FlatMap: add Umpire support #1544

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 8 commits into
base: develop
Choose a base branch
from

Conversation

publixsubfan
Copy link
Contributor

Summary

  • Adds an Allocator wrapper class to disambiguate between integer arguments and allocator ID

    This is a speculative change, and I'd like to hear feedback from others: my biggest complaint/regret about the Array constructor interface is that we aren't fully drop-in compatible with the std::vector constructor interface, as we've sprinkled in allocator_id function arguments haphazardly. I'm hoping this might provide an alternative path forward that we can test out in FlatMap for now, and maybe bring it over into Array proper once fleshed out.

  • Adds initial Umpire allocation support for FlatMap, supporting the following operations:

    • Creating a FlatMap within a custom Umpire allocator
    • Copying FlatMap between two different allocators, including device-based allocators
    • Destroying FlatMap created within a Umpire device-based allocator.
  • Fix move assignment operator by returning reference to current instance

@publixsubfan publixsubfan added Core Issues related to Axom's 'core' component GPU Issues related to GPU development labels Apr 8, 2025
MapType test_map_host = MapType(test_map_device, axom::Allocator {host_alloc_id});

// Check that maps are equivalent with the same allocator ID.
EXPECT_EQ(test_map_host.getAllocatorID(), host_alloc_id);
Copy link
Contributor

@bmhan12 bmhan12 Apr 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For consistency, I think it might make sense to make this getter a getAllocator() to compare Allocator's, and prevent the ability to access the int id's altogether.

@bmhan12
Copy link
Contributor

bmhan12 commented Apr 9, 2025

Adds an Allocator wrapper class to disambiguate between integer arguments and allocator ID

I think this makes a lot of sense and helps guard against coding errors.
For reference, allocating an axom::Array using Umpire typically looks like some form of the following:

    axom::Array<int> my_array(
      size,
      capacity,
      allocator_id);

size, capacity, and allocator_id are all of type IndexType/int.

Copy link
Member

@kennyweiss kennyweiss left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like this idea @publixsubfan -- it feels a lot cleaner than our current approach.

I'm eager to see what this would look like when applied to axom::Array and axom::ArrayView

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like this idea!

int get() const { return m_id; }

private:
int m_id {0};
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this be 0 by default? Or perhaps axom::INVALID_ALLOCATOR_ID
(which evaluates to -1 and is not a valid umpire allocator ID).


inline void setSentinel(axom::ArrayView<GroupBucket> metadata)
{
#if defined(AXOM_USE_UMPIRE) && defined(AXOM_USE_CUDA)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does/Will this also support hip? Or does hip not need extra handling?
If it's the latter, please add a comment.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, it's the latter case since with HIP we can access device memory directly from the host. I'll add a comment.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's true on some Hip platforms, but not all.

return;
}

#if defined(AXOM_USE_UMPIRE) && defined(AXOM_USE_CUDA)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similar comment about hip here (and presumably elsewhere in this file)

Comment on lines +706 to +707
// Copy from the host to the device.
MapType test_map_device = MapType(test_map, axom::Allocator {device_alloc_id});
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suppose this would be our own idiomatic way to copy, but it's interesting!

Comment on lines +711 to +712
EXPECT_EQ(test_map_device.bucket_count(), test_map.bucket_count());
EXPECT_EQ(test_map_device.size(), test_map.size());
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to make sure I understand -- test_map_device.size() is on the host, right?
Would we ever want the size to be on the device? If so, what would that look like?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In this case, yes -- after this PR, I'm planning on pushing out another follow-up with a FlatMapView type for device-side accesses.

It'll be similar to the Array/ArrayView dichotomy, where the device accesses the data through the "view" type that is value-captured in the lambda.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Core Issues related to Axom's 'core' component GPU Issues related to GPU development
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants