From eb0ea28b6a1b097ff33e9f2e1f94cfe0a2630a7e Mon Sep 17 00:00:00 2001 From: Nawrin Sultana Date: Wed, 25 Jan 2023 16:01:02 -0600 Subject: [PATCH] [OpenMP] Add check for target allocator regardless of the availability of libmemkind Current runtime implementation only checks for target allocator when libmemkind is not available. This patch adds checks for target allocator regardless of the presence of libmemkind library. Differential Revision: https://reviews.llvm.org/D142582 --- openmp/runtime/src/kmp_alloc.cpp | 58 +++++++++++++++++++++------------------- 1 file changed, 30 insertions(+), 28 deletions(-) diff --git a/openmp/runtime/src/kmp_alloc.cpp b/openmp/runtime/src/kmp_alloc.cpp index c602dc0..bf1d064 100644 --- a/openmp/runtime/src/kmp_alloc.cpp +++ b/openmp/runtime/src/kmp_alloc.cpp @@ -1576,6 +1576,36 @@ void *__kmp_alloc(int gtid, size_t algn, size_t size, // Use default allocator if libmemkind is not available int use_default_allocator = (__kmp_memkind_available) ? false : true; + if (KMP_IS_TARGET_MEM_ALLOC(allocator)) { + // Use size input directly as the memory may not be accessible on host. + // Use default device for now. + if (__kmp_target_mem_available) { + kmp_int32 device = + __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device; + if (allocator == llvm_omp_target_host_mem_alloc) + ptr = kmp_target_alloc_host(size, device); + else if (allocator == llvm_omp_target_shared_mem_alloc) + ptr = kmp_target_alloc_shared(size, device); + else // allocator == llvm_omp_target_device_mem_alloc + ptr = kmp_target_alloc_device(size, device); + } + return ptr; + } + + if (allocator >= kmp_max_mem_alloc && KMP_IS_TARGET_MEM_SPACE(al->memspace)) { + if (__kmp_target_mem_available) { + kmp_int32 device = + __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device; + if (al->memspace == llvm_omp_target_host_mem_space) + ptr = kmp_target_alloc_host(size, device); + else if (al->memspace == llvm_omp_target_shared_mem_space) + ptr = kmp_target_alloc_shared(size, device); + else // al->memspace == llvm_omp_target_device_mem_space + ptr = kmp_target_alloc_device(size, device); + } + return ptr; + } + if (__kmp_memkind_available) { if (allocator < kmp_max_mem_alloc) { // pre-defined allocator @@ -1645,22 +1675,6 @@ void *__kmp_alloc(int gtid, size_t algn, size_t size, } } } else if (allocator < kmp_max_mem_alloc) { - if (KMP_IS_TARGET_MEM_ALLOC(allocator)) { - // Use size input directly as the memory may not be accessible on host. - // Use default device for now. - if (__kmp_target_mem_available) { - kmp_int32 device = - __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device; - if (allocator == llvm_omp_target_host_mem_alloc) - ptr = kmp_target_alloc_host(size, device); - else if (allocator == llvm_omp_target_shared_mem_alloc) - ptr = kmp_target_alloc_shared(size, device); - else // allocator == llvm_omp_target_device_mem_alloc - ptr = kmp_target_alloc_device(size, device); - } - return ptr; - } - // pre-defined allocator if (allocator == omp_high_bw_mem_alloc) { KMP_WARNING(OmpNoAllocator, "omp_high_bw_mem_alloc"); @@ -1683,18 +1697,6 @@ void *__kmp_alloc(int gtid, size_t algn, size_t size, ptr = __kmp_thread_malloc(__kmp_thread_from_gtid(gtid), desc.size_a); use_default_allocator = false; } - } else if (KMP_IS_TARGET_MEM_SPACE(al->memspace)) { - if (__kmp_target_mem_available) { - kmp_int32 device = - __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device; - if (al->memspace == llvm_omp_target_host_mem_space) - ptr = kmp_target_alloc_host(size, device); - else if (al->memspace == llvm_omp_target_shared_mem_space) - ptr = kmp_target_alloc_shared(size, device); - else // al->memspace == llvm_omp_target_device_mem_space - ptr = kmp_target_alloc_device(size, device); - } - return ptr; } else if (al->pool_size > 0) { // custom allocator with pool size requested kmp_uint64 used = -- 2.7.4