diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index f8a48fabb3..4688517e34 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -53,6 +53,7 @@ #include "core/optimizer/nchwc_transformer.h" #include "core/optimizer/noop_elimination.h" #include "core/optimizer/not_where_fusion.h" +#include "core/optimizer/pre_shape_node_elimination.h" #ifdef MLAS_TARGET_AMD64_IX86 #include "core/optimizer/qdq_transformer/avx2_weight_s8_to_u8.h" #endif @@ -112,6 +113,7 @@ InlinedVector> GenerateRewriteRules( rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); + rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); diff --git a/onnxruntime/core/optimizer/pre_shape_node_elimination.cc b/onnxruntime/core/optimizer/pre_shape_node_elimination.cc new file mode 100644 index 0000000000..23980c9c10 --- /dev/null +++ b/onnxruntime/core/optimizer/pre_shape_node_elimination.cc @@ -0,0 +1,60 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/common/logging/logging.h" +#include "core/optimizer/rewrite_rule.h" +#include "core/optimizer/pre_shape_node_elimination.h" +#include "core/optimizer/utils.h" +#include "core/graph/graph.h" +#include "core/graph/graph_utils.h" + +namespace onnxruntime { + +Status PreShapeNodeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_effect, const logging::Logger&) const { + std::vector consumerIndices; + + // Save Consumer Nodes Indices (Shape) + for (auto it = node.OutputEdgesBegin(), end = node.OutputEdgesEnd(); it != end; ++it) { + const auto& consumer = (*it).GetNode(); + const auto consumer_idx = consumer.Index(); + consumerIndices.push_back(consumer_idx); + } + + // Remove output edges of the cast node + graph_utils::RemoveNodeOutputEdges(graph, node); + + // Change input defs of shape nodes + for (size_t consumer_idx : consumerIndices) { + Node& shape_node = *graph.GetNode(consumer_idx); + shape_node.MutableInputDefs()[0] = node.MutableInputDefs()[0]; + } + + graph.RemoveNode(node.Index()); + rule_effect = RewriteRuleEffect::kRemovedCurrentNode; + + return Status::OK(); +} + +bool PreShapeNodeElimination::SatisfyCondition(const Graph& graph, const Node& node, const logging::Logger& logger) const { + if (!graph_utils::CanRemoveNode(graph, node, logger)) { + return false; + } + + auto output_nodes = graph.GetConsumerNodes(node.OutputDefs()[0]->Name()); + + if (output_nodes.empty()) { + return false; + } + + for (const Node* next_node : output_nodes) { + // Check if the next node is not of type "Shape" + if (!graph_utils::IsSupportedOptypeVersionAndDomain(*next_node, "Shape", {13, 15, 19}, kOnnxDomain)) { + return false; + } + } + + // All output nodes are of type "Shape" + return true; +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/pre_shape_node_elimination.h b/onnxruntime/core/optimizer/pre_shape_node_elimination.h new file mode 100644 index 0000000000..91e9257bab --- /dev/null +++ b/onnxruntime/core/optimizer/pre_shape_node_elimination.h @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/rewrite_rule.h" + +namespace onnxruntime { + +/** +@class PreShapeNodeElimination + +Rewrite rule that eliminates some particular nodes if the next node is a Shape node. + +It is attempted to be triggered only on nodes with op type "Cast". +*/ +class PreShapeNodeElimination : public RewriteRule { + public: + PreShapeNodeElimination() noexcept : RewriteRule("PreShapeNodeElimination") {} + + std::vector TargetOpTypes() const noexcept override { + return {"Cast"}; + } + + private: + bool SatisfyCondition(const Graph& graph, const Node& node, const logging::Logger& logger) const override; + + Status Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_effect, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 6c614419dd..84f2fbcfd4 100755 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -81,6 +81,7 @@ #include "test/util/include/inference_session_wrapper.h" #include "test/util/include/temp_dir.h" #include "test/util/include/test_utils.h" +#include "core/optimizer/pre_shape_node_elimination.h" #ifdef ENABLE_TRAINING #include "orttraining/core/optimizer/bitmask_dropout_replacement.h" #endif @@ -3160,6 +3161,69 @@ TEST_F(GraphTransformationTests, CastElimination) { ASSERT_TRUE(op_to_count["Cast"] == 4); } +TEST_F(GraphTransformationTests, PreShapeNodeElimination) { + constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "pre_shape_node_elimination.onnx"; + std::shared_ptr model; + ASSERT_TRUE(Model::Load(model_uri, model, nullptr, *logger_).IsOK()); + Graph& graph = model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_TRUE(op_to_count["Cast"] == 3); + + auto rule_transformer_L1 = std::make_unique("RuleTransformer1"); + ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + op_to_count = CountOpsInGraph(graph); + + ASSERT_TRUE(op_to_count["Cast"] == 2); + + // Assert that the remaining "Cast" nodes have different names than "cast2" + bool names_are_different = true; + for (const Node& node : graph.Nodes()) { + if (node.OpType() == "Cast") { + const std::string& node_name = node.Name(); + if (node_name == "cast") { + names_are_different = false; + break; + } + } + } + + ASSERT_TRUE(names_are_different); + + auto pre_graph_checker = [](Graph& graph) { + TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Cast"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [&](Graph& graph) { + TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Cast"] == 0); + return Status::OK(); + }; + + // cast is the first node. + { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input_arg = builder.MakeInput({{2, 3, 3, 3}}); + auto* cast_out = builder.MakeIntermediate(); + auto* shape_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Cast", {input_arg}, {cast_out}) + .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); + builder.AddNode("Shape", {cast_out}, {shape_out}); + builder.AddNode("Identity", {shape_out}, {output}); + }; + + auto rule_transformer = std::make_unique("RuleTransformer"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::move(rule_transformer), TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); + } +} + #ifndef DISABLE_CONTRIB_OPS static void ValidateAttention(Graph& graph) { diff --git a/onnxruntime/test/testdata/transform/pre_shape_node_elimination.onnx b/onnxruntime/test/testdata/transform/pre_shape_node_elimination.onnx new file mode 100644 index 0000000000..f609450ca4 Binary files /dev/null and b/onnxruntime/test/testdata/transform/pre_shape_node_elimination.onnx differ diff --git a/onnxruntime/test/testdata/transform/pre_shape_node_elimination.py b/onnxruntime/test/testdata/transform/pre_shape_node_elimination.py new file mode 100644 index 0000000000..a6fdc09f95 --- /dev/null +++ b/onnxruntime/test/testdata/transform/pre_shape_node_elimination.py @@ -0,0 +1,40 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +import onnx +from onnx import OperatorSetIdProto, TensorProto, helper + +input_shape = [2, 2] + +input = helper.make_tensor_value_info("input", TensorProto.FLOAT, input_shape) +output = helper.make_tensor_value_info("output", TensorProto.FLOAT, input_shape) + +cast1 = helper.make_node("Cast", ["input"], ["cast1"], name="cast1", to=2) +shape1 = helper.make_node("Shape", ["cast1"], ["shape1"], name="shape1") +add = helper.make_node("Add", ["cast1", "cast1"], ["add"], name="add") +cast2 = helper.make_node("Cast", ["add"], ["cast2"], name="cast2", to=1) +shape2 = helper.make_node("Shape", ["cast2"], ["shape2"], name="shape2") +shape3 = helper.make_node("Shape", ["cast2"], ["shape3"], name="shape3") +add_shapes = helper.make_node("Add", ["shape2", "shape3"], ["add_shapes"], name="add_shapes") +cast3 = helper.make_node("Cast", ["add_shapes"], ["output"], name="cast3", to=1) + +graph_def = helper.make_graph( + [cast1, shape1, add, cast2, shape2, shape3, add_shapes, cast3], + "pre_shape_node_elimination.onnx", + [input], + [output], +) + +opsets = [] +onnxdomain = OperatorSetIdProto() +onnxdomain.version = 15 +onnxdomain.domain = "" +opsets.append(onnxdomain) + +kwargs = {} +kwargs["opset_imports"] = opsets + +model_def = helper.make_model(graph_def, producer_name="onnx-example", **kwargs) +onnx.save(model_def, "pre_shape_node_elimination.onnx") diff --git a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc index c23206c190..b5655e12c1 100644 --- a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc +++ b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc @@ -61,6 +61,7 @@ #include "orttraining/core/optimizer/qdq_fusion.h" #include "orttraining/core/optimizer/shape_optimizer.h" #include "orttraining/core/optimizer/transformer_layer_recompute.h" +#include "core/optimizer/pre_shape_node_elimination.h" #include "core/optimizer/compute_optimizer/upstream_gather.h" #include "core/optimizer/compute_optimizer/upstream_reshape.h" #include "orttraining/core/optimizer/compute_optimizer/padding_elimination.h" @@ -93,6 +94,7 @@ std::vector> GeneratePreTrainingTransformers( ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); + ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique()));