Memory Spaces

Related Links:

1. Main Issue: Memory Spaces and support for different kinds of memory #195
   • Subsection:1 Issue Determining the relationship between Teams, Contexts, and Spaces #455
2. OLD PROPOSAL: Initial High-level Ideas and APIs
3. Basic example
4. August, 2020 - OpenSHMEM spec Meeting Updates
5. October, 2020 - OpenSHMEM spec Meeting Updates
6. User examples:
   • Simple distributed sort
7. WG Meeting Notes
8. Specification Draft

What is Memory Space?

Spaces are Team-based symmetric heaps. In OpenSHMEM-1.5, we have a single symmetric heap implicitly created by the OpenSHMEM implementation based on the size as input from the users through the SHMEM_SYMMETRIC_SIZE environment variable.

What does Memory Space address?

1. Emerging system architectures support different kinds of memory, hence we need to allow users to provide more information to OpenSHMEM implementation while creating a symmetric heap. Information might include memory traits like symmetric heap size, type of memory, NUMA details, hugepage usage, and information specific to special type of memory like FAM, HBM.
2. Avoid unnecessarily allocating symmetric data objects across the global PEs, while the application is aware of symmetric data object within the sub-set of global PEs part of an OpenSHMEM Team.
3. Allow using the memory buffer used as a symmetric heap between multiple OpenSHMEM jobs
4. Allow using OpenSHMEM with other programming models and languages without the notion of PGAS style of programming, specifically no symmetric heap

What are different use-cases addressed?

1. Move data through RMA/AMO and Collectives between two different memory kinds

• Device to Host and Host to Device specifically using collectives like all-reduce
• Use FAM as an external memory and allow moving data into and from FAM from an on-node memory kind like DDR/HBM

2. Share Memory buffer between two different teams

• Row/Column 2D Teams sharing the data between themselves
• Duplicate Teams for performing multi-threaded collectives

3. Allow using the memory buffer used as a symmetric heap between multiple OpenSHMEM jobs

• Using FAM as flat memory shared between independent OpenSHMEM jobs

Proposal Overview

The proposal is divided into three sub-categories.

• Category:1 - Decide the core space syntax and semantics
• Category:2 - Introduce support for different kinds of memory, specifically FAM and HBM
• Category:3 - Extend FAM usage as a flat-memory

Category:1

Refer Presentation for more details. The core Memory Space syntax and semantics includes deciding the following:

1. Team-Context-Space relationship
2. Relationship of the spaces associated with the parent team through which a child team is split
3. Memory management operations like malloc, calloc, align, and free

To support the above use-cases, we need to support the following team - space relationships:

1. Multiple spaces per team (example: space for each memory kind)
2. Multiple teams per space (example: multithreaded collectives)

Parent-Child Team relationship with respect to the spaces associated with the parent team

1. Child teams should be able to access the parent team's space
2. By default, contexts associated with a particular team should be able to access all the associated spaces
3. As an optimization for RMA and AMO, contexts associated with a team can select a particular space on which the data movement operation can be performed. This would avoid using lookup tables or some-sort of space identification process during the RMA and AMO operation. This is not necessary for collectives, as the lookup operation can be performed once for all the internal RMA and AMO involved in the collective operation.

To achieve the above requirements, we can use the following APIs in OpenSHMEM,

Memory Space Objects

typedef shmem_space_t void *;
typedef struct {
   size_t sheap_size;
} shmem_space_config_t;

1. shmem_space_t is a new opaque handle for mapping the space
2. Different traits of a space will be provided using the shmem_space_config_t parameter
3. shmem_space_config_t will be updated with other traits when sub-category:2 of this proposal is discussed. For now, sheap_size is sufficient

Memory Space Creation Operation

typedef struct {
   // existing team configuration parameters
   int                    num_contexts;
   // new space-specific configurations
   shmem_space_t        * spaces;
   shmem_space_config_t * space_configs;
   shmem_team_t         * team_select;
   int                    num_spaces;
} shmem_team_config_t;

1. Spaces are created and destroyed with the Teams during the shmem_team_split_strided and shmem_team_split_2d operations
2. A Space can be created on only a single Team
3. All Teams split from a parent team, inherits the spaces accessible on the parent Team. Spaces from the parent team can be used only for RMA, AMO, and Collectives. They cannot be used for memory allocation operations.
4. The team_select input in the configuration denotes the team to which the space belongs. For strided-split, there is no need for this configuration option, while for a 2D-split there can be only one parent to the space. Hence, users need to set, whether the space belongs to the x-axis/y-axis team
5. EXTRA (separate proposal): We need to inherit a new Team-split API called shmem_team_duplicate to provide similar access guarantees on the spaces to the child Team as the parent Team

Cons:

1. We are introducing a new parent-child team hierarchy. To use the parent-teams space for RMA/AMO in the child-team, the parent-team has to be kept active through-out the lifetime of the child-team.
2. We need to introduce a new shmem_team_duplicate to share a space between two child-teams from the same-parent. Previously, we allowed the following sequence of creation:

parent_team(team_split)->child_1 
parent_team(team_split)->child_2
// where child_1 and child_2 have the same set of PEs and can be used for multithreaded collectives.

Now, we need the following sequence:

parent_team(team_split)->child_1->(team_duplicate)child_2

Space-based Context Creation

shmem_ctx_config_t {
  shmem_team_t  team;
  shmem_space_t space;
};
int shmem_ext_ctx_create(IN long options, IN shmem_ctx_config_t *config, OUT shmem_ctx_t *ctx);

1. Optimization to avoid space lookups during RMA and AMOs
2. Any context created using shmem_ext_ctx_create can be used only on the associated space buffers
3. Context created using shmem_team_create_ctx and shmem_ctx_create would inherit access to all spaces asssociated with the corresponding teams. Performance on these variables would depend on the number of attached spaces.

Memory Management Routines

void *shmem_space_malloc(IN shmem_space_t space, IN size_t size);
void *shmem_space_align(IN shmem_space_t space, IN size_t alignment, IN size_t size); 

1. Malloc and Align operations will take space handle as one of the input in shmem_space_malloc and shmem_space_align.
2. Space attach would automatically perform a barrier-like operation on the associated parent team
3. Allocation operations are always symmetric and follows the existing shmem_malloc and shmem_align semantics
4. There is no need for a space-based realloc and free operations, as the realloc is always with the same space and free can determine the space associated internally in the implementation

Category:2 (TODO)

Category:3 (TODO)