Skip to content

Commit 8100bd5

Browse files
hansangbaehansangbae
authored
[OpenMP] 6.0 (TR11) Memory Management Update (#97106)
TR11 introduced changes to support target memory management in a unified way by defining a series of API routines and additional traits. Host runtime is oblivious to how actual memory resources are mapped when using the new API routines, so it can only support how the composed memory space is maintained, and the offload backend must handle which memory resources are actually used to allocate memory from the memory space. Here is summary of the implementation. * Implemented 12 API routines to get/mainpulate memory space/allocator. * Memory space composed with a list of devices has a state with resource description, and runtime is responsible for maintaining the allocated memory space objects. * Defined interface with offload runtime to access memory resource list, and to redirect calls to omp_alloc/omp_free since it requires backend-specific information. * Value of omp_default_mem_space changed from 0 to 99, and omp_null_mem_space took the value 0 as defined in the language. * New allocator traits were introduced, but how to use them is up to the offload backend. * Added basic tests for the new API routines.
1 parent acc6bcd commit 8100bd5

File tree

12 files changed

+1006
-20
lines changed

12 files changed

+1006
-20
lines changed

‎openmp/runtime/src/dllexports

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -532,6 +532,18 @@ kmp_set_disp_num_buffers 890
532532
omp_get_device_num 896
533533
omp_init_allocator 897
534534
omp_destroy_allocator 898
535+
omp_get_devices_memspace 810
536+
omp_get_device_memspace 811
537+
omp_get_devices_and_host_memspace 812
538+
omp_get_device_and_host_memspace 813
539+
omp_get_devices_all_memspace 814
540+
omp_get_devices_allocator 815
541+
omp_get_device_allocator 816
542+
omp_get_devices_and_host_allocator 817
543+
omp_get_device_and_host_allocator 818
544+
omp_get_devices_all_allocator 819
545+
omp_get_memspace_num_resources 820
546+
omp_get_submemspace 821
535547
%ifndef stub
536548
__kmpc_set_default_allocator
537549
__kmpc_get_default_allocator
@@ -592,6 +604,7 @@ kmp_set_disp_num_buffers 890
592604
llvm_omp_target_host_mem_space DATA
593605
llvm_omp_target_shared_mem_space DATA
594606
llvm_omp_target_device_mem_space DATA
607+
omp_null_mem_space DATA
595608

596609
%ifndef stub
597610
# Ordinals between 900 and 999 are reserved

‎openmp/runtime/src/include/omp.h.var

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -339,7 +339,13 @@
339339
omp_atk_fallback = 5,
340340
omp_atk_fb_data = 6,
341341
omp_atk_pinned = 7,
342-
omp_atk_partition = 8
342+
omp_atk_partition = 8,
343+
omp_atk_pin_device = 9,
344+
omp_atk_preferred_device = 10,
345+
omp_atk_device_access = 11,
346+
omp_atk_target_access = 12,
347+
omp_atk_atomic_scope = 13,
348+
omp_atk_part_size = 14
343349
} omp_alloctrait_key_t;
344350

345351
typedef enum {
@@ -350,7 +356,7 @@
350356
omp_atv_serialized = 5,
351357
omp_atv_sequential = omp_atv_serialized, // (deprecated)
352358
omp_atv_private = 6,
353-
omp_atv_all = 7,
359+
omp_atv_device = 7,
354360
omp_atv_thread = 8,
355361
omp_atv_pteam = 9,
356362
omp_atv_cgroup = 10,
@@ -361,7 +367,11 @@
361367
omp_atv_environment = 15,
362368
omp_atv_nearest = 16,
363369
omp_atv_blocked = 17,
364-
omp_atv_interleaved = 18
370+
omp_atv_interleaved = 18,
371+
omp_atv_all = 19,
372+
omp_atv_single = 20,
373+
omp_atv_multiple = 21,
374+
omp_atv_memspace = 22
365375
} omp_alloctrait_value_t;
366376
#define omp_atv_default ((omp_uintptr_t)-1)
367377

@@ -387,6 +397,7 @@
387397
extern __KMP_IMP omp_allocator_handle_t const llvm_omp_target_device_mem_alloc;
388398

