[JSEP] allow DataTransfer to deal with zero sized input (#17661)

### Description
allow DataTransfer to deal with zero sized input.

This is a standalone fix for zero-sized tensor handling for JSEP
DataTransfer. There are other components in JSEP not supporting
zero-sized tensors need to be fixed.
This commit is contained in:
Yulong Wang 2023-09-25 12:21:20 -07:00 committed by GitHub
parent fcfc2391b8
commit f50fa46fe0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -20,23 +20,25 @@ bool DataTransfer::CanCopy(const OrtDevice& src_device, const OrtDevice& dst_dev
common::Status DataTransfer::CopyTensor(const Tensor& src, Tensor& dst) const {
size_t bytes = src.SizeInBytes();
const void* src_data = src.DataRaw();
void* dst_data = dst.MutableDataRaw();
if (bytes > 0) {
const void* src_data = src.DataRaw();
void* dst_data = dst.MutableDataRaw();
auto& src_device = src.Location().device;
auto& dst_device = dst.Location().device;
auto& src_device = src.Location().device;
auto& dst_device = dst.Location().device;
if (dst_device.Type() == OrtDevice::GPU) {
if (src_device.Type() == OrtDevice::GPU) {
// copy from GPU to GPU
EM_ASM({ Module.jsepCopy($0, $1, $2, true); }, src_data, dst_data, bytes);
} else {
// copy from CPU to GPU
EM_ASM({ Module.jsepCopy($0, $1, $2); }, src_data, dst_data, bytes);
if (dst_device.Type() == OrtDevice::GPU) {
if (src_device.Type() == OrtDevice::GPU) {
// copy from GPU to GPU
EM_ASM({ Module.jsepCopy($0, $1, $2, true); }, src_data, dst_data, bytes);
} else {
// copy from CPU to GPU
EM_ASM({ Module.jsepCopy($0, $1, $2); }, src_data, dst_data, bytes);
}
} else /* if (src_device.Type() == OrtDevice::GPU) */ {
// copy from GPU to CPU
jsepDownload(src_data, dst_data, bytes);
}
} else /* if (src_device.Type() == OrtDevice::GPU) */ {
// copy from GPU to CPU
jsepDownload(src_data, dst_data, bytes);
}
return Status::OK();