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

Add cast before copy for dissimilar scalar type (#12391)

Browse source 
* Add proper cast/copy callflow for ORT and non-ORT devices
This commit is contained in:
smrkatte 2022-08-02 18:32:58 -07:00 committed by GitHub
parent c9e0d0f8b6
commit 54d5e86981
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 19 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
orttraining/orttraining/eager/ort_aten.cpp
View file

@ -663,20 +663,28 @@ at::Tensor& copy_(
auto& invoker = GetORTInvoker(self.device().type() == at::kORT ? self.device() : src.device());
const auto ort_src = create_ort_value(invoker, src);
auto ort_self = create_ort_value(invoker, self);
if (self.scalar_type() != src.scalar_type()) {
// invoke cast first
std::vector<OrtValue> ort_cast_output(1);
onnxruntime::NodeAttributes attrs(1);
attrs["to"] = create_ort_attribute(
"to", (int64_t)GetONNXTensorProtoDataType(self.scalar_type()), at::kLong);
if(src.device().type() != at::kORT)
{
// invoke cast first and then copy for non-ORT device types
auto val = at::native::to(src, self.scalar_type());
const auto ort_val = create_ort_value(invoker, val);
copy(invoker, ort_val, ort_self);
}else{
// For ORT device type, the cast operation will perform the copy as well
std::vector<OrtValue> ort_cast_output(1);
ort_cast_output[0] = ort_self;
onnxruntime::NodeAttributes attrs(1);
attrs["to"] = create_ort_attribute(
"to", (int64_t)GetONNXTensorProtoDataType(self.scalar_type()), at::kLong);
auto status = invoker.Invoke("Cast",
{std::move(ort_src)},
ort_cast_output, &attrs);
auto status = invoker.Invoke("Cast",
{std::move(ort_src)},
ort_cast_output, &attrs);
CHECK_STATUS(status);
copy(invoker, ort_cast_output[0], ort_self);
CHECK_STATUS(status);
}
} else {
copy(invoker, ort_src, ort_self);
}