Skip to content

Commit

Permalink
[JSEP] allow DataTransfer to deal with zero sized input (microsoft#17661
Browse files Browse the repository at this point in the history
)

### 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.
  • Loading branch information
fs-eire authored Sep 25, 2023
1 parent fcfc239 commit f50fa46
Showing 1 changed file with 18 additions and 16 deletions.
34 changes: 18 additions & 16 deletions onnxruntime/core/providers/js/data_transfer.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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();

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 (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;

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();
Expand Down

0 comments on commit f50fa46

Please sign in to comment.