Local allocator

lib/CMakeLists.txt

@@ -3,6 +3,7 @@ add_subdirectory(compiler)
 add_subdirectory(runtime)
 add_subdirectory(op-attrs)
 add_subdirectory(kernels)
+add_subdirectory(local-execution)
 add_subdirectory(utils)
 add_subdirectory(ffi)
 add_subdirectory(substitutions)

lib/local-execution/CMakeLists.txt

@@ -0,0 +1,15 @@
+ff_add_library(
+  NAME
+    local-execution
+  SRC_PATTERNS
+    src/*.cc
+  PUBLIC_INCLUDE
+    include/
+  PRIVATE_INCLUDE
+    src/
+  DEPS
+    op-attrs
+    utils
+    kernels
+    pcg
+)

lib/local-execution/include/local_allocator.h

@@ -0,0 +1,24 @@
+#ifndef _FLEXFLOW_RUNTIME_SRC_LOCAL_ALLOCATOR_H
+#define _FLEXFLOW_RUNTIME_SRC_LOCAL_ALLOCATOR_H
+
+#include "kernels/allocation.h"
+#include <unordered_set>
+
+namespace FlexFlow {
+
+struct LocalAllocator : public IAllocator {
+  LocalAllocator() = default;
+  LocalAllocator(LocalAllocator const &) = delete;
+  LocalAllocator(LocalAllocator &&) = delete;
+  ~LocalAllocator() = default;
+
+  void *allocate(size_t) override;
+  void deallocate(void *) override;
+};
+CHECK_RC_COPY_VIRTUAL_COMPLIANT(LocalAllocator);
+
+Allocator get_local_memory_allocator();
+
+} // namespace FlexFlow
+
+#endif

lib/local-execution/include/tracked_allocator.h

@@ -0,0 +1,29 @@
+#ifndef _FLEXFLOW_LOCAL_EXECUTION_TRACKED_ALLOCATOR_H
+#define _FLEXFLOW_LOCAL_EXECUTION_TRACKED_ALLOCATOR_H
+
+#include "kernels/allocation.h"
+#include "local_allocator.h"
+
+namespace FlexFlow {
+
+struct TrackedAllocator : public IAllocator {
+  TrackedAllocator(Allocator);
+  TrackedAllocator(TrackedAllocator const &) = delete;
+  TrackedAllocator(TrackedAllocator &&) = delete;
+  ~TrackedAllocator() = default;
+
+  void *allocate(size_t) override;
+  void deallocate(void *) override;
+  size_t get_current_mem_usage();
+
+private:
+  size_t current_mem_usage = 0;
+  Allocator allocator;
+};
+CHECK_RC_COPY_VIRTUAL_COMPLIANT(TrackedAllocator);
+
+Allocator get_tracked_memory_allocator(Allocator const &base_allocator);
+
+} // namespace FlexFlow
+
+#endif

lib/local-execution/src/local_allocator.cc

@@ -0,0 +1,20 @@
+#include "local_allocator.h"
+#include "kernels/device.h"
+
+namespace FlexFlow {
+
+void *LocalAllocator::allocate(size_t requested_memory_size) {
+  void *ptr;
+  checkCUDA(cudaMalloc(&ptr, requested_memory_size));
+  return ptr;
+}
+
+void LocalAllocator::deallocate(void *ptr) {
+  checkCUDA(cudaFree(ptr));
+}
+
+Allocator get_local_memory_allocator() {
+  return Allocator::create<LocalAllocator>();
+}
+
+} // namespace FlexFlow

lib/local-execution/src/tracked_allocator.cc

@@ -0,0 +1,29 @@
+#include "tracked_allocator.h"
+#include "kernels/device.h"
+
+namespace FlexFlow {
+
+TrackedAllocator::TrackedAllocator(Allocator a) : allocator(a) {}
+
+void *TrackedAllocator::allocate(size_t requested_memory_size) {
+  void *ptr = this->allocator.allocate(requested_memory_size);
+  this->current_mem_usage += requested_memory_size;
+  return ptr;
+}
+
+void TrackedAllocator::deallocate(void *ptr) {
+  size_t psize;
+  checkCUDA(cudaGetSymbolSize(&psize, ptr));
+  this->allocator.deallocate(ptr);
+  this->current_mem_usage -= psize;
+}
+
+size_t TrackedAllocator::get_current_mem_usage() {
+  return this->current_mem_usage;
+}
+
+Allocator get_tracked_memory_allocator(Allocator const &base_allocator) {
+  return Allocator::create<TrackedAllocator>(base_allocator);
+}
+
+} // namespace FlexFlow