From fd5782920ab0fa9bd9d008436b850673c2fb321c Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 3 Sep 2026 15:19:14 -0700 Subject: [PATCH] Use the largest device memory available under Realm. --- .../dynamic_tensor_accessor_from_instance.cc | 1 + .../src/realm-execution/realm_context.cc | 15 ++++++++++++--- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/lib/realm-execution/src/realm-execution/dynamic_tensor_accessor_from_instance.cc b/lib/realm-execution/src/realm-execution/dynamic_tensor_accessor_from_instance.cc index a2a40e3752..9c76aeac13 100644 --- a/lib/realm-execution/src/realm-execution/dynamic_tensor_accessor_from_instance.cc +++ b/lib/realm-execution/src/realm-execution/dynamic_tensor_accessor_from_instance.cc @@ -15,6 +15,7 @@ static DeviceType infer_device_type_from_memory_and_processor( device_type = DeviceType::CPU; break; case Realm::Memory::GPU_FB_MEM: + case Realm::Memory::GPU_DYNAMIC_MEM: // Only accessible on GPU device_type = DeviceType::GPU; break; diff --git a/lib/realm-execution/src/realm-execution/realm_context.cc b/lib/realm-execution/src/realm-execution/realm_context.cc index 1e90d4573f..e722d4a65d 100644 --- a/lib/realm-execution/src/realm-execution/realm_context.cc +++ b/lib/realm-execution/src/realm-execution/realm_context.cc @@ -158,12 +158,21 @@ Realm::Memory RealmContext::get_nearest_memory(Realm::Processor proc) { return Realm::Memory::NO_MEMORY; } - // FIMXE: this isn't going to do what you expect until - // https://github.com/StanfordLegion/realm/pull/392 merges Realm::Machine::MemoryQuery mq(Realm::Machine::get_machine()); mq.best_affinity_to(proc); ASSERT(mq.count() > 0); - return mq.first(); + + // Among the best memories, pick the one with the largest capacity + Realm::Memory result = Realm::Memory::NO_MEMORY; + for (Realm::Machine::MemoryQuery::iterator it = mq.begin(); it != mq.end(); + ++it) { + if (!result.exists() || it->capacity() > result.capacity()) { + result = *it; + } + } + + ASSERT(result.exists()); + return result; } Realm::Processor RealmContext::get_current_processor() const {