saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support ORTModule on ROCm EP (#6945)

Browse source
This commit is contained in:
Weixing Zhang 2021-03-09 10:10:57 -08:00 committed by GitHub
parent 3b2847b2d8
commit 534adbb065
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/python/dlpack_convertor.cc
View file

@ -84,7 +84,11 @@ DLContext get_dlpack_context(const OrtValue& ort_value, const int64_t& device_id
ctx.device_type = DLDeviceType::kDLCPU;
break;
case OrtDevice::GPU:
#ifdef USE_ROCM
ctx.device_type = DLDeviceType::kDLROCM;
#else
ctx.device_type = DLDeviceType::kDLGPU;
#endif
break;
default:
ORT_THROW("Cannot pack tensors on this device.");

2
orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc
View file

@ -124,6 +124,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, Recv
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, RecordEvent);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, WaitEvent);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, YieldOp);
#ifdef ORT_USE_NCCL
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, NcclAllReduce);
@ -247,6 +248,7 @@ Status RegisterRocmTrainingKernels(KernelRegistry& kernel_registry) {
// BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, RecordEvent)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, WaitEvent)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, YieldOp)>,
#ifdef ORT_USE_NCCL
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, NcclAllReduce)>,