feat: add func get_proc_mem to process memory allocation

include/flexflow/utils/memory_allocator.h

@@ -62,6 +62,8 @@ class MemoryAllocator {
   size_t instance_total_size, instance_allocated_size;
 };
 
+Legion::Memory get_proc_mem(Legion::Processor proc);
+
 }; // namespace FlexFlow
 
 #endif // _FLEXFLOW_RUNTIME_H_

src/utils/memory_allocator.cc

@@ -0,0 +1,42 @@
+/* Copyright 2023 CMU, Facebook, LANL, MIT, NVIDIA, and Stanford (alphabetical)
+ *
+ * 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 "flexflow/utils/memory_allocator.h"
+
+namespace FlexFlow {
+
+using Legion::Machine;
+using Legion::Memory;
+
+Memory get_proc_mem(Legion::Processor proc) {
+    // First try to allocate a managed memory (cudaMallocManaged)
+    Memory proc_mem = Machine::MemoryQuery(Machine::get_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())
+                        .only_kind(Memory::GPU_FB_MEM)
+                        .best_affinity_to(proc)
+                        .first();
+    assert(proc_mem.capacity() > 0);
+    return proc_mem;
+}
+
+} // namespace FlexFlow
+