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

Fix LORT CI failure due to PyTorch change (#14367)

Browse source 
As title. The fuser in LORT doesn't like "scalar". With a recent PyTorch
change, scalar is intorduced somewhere it was there before. Now, a
simple fix is to check if all inputs are tensors or some specially
allowed cases before sending ops to ORT.
This commit is contained in:
Wei-Sheng Chin 2023-01-19 16:02:40 -08:00 committed by GitHub
parent 36ba3d8d21
commit 432a9912a3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 34 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
orttraining/orttraining/lazy_tensor/accelerator.cc
View file

@ -36,6 +36,39 @@ namespace py = pybind11;
namespace aten = torch::jit::aten;
namespace prim = torch::jit::prim;
bool IsFusable(const torch::jit::Node* node) {
// This function checks the fusion restriction inside
// mergeNodeIntoGroup(...). When selecting onnx-supported aten ops,
// we need to call this function to make sure they are fusable
// in mergeNodeIntoGroup(...).
// Not all inputs are fusable. For example, the fuser
// expects all inputs are tensors with a few exceptions.
// This flag denotes if we find an unsupported input.
bool found_not_fusable = false;
for (auto input : node->inputs()) {
if (input->type()->isSubtypeOf(*c10::TensorType::get())) {
continue;
} else if (
(input->type()->isSubtypeOf(*c10::FloatType::get()) &&
input->node()->kind() != torch::jit::prim::Constant) ||
(node->kind() == torch::jit::aten::_grad_sum_to_size &&
input->type()->isSubtypeOf(*c10::ListType::ofInts()))) {
continue;
} else if (
input->type()->isSubtypeOf(*c10::IntType::get()) &&
input->node()->kind() != torch::jit::prim::Constant) {
continue;
} else{
if (input->node()->kind() == torch::jit::prim::Constant) {
continue;
}
found_not_fusable = true;
}
}
return !found_not_fusable;
}
bool Accelerator::Supported(const torch::jit::Node* node) {
if (!node) {
return false;
@ -74,7 +107,7 @@ bool Accelerator::Supported(const torch::jit::Node* node) {
std::cout << "Supported op: "
<< ToString(*node) << std::endl;
}
return true;
return IsFusable(node);
}
default: {
if (DumpAtenOpHistory()) {