chore: minor

flexflow · Aug 11, 2024 · 6a0bd4a · 6a0bd4a
1 parent ef118f9
commit 6a0bd4a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/include/flexflow/utils/memory_allocator.h b/include/flexflow/utils/memory_allocator.h
@@ -62,7 +62,7 @@ class MemoryAllocator {
   size_t instance_total_size, instance_allocated_size;
 };
 
-Legion::Memory get_proc_mem(Legion::Processor proc);
+Legion::Memory get_proc_mem(Legion::Machine machine, Legion::Processor proc);
 
 }; // namespace FlexFlow
 

diff --git a/src/utils/memory_allocator.cc b/src/utils/memory_allocator.cc
@@ -19,18 +19,19 @@ namespace FlexFlow {
 
 using Legion::Machine;
 using Legion::Memory;
+using Legion::Processor;
 
-Memory get_proc_mem(Legion::Processor proc) {
+Memory get_proc_mem(Machine machine, Processor proc) {
     // First try to allocate a managed memory (cudaMallocManaged)
-    Memory proc_mem = Machine::MemoryQuery(Machine::get_machine())
+    Memory proc_mem = Machine::MemoryQuery(machine)
                         .only_kind(Memory::GPU_MANAGED_MEM)
                         .best_affinity_to(proc)
                         .first();
     if (proc_mem.capacity() > 0) {
       return proc_mem;
     }
     // If managed memory is not available, try to allocate a framebuffer memory
-    proc_mem = Machine::MemoryQuery(Machine::get_machine())
+    proc_mem = Machine::MemoryQuery(machine)
                         .only_kind(Memory::GPU_FB_MEM)
                         .best_affinity_to(proc)
                         .first();