offload/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 "Shared/Environment.h"
  12
  13 #include "Allocator.h"
  14 #include "Configuration.h"
  15 #include "DeviceTypes.h"
  16 #include "DeviceUtils.h"
  17 #include "Mapping.h"
  18 #include "Synchronization.h"
  19
  20 using namespace ompx;
  21
  22 #pragma omp begin declare target device_type(nohost)
  23
  24 [[gnu::used, gnu::retain, gnu::weak,
  25   gnu::visibility(
  26       "protected")]] DeviceMemoryPoolTy __omp_rtl_device_memory_pool;
  27 [[gnu::used, gnu::retain, gnu::weak,
  28   gnu::visibility("protected")]] DeviceMemoryPoolTrackingTy
  29     __omp_rtl_device_memory_pool_tracker;
  30
  31 /// Stateless bump allocator that uses the __omp_rtl_device_memory_pool
  32 /// directly.
  33 struct BumpAllocatorTy final {
  34
  35   void *alloc(uint64_t Size) {
  36     Size = utils::roundUp(Size, uint64_t(allocator::ALIGNMENT));
  37
  38     if (config::isDebugMode(DeviceDebugKind::AllocationTracker)) {
  39       atomic::add(&__omp_rtl_device_memory_pool_tracker.NumAllocations, 1,
  40                   atomic::seq_cst);
  41       atomic::add(&__omp_rtl_device_memory_pool_tracker.AllocationTotal, Size,
  42                   atomic::seq_cst);
  43       atomic::min(&__omp_rtl_device_memory_pool_tracker.AllocationMin, Size,
  44                   atomic::seq_cst);
  45       atomic::max(&__omp_rtl_device_memory_pool_tracker.AllocationMax, Size,
  46                   atomic::seq_cst);
  47     }
  48
  49     uint64_t *Data =
  50         reinterpret_cast<uint64_t *>(&__omp_rtl_device_memory_pool.Ptr);
  51     uint64_t End =
  52         reinterpret_cast<uint64_t>(Data) + __omp_rtl_device_memory_pool.Size;
  53
  54     uint64_t OldData = atomic::add(Data, Size, atomic::seq_cst);
  55     if (OldData + Size > End)
  56       __builtin_trap();
  57
  58     return reinterpret_cast<void *>(OldData);
  59   }
  60
  61   void free(void *) {}
  62 };
  63
  64 BumpAllocatorTy BumpAllocator;
  65
  66 /// allocator namespace implementation
  67 ///
  68 ///{
  69
  70 void allocator::init(bool IsSPMD, KernelEnvironmentTy &KernelEnvironment) {
  71   // TODO: Check KernelEnvironment for an allocator choice as soon as we have
  72   // more than one.
  73 }
  74
  75 void *allocator::alloc(uint64_t Size) { return BumpAllocator.alloc(Size); }
  76
  77 void allocator::free(void *Ptr) { BumpAllocator.free(Ptr); }
  78
  79 ///}
  80
  81 #pragma omp end declare target