[flang-rt] Return default allocator directly in device rt (#197009)

flang-rt/lib/runtime/allocator-registry.cpp

diff --git a/flang-rt/lib/runtime/allocator-registry.cpp b/flang-rt/lib/runtime/allocator-registry.cpp
index f8a8daa..d1e85c2 100644
--- a/flang-rt/lib/runtime/allocator-registry.cpp
+++ b/flang-rt/lib/runtime/allocator-registry.cpp
@@ -27,6 +27,11 @@
 
 RT_API_ATTRS AllocFct AllocatorRegistry::GetAllocator(int pos) {
   INTERNAL_CHECK(pos >= 0 && pos < MAX_ALLOCATOR);
+#ifdef RT_DEVICE_COMPILATION
+  if (pos == kDefaultAllocator) {
+    return &MallocWrapper;
+  }
+#endif
   AllocFct f{allocators[pos].alloc};
   INTERNAL_CHECK(f != nullptr);
   return f;
@@ -34,6 +39,11 @@
 
 RT_API_ATTRS FreeFct AllocatorRegistry::GetDeallocator(int pos) {
   INTERNAL_CHECK(pos >= 0 && pos < MAX_ALLOCATOR);
+#ifdef RT_DEVICE_COMPILATION
+  if (pos == kDefaultAllocator) {
+    return &FreeWrapper;
+  }
+#endif
   FreeFct f{allocators[pos].free};
   INTERNAL_CHECK(f != nullptr);
   return f;