openmp/libomptarget/DeviceRTL/src/Allocator.cpp

   1 //===------ State.cpp - OpenMP State & ICV interface ------------- C++ -*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 //===----------------------------------------------------------------------===//
  10
  11 #include "Allocator.h"
  12 #include "Configuration.h"
  13 #include "Environment.h"
  14 #include "Mapping.h"
  15 #include "Synchronization.h"
  16 #include "Types.h"
  17 #include "Utils.h"
  18
  19 using namespace ompx;
  20
  21 #pragma omp begin declare target device_type(nohost)
  22
  23 [[gnu::used, gnu::retain, gnu::weak,
  24   gnu::visibility(
  25       "protected")]] DeviceMemoryPoolTy __omp_rtl_device_memory_pool;
  26 [[gnu::used, gnu::retain, gnu::weak,
  27   gnu::visibility("protected")]] DeviceMemoryPoolTrackingTy
  28     __omp_rtl_device_memory_pool_tracker;
  29
  30 /// Stateless bump allocator that uses the __omp_rtl_device_memory_pool
  31 /// directly.
  32 struct BumpAllocatorTy final {
  33
  34   void *alloc(uint64_t Size) {
  35     Size = utils::roundUp(Size, uint64_t(allocator::ALIGNMENT));
  36
  37     if (config::isDebugMode(DeviceDebugKind::AllocationTracker)) {
  38       atomic::add(&__omp_rtl_device_memory_pool_tracker.NumAllocations, 1,
  39                   atomic::seq_cst);
  40       atomic::add(&__omp_rtl_device_memory_pool_tracker.AllocationTotal, Size,
  41                   atomic::seq_cst);
  42       atomic::min(&__omp_rtl_device_memory_pool_tracker.AllocationMin, Size,
  43                   atomic::seq_cst);
  44       atomic::max(&__omp_rtl_device_memory_pool_tracker.AllocationMax, Size,
  45                   atomic::seq_cst);
  46     }
  47
  48     uint64_t *Data =
  49         reinterpret_cast<uint64_t *>(&__omp_rtl_device_memory_pool.Ptr);
  50     uint64_t End =
  51         reinterpret_cast<uint64_t>(Data) + __omp_rtl_device_memory_pool.Size;
  52
  53     uint64_t OldData = atomic::add(Data, Size, atomic::seq_cst);
  54     if (OldData + Size > End)
  55       __builtin_trap();
  56
  57     return reinterpret_cast<void *>(OldData);
  58   }
  59
  60   void free(void *) {}
  61 };
  62
  63 BumpAllocatorTy BumpAllocator;
  64
  65 /// allocator namespace implementation
  66 ///
  67 ///{
  68
  69 void allocator::init(bool IsSPMD, KernelEnvironmentTy &KernelEnvironment) {
  70   // TODO: Check KernelEnvironment for an allocator choice as soon as we have
  71   // more than one.
  72 }
  73
  74 void *allocator::alloc(uint64_t Size) { return BumpAllocator.alloc(Size); }
  75
  76 void allocator::free(void *Ptr) { BumpAllocator.free(Ptr); }
  77
  78 ///}
  79
  80 #pragma omp end declare target