diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index a73da48fac..b3735ed19f 100644 --- a/lib/CMakeLists.txt +++ b/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) diff --git a/lib/local-execution/CMakeLists.txt b/lib/local-execution/CMakeLists.txt new file mode 100644 index 0000000000..ee1d8fecdc --- /dev/null +++ b/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 +) \ No newline at end of file diff --git a/lib/local-execution/include/local_allocator.h b/lib/local-execution/include/local_allocator.h new file mode 100644 index 0000000000..f4b253b281 --- /dev/null +++ b/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 + +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 diff --git a/lib/local-execution/include/tracked_allocator.h b/lib/local-execution/include/tracked_allocator.h new file mode 100644 index 0000000000..4f51670426 --- /dev/null +++ b/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 diff --git a/lib/local-execution/src/local_allocator.cc b/lib/local-execution/src/local_allocator.cc new file mode 100644 index 0000000000..0bb7d04574 --- /dev/null +++ b/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(); +} + +} // namespace FlexFlow diff --git a/lib/local-execution/src/tracked_allocator.cc b/lib/local-execution/src/tracked_allocator.cc new file mode 100644 index 0000000000..6d06714252 --- /dev/null +++ b/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(base_allocator); +} + +} // namespace FlexFlow