389399
typedef omp_uintptr_t omp_memspace_handle_t;
400+
extern __KMP_IMP omp_memspace_handle_t const omp_null_mem_space;
390401
extern __KMP_IMP omp_memspace_handle_t const omp_default_mem_space;
391402
extern __KMP_IMP omp_memspace_handle_t const omp_large_cap_mem_space;
392403
extern __KMP_IMP omp_memspace_handle_t const omp_const_mem_space;
@@ -422,7 +433,8 @@
422433
typedef enum omp_memspace_handle_t
423434
# endif
424435
{
425-
omp_default_mem_space = 0,
436+
omp_null_mem_space = 0,
437+
omp_default_mem_space = 99,
426438
omp_large_cap_mem_space = 1,
427439
omp_const_mem_space = 2,
428440
omp_high_bw_mem_space = 3,
@@ -463,6 +475,20 @@
463475
extern void __KAI_KMPC_CONVENTION omp_free(void *ptr, omp_allocator_handle_t a);
464476
# endif
465477

478+
/* OpenMP TR11 routines to get memory spaces and allocators */
479+
extern omp_memspace_handle_t omp_get_devices_memspace(int ndevs, const int *devs, omp_memspace_handle_t memspace);
480+
extern omp_memspace_handle_t omp_get_device_memspace(int dev, omp_memspace_handle_t memspace);
481+
extern omp_memspace_handle_t omp_get_devices_and_host_memspace(int ndevs, const int *devs, omp_memspace_handle_t memspace);
482+
extern omp_memspace_handle_t omp_get_device_and_host_memspace(int dev, omp_memspace_handle_t memspace);
483+
extern omp_memspace_handle_t omp_get_devices_all_memspace(omp_memspace_handle_t memspace);
484+
extern omp_allocator_handle_t omp_get_devices_allocator(int ndevs, const int *devs, omp_memspace_handle_t memspace);
485+
extern omp_allocator_handle_t omp_get_device_allocator(int dev, omp_memspace_handle_t memspace);
486+
extern omp_allocator_handle_t omp_get_devices_and_host_allocator(int ndevs, const int *devs, omp_memspace_handle_t memspace);
487+
extern omp_allocator_handle_t omp_get_device_and_host_allocator(int dev, omp_memspace_handle_t memspace);
488+
extern omp_allocator_handle_t omp_get_devices_all_allocator(omp_memspace_handle_t memspace);
489+
extern int omp_get_memspace_num_resources(omp_memspace_handle_t memspace);
490+
extern omp_memspace_handle_t omp_get_submemspace(omp_memspace_handle_t memspace, int num_resources, int *resources);
491+
466492
/* OpenMP 5.0 Affinity Format */
467493
extern void __KAI_KMPC_CONVENTION omp_set_affinity_format(char const *);
468494
extern size_t __KAI_KMPC_CONVENTION omp_get_affinity_format(char *, size_t);

‎openmp/runtime/src/include/omp_lib.F90.var

Lines changed: 104 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,12 @@
145145
integer (kind=omp_alloctrait_key_kind), parameter, public :: omp_atk_fb_data = 6
146146
integer (kind=omp_alloctrait_key_kind), parameter, public :: omp_atk_pinned = 7
147147
integer (kind=omp_alloctrait_key_kind), parameter, public :: omp_atk_partition = 8
148+
integer (kind=omp_alloctrait_key_kind), parameter, public :: omp_atk_pin_device = 9
149+
integer (kind=omp_alloctrait_key_kind), parameter, public :: omp_atk_preferred_device = 10
150+
integer (kind=omp_alloctrait_key_kind), parameter, public :: omp_atk_device_access = 11
151+
integer (kind=omp_alloctrait_key_kind), parameter, public :: omp_atk_target_access = 12
152+
integer (kind=omp_alloctrait_key_kind), parameter, public :: omp_atk_atomic_scope = 13
153+
integer (kind=omp_alloctrait_key_kind), parameter, public :: omp_atk_part_size = 14
148154

149155
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_default = -1
150156
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_false = 0
@@ -154,7 +160,7 @@
154160
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_serialized = 5
155161
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_sequential = omp_atv_serialized
156162
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_private = 6
157-
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_all = 7
163+
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_device = 7
158164
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_thread = 8
159165
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_pteam = 9
160166
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_cgroup = 10
@@ -166,6 +172,10 @@
166172
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_nearest = 16
167173
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_blocked = 17
168174
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_interleaved = 18
175+
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_all = 19
176+
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_single = 20
177+
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_multiple = 21
178+
integer (kind=omp_alloctrait_val_kind), parameter, public :: omp_atv_memspace = 22
169179

170180
integer (kind=omp_allocator_handle_kind), parameter, public :: omp_null_allocator = 0
171181
integer (kind=omp_allocator_handle_kind), parameter, public :: omp_default_mem_alloc = 1
@@ -180,7 +190,8 @@
180190
integer (kind=omp_allocator_handle_kind), parameter, public :: llvm_omp_target_shared_mem_alloc = 101
181191
integer (kind=omp_allocator_handle_kind), parameter, public :: llvm_omp_target_device_mem_alloc = 102
182192

183-
integer (kind=omp_memspace_handle_kind), parameter, public :: omp_default_mem_space = 0
193+
integer (kind=omp_memspace_handle_kind), parameter, public :: omp_null_mem_space = 0
194+
integer (kind=omp_memspace_handle_kind), parameter, public :: omp_default_mem_space = 99
184195
integer (kind=omp_memspace_handle_kind), parameter, public :: omp_large_cap_mem_space = 1
185196
integer (kind=omp_memspace_handle_kind), parameter, public :: omp_const_mem_space = 2
186197
integer (kind=omp_memspace_handle_kind), parameter, public :: omp_high_bw_mem_space = 3
@@ -802,6 +813,97 @@
802813
logical (kind=omp_logical_kind) omp_in_explicit_task
803814
end function omp_in_explicit_task
804815

816+
function omp_get_devices_memspace(ndevs, devs, memspace)
817+
use omp_lib_kinds
818+
integer(omp_memspace_handle_kind) :: omp_get_devices_memspace
819+
integer, intent(in) :: ndevs
820+
integer, intent(in) :: devs(*)
821+
integer(omp_memspace_handle_kind), intent(in) :: memspace
822+
end function omp_get_devices_memspace
823+
824+
function omp_get_device_memspace(dev, memspace)
825+
use omp_lib_kinds
826+
integer(omp_memspace_handle_kind) :: omp_get_device_memspace
827+
integer, intent(in) :: dev
828+
integer(omp_memspace_handle_kind), intent(in) :: memspace
829+
end function omp_get_device_memspace
830+
831+
function omp_get_devices_and_host_memspace(ndevs, devs, memspace)
832+
use omp_lib_kinds
833+
integer(omp_memspace_handle_kind) :: &
834+
omp_get_devices_and_host_memspace
835+
integer, intent(in) :: ndevs
836+
integer, intent(in) :: devs(*)
837+
integer(omp_memspace_handle_kind), intent(in) :: memspace
838+
end function omp_get_devices_and_host_memspace
839+
840+
function omp_get_device_and_host_memspace(dev, memspace)
841+
use omp_lib_kinds
842+
integer(omp_memspace_handle_kind) :: &
843+
omp_get_device_and_host_memspace
844+
integer, intent(in) :: dev
845+
integer(omp_memspace_handle_kind), intent(in) :: memspace
846+
end function omp_get_device_and_host_memspace
847+
848+
function omp_get_devices_all_memspace(memspace)
849+
use omp_lib_kinds
850+
integer(omp_memspace_handle_kind) :: omp_get_devices_all_memspace
851+
integer(omp_memspace_handle_kind), intent(in) :: memspace
852+
end function omp_get_devices_all_memspace
853+
854+
function omp_get_devices_allocator(ndevs, devs, memspace)
855+
use omp_lib_kinds
856+
integer(omp_allocator_handle_kind) :: omp_get_devices_allocator
857+
integer, intent(in) :: ndevs
858+
integer, intent(in) :: devs(*)
859+
integer(omp_memspace_handle_kind), intent(in) :: memspace
860+
end function omp_get_devices_allocator
861+
862+
function omp_get_device_allocator(dev, memspace)
863+
use omp_lib_kinds
864+
integer(omp_allocator_handle_kind) :: omp_get_device_allocator
865+
integer, intent(in) :: dev
866+
integer(omp_memspace_handle_kind), intent(in) :: memspace
867+
end function omp_get_device_allocator
868+
869+
function omp_get_devices_and_host_allocator(ndevs, devs, memspace)
870+
use omp_lib_kinds
871+
integer(omp_allocator_handle_kind) :: &
872+
omp_get_devices_and_host_allocator
873+
integer, intent(in) :: ndevs
874+
integer, intent(in) :: devs(*)
875+
integer(omp_memspace_handle_kind), intent(in) :: memspace
876+
end function omp_get_devices_and_host_allocator
877+
878+
function omp_get_device_and_host_allocator(dev, memspace)
879+
use omp_lib_kinds
880+
integer(omp_allocator_handle_kind) :: &
881+
omp_get_device_and_host_allocator
882+
integer, intent(in) :: dev
883+
integer(omp_memspace_handle_kind), intent(in) :: memspace
884+
end function omp_get_device_and_host_allocator
885+
886+
function omp_get_devices_all_allocator(memspace)
887+
use omp_lib_kinds
888+
integer(omp_allocator_handle_kind) :: &
889+
omp_get_devices_all_allocator
890+
integer(omp_memspace_handle_kind), intent(in) :: memspace
891+
end function omp_get_devices_all_allocator
892+
893+
function omp_get_memspace_num_resources(memspace)
894+
use omp_lib_kinds
895+
integer omp_get_memspace_num_resources
896+
integer(omp_memspace_handle_kind), intent(in) :: memspace
897+
end function omp_get_memspace_num_resources
898+
899+
function omp_get_submemspace(memspace, num_resources, resources)
900+
use omp_lib_kinds
901+
integer(omp_memspace_handle_kind) omp_get_submemspace
902+
integer(omp_memspace_handle_kind), intent(in) :: memspace
903+
integer, intent(in) :: num_resources
904+
integer, intent(in) :: resources(*)
905+
end function omp_get_submemspace
906+
805907
! ***
806908
! *** kmp_* entry points
807909
! ***

‎openmp/runtime/src/include/omp_lib.h.var

Lines changed: 117 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,18 @@
151151
parameter(omp_atk_pinned=7)
152152
integer(kind=omp_alloctrait_key_kind)omp_atk_partition
153153
parameter(omp_atk_partition=8)
154+
integer(kind=omp_alloctrait_key_kind)omp_atk_pin_device
155+
parameter(omp_atk_pin_device=9)
156+
integer(kind=omp_alloctrait_key_kind)omp_atk_preferred_device
157+
parameter(omp_atk_preferred_device=10)
158+
integer(kind=omp_alloctrait_key_kind)omp_atk_device_access
159+
parameter(omp_atk_device_access=11)
160+
integer(kind=omp_alloctrait_key_kind)omp_atk_target_access
161+
parameter(omp_atk_target_access=12)
162+
integer(kind=omp_alloctrait_key_kind)omp_atk_atomic_scope
163+
parameter(omp_atk_atomic_scope=13)
164+
integer(kind=omp_alloctrait_key_kind)omp_atk_part_size
165+
parameter(omp_atk_part_size=14)
154166

155167
integer(kind=omp_alloctrait_val_kind)omp_atv_default
156168
parameter(omp_atv_default=-1)
@@ -170,8 +182,8 @@
170182
parameter(omp_atv_sequential=5)
171183
integer(kind=omp_alloctrait_val_kind)omp_atv_private
172184
parameter(omp_atv_private=6)
173-
integer(kind=omp_alloctrait_val_kind)omp_atv_all
174-
parameter(omp_atv_all=7)
185+
integer(kind=omp_alloctrait_val_kind)omp_atv_device
186+
parameter(omp_atv_device=7)
175187
integer(kind=omp_alloctrait_val_kind)omp_atv_thread
176188
parameter(omp_atv_thread=8)
177189
integer(kind=omp_alloctrait_val_kind)omp_atv_pteam
@@ -194,6 +206,14 @@
194206
parameter(omp_atv_blocked=17)
195207
integer(kind=omp_alloctrait_val_kind)omp_atv_interleaved
196208
parameter(omp_atv_interleaved=18)
209+
integer(kind=omp_alloctrait_val_kind)omp_atv_all
210+
parameter(omp_atv_all=19)
211+
integer(kind=omp_alloctrait_val_kind)omp_atv_single
212+
parameter(omp_atv_single=20)
213+
integer(kind=omp_alloctrait_val_kind)omp_atv_multiple
214+
parameter(omp_atv_multiple=21)
215+
integer(kind=omp_alloctrait_val_kind)omp_atv_memspace
216+
parameter(omp_atv_memspace=22)
197217

198218
type omp_alloctrait
199219
integer (kind=omp_alloctrait_key_kind) key
@@ -225,8 +245,10 @@
225245
integer(omp_allocator_handle_kind)llvm_omp_target_device_mem_alloc
226246
parameter(llvm_omp_target_device_mem_alloc=102)
227247

248+
integer(kind=omp_memspace_handle_kind)omp_null_mem_space
249+
parameter(omp_null_mem_space=0)
228250
integer(kind=omp_memspace_handle_kind)omp_default_mem_space
229-
parameter(omp_default_mem_space=0)
251+
parameter(omp_default_mem_space=99)
230252
integer(kind=omp_memspace_handle_kind)omp_large_cap_mem_space
231253
parameter(omp_large_cap_mem_space=1)
232254
integer(kind=omp_memspace_handle_kind)omp_const_mem_space
@@ -863,6 +885,98 @@
863885
logical (kind=omp_logical_kind) omp_in_explicit_task
864886
end function omp_in_explicit_task
865887

888+
function omp_get_devices_memspace(ndevs, devs, memspace)
889+
import
890+
integer(omp_memspace_handle_kind) :: omp_get_devices_memspace
891+
integer, intent(in) :: ndevs
892+
integer, intent(in) :: devs(*)
893+
integer(omp_memspace_handle_kind), intent(in) :: memspace
894+
end function omp_get_devices_memspace
895+
896+
function omp_get_device_memspace(dev, memspace)
897+
import
898+
integer(omp_memspace_handle_kind) :: omp_get_device_memspace
899+
integer, intent(in) :: dev
900+
integer(omp_memspace_handle_kind), intent(in) :: memspace
901+
end function omp_get_device_memspace
902+
903+
function omp_get_devices_and_host_memspace(ndevs,devs,memspace)
904+
import
905+
integer(omp_memspace_handle_kind) :: &
906+
& omp_get_devices_and_host_memspace
907+
integer, intent(in) :: ndevs
908+
integer, intent(in) :: devs(*)
909+
integer(omp_memspace_handle_kind), intent(in) :: memspace
910+
end function omp_get_devices_and_host_memspace
911+
912+
function omp_get_device_and_host_memspace(dev, memspace)
913+
import
914+
integer(omp_memspace_handle_kind) :: &
915+
& omp_get_device_and_host_memspace
916+
integer, intent(in) :: dev
917+
integer(omp_memspace_handle_kind), intent(in) :: memspace
918+
end function omp_get_device_and_host_memspace
919+
920+
function omp_get_devices_all_memspace(memspace)
921+
import
922+
integer(omp_memspace_handle_kind)::omp_get_devices_all_memspace
923+
integer(omp_memspace_handle_kind), intent(in) :: memspace
924+
end function omp_get_devices_all_memspace
925+
926+
function omp_get_devices_allocator(ndevs, devs, memspace)
927+
import
928+
integer(omp_allocator_handle_kind)::omp_get_devices_allocator
929+
integer, intent(in) :: ndevs
930+
integer, intent(in) :: devs(*)
931+
integer(omp_memspace_handle_kind), intent(in) :: memspace
932+
end function omp_get_devices_allocator
933+
934+
function omp_get_device_allocator(dev, memspace)
935+
import
936+
integer(omp_allocator_handle_kind) :: omp_get_device_allocator
937+
integer, intent(in) :: dev
938+
integer(omp_memspace_handle_kind), intent(in) :: memspace
939+
end function omp_get_device_allocator
940+
941+
function omp_get_devices_and_host_allocator(ndevs,devs,memspace)
942+
import
943+
integer(omp_allocator_handle_kind) :: &
944+
& omp_get_devices_and_host_allocator
945+
integer, intent(in) :: ndevs
946+
integer, intent(in) :: devs(*)
947+
integer(omp_memspace_handle_kind), intent(in) :: memspace
948+
end function omp_get_devices_and_host_allocator
949+
950+
function omp_get_device_and_host_allocator(dev, memspace)
951+
import
952+
integer(omp_allocator_handle_kind) :: &
953+
& omp_get_device_and_host_allocator
954+
integer, intent(in) :: dev
955+
integer(omp_memspace_handle_kind), intent(in) :: memspace
956+
end function omp_get_device_and_host_allocator
957+
958+
function omp_get_devices_all_allocator(memspace)
959+
import
960+
integer(omp_allocator_handle_kind) :: &
961+
& omp_get_devices_all_allocator
962+
integer(omp_memspace_handle_kind), intent(in) :: memspace
963+
end function omp_get_devices_all_allocator
964+
965+
function omp_get_memspace_num_resources(memspace)
966+
import
967+
integer omp_get_memspace_num_resources
968+
integer(omp_memspace_handle_kind), intent(in) :: memspace
969+
end function omp_get_memspace_num_resources
970+
971+
function omp_get_submemspace(memspace, num_resources, resources)
972+
import
973+
integer(omp_memspace_handle_kind) omp_get_submemspace
974+
integer(omp_memspace_handle_kind), intent(in) :: memspace
975+
integer, intent(in) :: num_resources
976+
integer, intent(in) :: resources(*)
977+
end function omp_get_submemspace
978+
979+
866980
! ***
867981
! *** kmp_* entry points
868982
! ***

0 commit comments

Comments
 (0)