Skip to content

Commit

Permalink
Fix pre-main crash in CUDA builds (#905)
Browse files Browse the repository at this point in the history
  • Loading branch information
adayton1 authored Aug 30, 2024
1 parent abd729f commit 7d667b2
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 4 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -7,3 +7,4 @@ CMakeFiles
.devcontainer
*spack.lock*
.spack_env/*/.spack-env
*.orig
19 changes: 15 additions & 4 deletions src/umpire/resource/MemoryResourceTypes.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,15 @@
#define UMPIRE_MemoryResourceTypes_HPP

#include <cstddef>
#include <regex>
#include <string>

#include "umpire/config.hpp"
#include "umpire/util/error.hpp"

#if defined(UMPIRE_ENABLE_CUDA)
#include <cuda_runtime_api.h>
#else
#include <regex>
#endif /* UMPIRE_ENABLE_CUDA */

#if defined(UMPIRE_ENABLE_HIP)
Expand Down Expand Up @@ -106,14 +107,23 @@ inline MemoryResourceType string_to_resource(const std::string& resource)

inline int resource_to_device_id(const std::string& resource)
{
int device_id{0};

#if defined(UMPIRE_ENABLE_CUDA)
if (resource.find("::") != std::string::npos) {
device_id = std::stoi(resource.substr(resource.find("::") + 2));
}
#else
const std::regex id_regex{R"(.*::(\d+))", std::regex_constants::ECMAScript | std::regex_constants::optimize};
std::smatch m;

int device_id{0};
if (std::regex_match(resource, m, id_regex)) {
device_id = std::stoi(m[1]);
} else {
// get the device bound to the current process
}
#endif
else {
// get the device bound to the current process

#if defined(UMPIRE_ENABLE_CUDA)
cudaGetDevice(&device_id);
#endif /* UMPIRE_ENABLE_CUDA */
Expand All @@ -122,6 +132,7 @@ inline int resource_to_device_id(const std::string& resource)
hipGetDevice(&device_id);
#endif /* UMPIRE_ENABLE_HIP */
}

return device_id;
}

Expand Down

0 comments on commit 7d667b2

Please sign in to comment.