Skip to content

Commit

Permalink
[Runtime] Add Stringification to DLDeviceType and DLDevice
Browse files Browse the repository at this point in the history
This PR adds `DLDeviceType2String` and overloading of ostream's
`operator<<` on `DLDevice` so that they could be printed as strings more
with better developer ergonomics.
  • Loading branch information
junrushao committed Aug 27, 2023
1 parent aa805f2 commit b80177d
Showing 1 changed file with 41 additions and 0 deletions.
41 changes: 41 additions & 0 deletions include/tvm/runtime/packed_func.h
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
#include <functional>
#include <limits>
#include <memory>
#include <sstream>
#include <string>
#include <tuple>
#include <type_traits>
Expand Down Expand Up @@ -1259,6 +1260,46 @@ inline const char* ArgTypeCode2Str(int type_code) {
}
}

/*! \brief Convert a DLDeviceType to string */
inline std::string DLDeviceType2Str(DLDeviceType ty) {
switch (ty) {
case DLDeviceType::kDLCPU:
return "cpu";
case DLDeviceType::kDLCUDA:
return "cuda";
case DLDeviceType::kDLCUDAHost:
return "cuda-host";
case DLDeviceType::kDLOpenCL:
return "opencl";
case DLDeviceType::kDLVulkan:
return "vulkan";
case DLDeviceType::kDLMetal:
return "metal";
case DLDeviceType::kDLVPI:
return "vpi";
case DLDeviceType::kDLROCM:
return "rocm";
case DLDeviceType::kDLROCMHost:
return "rocm-host";
case DLDeviceType::kDLCUDAManaged:
return "cuda-managed";
case DLDeviceType::kDLOneAPI:
return "oneapi";
case DLDeviceType::kDLWebGPU:
return "webgpu";
case DLDeviceType::kDLHexagon:
return "hexagon";
default:
return "Device(" + std::to_string(ty) + ")";
}
throw;
}

inline std::ostringstream& operator<<(std::ostringstream& os, const DLDevice& device) {
os << DLDeviceType2Str(device.device_type) << ":" << device.device_id;
return os;
}

namespace detail {

template <bool stop, std::size_t I, typename F>
Expand Down

0 comments on commit b80177d

Please sign in to comment.