diff --git a/flang/runtime/descriptor.cpp b/flang/runtime/descriptor.cpp index 34f7a02ea8c7bd..32f43e89dc7a36 100644 --- a/flang/runtime/descriptor.cpp +++ b/flang/runtime/descriptor.cpp @@ -154,6 +154,15 @@ RT_API_ATTRS std::size_t Descriptor::Elements() const { return elements; } +RT_API_ATTRS static inline int MapAllocIdx(const Descriptor &desc) { +#ifdef RT_DEVICE_COMPILATION + // Force default allocator in device code. + return kDefaultAllocator; +#else + return desc.GetAllocIdx(); +#endif +} + RT_API_ATTRS int Descriptor::Allocate() { std::size_t elementBytes{ElementBytes()}; if (static_cast(elementBytes) < 0) { @@ -162,11 +171,10 @@ RT_API_ATTRS int Descriptor::Allocate() { elementBytes = raw_.elem_len = 0; } std::size_t byteSize{Elements() * elementBytes}; + AllocFct alloc{allocatorRegistry.GetAllocator(MapAllocIdx(*this))}; // Zero size allocation is possible in Fortran and the resulting // descriptor must be allocated/associated. Since std::malloc(0) // result is implementation defined, always allocate at least one byte. - - AllocFct alloc{allocatorRegistry.GetAllocator(GetAllocIdx())}; void *p{alloc(byteSize ? byteSize : 1)}; if (!p) { return CFI_ERROR_MEM_ALLOCATION; @@ -209,7 +217,7 @@ RT_API_ATTRS int Descriptor::Deallocate() { if (!descriptor.base_addr) { return CFI_ERROR_BASE_ADDR_NULL; } else { - FreeFct free{allocatorRegistry.GetDeallocator(GetAllocIdx())}; + FreeFct free{allocatorRegistry.GetDeallocator(MapAllocIdx(*this))}; free(descriptor.base_addr); descriptor.base_addr = nullptr; return CFI_SUCCESS;