Eliminate safe nodes that are followed by a shape node. (#16065)

### Description
Eliminate Cast operator if Shape is the next one.

### Motivation and Context
#### Cast
When working with onnx opset 15 and above, the shape operator now
accepts all types of variables.
This change is documented in the [onnx
Changelog](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Shape-15).

As a result, casting variables right before the shape operation becomes
unnecessary.
Removing these unnecessary casts will improve the graph and potentially
provide better performance gains.


## Results
On :
torchrun examples/onnxruntime/training/language-modeling/run_clm.py
--model_name_or_path gpt2 --do_train --overwrite_output_dir --output_dir
./outputs/ --seed 1337 --fp16 True --per_device_train_batch_size 4
--num_train_epochs 1 --dataset_name wikitext --dataset_config_name
wikitext-2-raw-v1 --learning_rate 2e-5 --report_to none --optim
adamw_ort_fused

without changes:
***** train metrics *****
  epoch                    =        1.0
  train_loss               =     3.2981
  train_runtime            = 0:02:13.29
  train_samples            =       2318
  train_samples_per_second =      17.39
  train_steps_per_second   =      4.351

With my changes:
***** train metrics *****
  epoch                    =        1.0
  train_loss               =     3.2981
  train_runtime            = 0:02:08.98
  train_samples            =       2318
  train_samples_per_second =     17.971
  train_steps_per_second   =      4.497

We see around 3% gain.

---------

Co-authored-by: Adam Louly <adamlouly@microsoft.com@orttrainingdev9.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>

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<std::unique_ptr<RewriteRule>> GenerateRewriteRules(
       rules.push_back(std::make_unique<EliminateDropout>());
       rules.push_back(std::make_unique<ExpandElimination>());
       rules.push_back(std::make_unique<CastElimination>());
+      rules.push_back(std::make_unique<PreShapeNodeElimination>());
       rules.push_back(std::make_unique<NoopElimination>());
       rules.push_back(std::make_unique<DivMulFusion>());
       rules.push_back(std::make_unique<FuseReluClip>());

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<size_t> 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

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<std::string> 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

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> model;
+  ASSERT_TRUE(Model::Load(model_uri, model, nullptr, *logger_).IsOK());
+  Graph& graph = model->MainGraph();
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+  ASSERT_TRUE(op_to_count["Cast"] == 3);
+
+  auto rule_transformer_L1 = std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1");
+  ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique<PreShapeNodeElimination>()));
+  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<float>({{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<int64_t>(ONNX_NAMESPACE::TensorProto_DataType_FLOAT));
+      builder.AddNode("Shape", {cast_out}, {shape_out});
+      builder.AddNode("Identity", {shape_out}, {output});
+    };
+
+    auto rule_transformer = std::make_unique<RuleBasedGraphTransformer>("RuleTransformer");
+    ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique<PreShapeNodeElimination>()));
+    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) {

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")

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<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<UnsqueezeElimination>()));
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<ExpandElimination>()));
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<CastElimination>()));
+      ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<PreShapeNodeElimination>()));
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<NoopElimination>()));
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<DivMulFusion>()));
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<EliminateDropout>()));