<commit_before>/* Copyright 2016 Car<NAME> University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "scanner/util/memory.h"

#include <cassert>
#include <mutex>
#include <sys/syscall.h>
#include <unistd.h>

#ifdef HAVE_CUDA
#include <cuda.h>
#include <cuda_runtime_api.h>
#include "scanner/util/cuda.h"
#endif

namespace scanner {

// Allocations in Scanner differ from manual memory management in three respects:
//
// 1. Scanner uses different allocators for different devices, e.g. the `new`
//    keyword for CPU allocations and cudaMalloc for GPU memory. The System and
//    Pool allocators are parameterized by a DeviceHandle to determine which
//    device they are allocating for.
//
// 2. The pool allocator up-front allocates a large memory pool and then does its
//    own simple memory allocation in place of the system allocator for the
//    respective device. Memory pools avoid potentially expensive and synchronous
//    calls to the system allocator, e.g. cudaMalloc which synchronizes the whole
//    GPU.
//
// 3. Block allocations allow evaluators to allocate a single block of memory for
//    their returned rows instead of allocating individually for each row. This
//    again reduces the number of cudaMallocs if not using a memory pool.
//    Regardless of pool usage, blocks can also be copied in a single memcpy
//    instead of many, which reduces memcpy calls. To avoid complexity in the
//    core Scanner engine, it is oblivious to whether a u8* in an output row
//    is from a block or an individual allocation. Instead, the allocation
//    runtime does reference counting when the engine calls free on a memory
//    block, e.g. if a memory block is allocated for 96 rows (96 different
//    pointers in the same block), then each free to a pointer into the block
//    decrements a reference counter until freeing the block at 0 refs.
//
// The user can dictate usage of the memory pool with the MemoryPoolConfig, but
// cannot directly call into it. Users can only ask for normal memory segments
// or block memory segments, the former of which is allocated by the system
// and the latter by the pool if it exists.

class Allocator {
public:
  virtual ~Allocator(){};

  virtual u8* allocate(size_t size) = 0;
  virtual void free(u8* buffer) = 0;
};

class SystemAllocator : public Allocator {
public:
  SystemAllocator(DeviceHandle device) :
    device_(device) {
  }

  u8* allocate(size_t size) {
    if (device_.type == DeviceType::CPU) {
      return new u8[size];
    } else if (device_.type == DeviceType::GPU) {
      u8* buffer;
#ifdef HAVE_CUDA
      CU_CHECK(cudaSetDevice(device_.id));
      CU_CHECK(cudaMalloc((void**) &buffer, size));
#else
      LOG(FATAL) << "Cuda not enabled.";
#endif
      return buffer;
    }
  }

  void free(u8* buffer) {
    if (device_.type == DeviceType::CPU) {
      delete buffer;
    } else if (device_.type == DeviceType::GPU) {
      CU_CHECK(cudaSetDevice(device_.id));
      CU_CHECK(cudaFree(buffer));
    }
  }

private:
  DeviceHandle device_;
};

bool pointer_in_buffer(u8* ptr, u8* buf_start, u8* buf_end) {
  return (size_t) ptr >= (size_t) buf_start &&
    (size_t) ptr < (size_t) buf_end;
}

class PoolAllocator : public Allocator {
public:
  PoolAllocator(DeviceHandle device, SystemAllocator* allocator, i64 pool_size) :
    device_(device),
    system_allocator(allocator),
    pool_size_(pool_size) {
    pool_ = system_allocator->allocate(pool_size_);
  }

  ~PoolAllocator() {
    system_allocator->free(pool_);
  }

  u8* allocate(size_t size) {
    Allocation alloc;
    alloc.length = size;

    std::lock_guard<std::mutex> guard(lock_);
    bool found = false;
    i32 num_alloc = allocations_.size();
    for (i32 i = 0; i < num_alloc; ++i) {
      Allocation lower;
      if (i == 0) {
        lower.offset = 0;
        lower.length = 0;
      } else {
        lower = allocations_[i-1];
      }
      Allocation higher = allocations_[i];
      assert(higher.offset >= lower.offset + lower.length);
      if ((higher.offset - (lower.offset + lower.length)) >= size) {
        alloc.offset = lower.offset + lower.length;
        allocations_.insert(allocations_.begin() + i, alloc);
        found = true;
        break;
      }
    }

    if (!found) {
      if (num_alloc > 0) {
        Allocation& last = allocations_[num_alloc - 1];
        alloc.offset = last.offset + last.length;
      } else {
        alloc.offset = 0;
      }
      allocations_.push_back(alloc);
    }

    u8* buffer = pool_ + alloc.offset;
    LOG_IF(FATAL, alloc.offset + alloc.length >= pool_size_)
      << "Exceeded pool size";

    return buffer;
  }

  void free(u8* buffer) {
    LOG_IF(FATAL, !pointer_in_buffer(buffer, pool_, pool_ + pool_size_))
      << "Pool allocator tried to free buffer not in pool";

    std::lock_guard<std::mutex> guard(lock_);
    i32 index;
    bool found = find_buffer(buffer, index);
    LOG_IF(FATAL, !found)
      << "Attempted to free unallocated buffer in pool";


    Allocation& alloc = allocations_[index];
    allocations_.erase(allocations_.begin() + index);
  }

private:
  bool find_buffer(u8* buffer, i32& index) {
    i32 num_alloc = allocations_.size();
    for (i32 i = 0; i < num_alloc; ++i) {
      Allocation alloc = allocations_[i];
      if ((size_t) buffer == (size_t) pool_ + alloc.offset) {
        index = i;
        return true;
      }
    }
    return false;
  }

  typedef struct {
    i64 offset;
    i64 length;
  } Allocation;

  DeviceHandle device_;
  u8* pool_ = nullptr;
  i64 pool_size_;
  std::mutex lock_;
  std::vector<Allocation> allocations_;

  SystemAllocator* system_allocator;
};

class BlockAllocator {
public:
  BlockAllocator(Allocator* allocator) : allocator_(allocator) {
  }

  u8* allocate(size_t size, i32 refs) {
    u8* buffer = allocator_->allocate(size);

    Allocation alloc;
    alloc.buffer = buffer;
    alloc.size = size;
    alloc.refs = refs;

    std::lock_guard<std::mutex> guard(lock_);
    allocations_.push_back(alloc);

    return buffer;
  }<commit_msg>Added memory alignment to pool allocator<commit_after>/* Copyright 2016 Carnegie Mellon University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "scanner/util/memory.h"

#include <cassert>
#include <mutex>
#include <sys/syscall.h>
#include <unistd.h>

#ifdef HAVE_CUDA
#include <cuda.h>
#include <cuda_runtime_api.h>
#include "scanner/util/cuda.h"
#endif

namespace scanner {

// Allocations in Scanner differ from manual memory management in three respects:
//
// 1. Scanner uses different allocators for different devices, e.g. the `new`
//    keyword for CPU allocations and cudaMalloc for GPU memory. The System and
//    Pool allocators are parameterized by a DeviceHandle to determine which
//    device they are allocating for.
//
// 2. The pool allocator up-front allocates a large memory pool and then does its
//    own simple memory allocation in place of the system allocator for the
//    respective device. Memory pools avoid potentially expensive and synchronous
//    calls to the system allocator, e.g. cudaMalloc which synchronizes the whole
//    GPU.
//
// 3. Block allocations allow evaluators to allocate a single block of memory for
//    their returned rows instead of allocating individually for each row. This
//    again reduces the number of cudaMallocs if not using a memory pool.
//    Regardless of pool usage, blocks can also be copied in a single memcpy
//    instead of many, which reduces memcpy calls. To avoid complexity in the
//    core Scanner engine, it is oblivious to whether a u8* in an output row
//    is from a block or an individual allocation. Instead, the allocation
//    runtime does reference counting when the engine calls free on a memory
//    block, e.g. if a memory block is allocated for 96 rows (96 different
//    pointers in the same block), then each free to a pointer into the block
//    decrements a reference counter until freeing the block at 0 refs.
//
// The user can dictate usage of the memory pool with the MemoryPoolConfig, but
// cannot directly call into it. Users can only ask for normal memory segments
// or block memory segments, the former of which is allocated by the system
// and the latter by the pool if it exists.

class Allocator {
public:
  virtual ~Allocator(){};

  virtual u8* allocate(size_t size) = 0;
  virtual void free(u8* buffer) = 0;
};

class SystemAllocator : public Allocator {
public:
  SystemAllocator(DeviceHandle device) :
    device_(device) {
  }

  u8* allocate(size_t size) {
    if (device_.type == DeviceType::CPU) {
      return new u8[size];
    } else if (device_.type == DeviceType::GPU) {
      u8* buffer;
#ifdef HAVE_CUDA
      CU_CHECK(cudaSetDevice(device_.id));
      CU_CHECK(cudaMalloc((void**) &buffer, size));
#else
      LOG(FATAL) << "Cuda not enabled.";
#endif
      return buffer;
    }
  }

  void free(u8* buffer) {
    if (device_.type == DeviceType::CPU) {
      delete buffer;
    } else if (device_.type == DeviceType::GPU) {
      CU_CHECK(cudaSetDevice(device_.id));
      CU_CHECK(cudaFree(buffer));
    }
  }

  size_t alignment() {
    if (device_.type == DeviceType::CPU) {
      return 16;
    } else if (device_.type == DeviceType::GPU) {
      return 256;
    }
  }

private:
  DeviceHandle device_;
};

bool pointer_in_buffer(u8* ptr, u8* buf_start, u8* buf_end) {
  return (size_t) ptr >= (size_t) buf_start &&
    (size_t) ptr < (size_t) buf_end;
}

class PoolAllocator : public Allocator {
public:
  PoolAllocator(DeviceHandle device, SystemAllocator* allocator, size_t pool_size) :
    device_(device),
    system_allocator(allocator),
    pool_size_(pool_size) {
    pool_ = system_allocator->allocate(pool_size_);
  }

  ~PoolAllocator() {
    system_allocator->free(pool_);
  }

  u8* allocate(size_t size) {
    Allocation alloc;
    alloc.length = size;

    std::lock_guard<std::mutex> guard(lock_);
    bool found = false;
    i32 num_alloc = allocations_.size();
    for (i32 i = 0; i < num_alloc; ++i) {
      Allocation lower;
      if (i == 0) {
        lower.offset = 0;
        lower.length = 0;
      } else {
        lower = allocations_[i-1];
      }
      Allocation higher = allocations_[i];
      assert(higher.offset >= lower.offset + lower.length);
      size_t base = align(lower.offset + lower.length);
      if ((higher.offset - base) >= size) {
        alloc.offset = base;
        allocations_.insert(allocations_.begin() + i, alloc);
        found = true;
        break;
      }
    }

    if (!found) {
      if (num_alloc > 0) {
        Allocation& last = allocations_[num_alloc - 1];
        alloc.offset = align(last.offset + last.length);
      } else {
        alloc.offset = 0;
      }
      allocations_.push_back(alloc);
    }

    LOG_IF(FATAL, alloc.offset + alloc.length >= pool_size_)
      << "Exceeded pool size";

    u8* buffer = pool_ + alloc.offset;
    return buffer;
  }

  size_t align(size_t ptr) {
    size_t alignment = system_allocator->alignment();
    size_t remainder = ptr % alignment;
    if (remainder != 0) {
      return ptr + (alignment - remainder);
    } else {
      return ptr;
    }
  }

  void free(u8* buffer) {
    LOG_IF(FATAL, !pointer_in_buffer(buffer, pool_, pool_ + pool_size_))
      << "Pool allocator tried to free buffer not in pool";

    std::lock_guard<std::mutex> guard(lock_);
    i32 index;
    bool found = find_buffer(buffer, index);
    LOG_IF(FATAL, !found)
      << "Attempted to free unallocated buffer in pool";

    Allocation& alloc = allocations_[index];
    allocations_.erase(allocations_.begin() + index);
  }

private:
  bool find_buffer(u8* buffer, i32& index) {
    i32 num_alloc = allocations_.size();
    for (i32 i = 0; i < num_alloc; ++i) {
      Allocation alloc = allocations_[i];
      if ((size_t) buffer == (size_t) pool_ + alloc.offset) {
        index = i;
        return true;
      }
    }
    return false;
  }

  typedef struct {
    size_t offset;
    size_t length;
  } Allocation;

  DeviceHandle device_;
  u8* pool_ = nullptr;
  size_t pool_size_;
  std::mutex lock_;
  std::vector<Allocation> allocations_;

  SystemAllocator* system_allocator;
};

class BlockAllocator {
public:
  BlockAllocator(Allocator* allocator) : allocator_(allocator) {
  }

  u8* allocate(size_t size, i32 refs) {
    u8* buffer = allocator_->allocate(size);

    Allocation alloc;
    alloc.buffer = buffer;
    alloc.size = size;
    alloc.refs = refs;

    std::lock_guard<std::mutex> guard(lock_);
    allocations_.push_back(alloc);

    return buffer;
  }
{'fraction_non_alphanumeric': 0.10373200117543344, 'fraction_numerical': 0.009550396708786366, 'mean_word_length': 3.0672243800418286, 'pattern_counts': {'":': 0, '<': 46, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4788127', 'n_tokens_mistral': 4295, 'n_tokens_neox': 4143, 'n_words': 1417}