处理kDLROCM有关的函数注册

include/dgl/runtime/ndarray.h

@@ -600,7 +600,7 @@ inline const char* TypeCode2Str(int type_code) {
 inline const char* DeviceTypeCode2Str(DLDeviceType device_type) {
   switch (device_type) {
     case kDLCPU: return "cpu";
-    case kDLROCM: return "cuda";
+    case kDLGPU: return "cuda";
     case kDLCPUPinned: return "cpu_pinned";
     case kDLOpenCL: return "opencl";
     case kDLVulkan: return "vulkan";

src/runtime/c_runtime_api.cc

@@ -27,7 +27,7 @@ namespace runtime {
 inline std::string DeviceName(int type) {
   switch (type) {
     case kDLCPU: return "cpu";
-    case kDLROCM: return "gpu";
+    case kDLGPU: return "gpu";
     case kDLOpenCL: return "opencl";
     case kDLSDAccel: return "sdaccel";
     case kDLAOCL: return "aocl";

src/runtime/cuda/cuda_device_api.cc

@@ -335,5 +335,11 @@ DGL_REGISTER_GLOBAL("device_api.gpu")
     *rv = static_cast<void*>(ptr);
   });
 
+DGL_REGISTER_GLOBAL("device_api.rocm")
+.set_body([](DGLArgs args, DGLRetValue* rv) {
+    DeviceAPI* ptr = CUDADeviceAPI::Global().get();
+    *rv = static_cast<void*>(ptr);
+  });
+
 }  // namespace runtime
 }  // namespace dgl