61 lines
2.6 KiB
Diff
61 lines
2.6 KiB
Diff
From e6340f0665d53716ef3197ada88936c2a5f7a2d3 Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
Date: Mon, 19 Apr 2021 09:57:01 -0700
|
|
Subject: [PATCH] Handle a special grappler case resulting in crash.
|
|
|
|
It might happen that a malformed input could be used to trick Grappler into trying to optimize a node with no inputs. This, in turn, would produce a null pointer dereference and a segfault.
|
|
|
|
PiperOrigin-RevId: 369242852
|
|
Change-Id: I2e5cbe7aec243d34a6d60220ac8ac9b16f136f6b
|
|
---
|
|
.../core/grappler/optimizers/arithmetic_optimizer.cc | 11 +++++++++++
|
|
.../core/grappler/optimizers/dependency_optimizer.cc | 6 ++++++
|
|
2 files changed, 17 insertions(+)
|
|
|
|
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
|
|
index 484983a7a549c..774fdf3484788 100644
|
|
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
|
|
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
|
|
@@ -2047,6 +2047,12 @@ class ReorderCastLikeAndValuePreserving : public ArithmeticOptimizerStage {
|
|
|
|
Status TrySimplify(NodeDef* consumer, string* simplified_node_name) override {
|
|
NodeDef* producer;
|
|
+
|
|
+ if (consumer->input_size() < 1) {
|
|
+ return errors::FailedPrecondition("Node ", simplified_node_name,
|
|
+ " lacks inputs");
|
|
+ }
|
|
+
|
|
TF_RETURN_IF_ERROR(GetInputNode(consumer->input(0), &producer));
|
|
const bool producer_is_cast = IsCastLike(*producer);
|
|
const bool can_optimize =
|
|
@@ -2538,6 +2544,11 @@ class ReplaceMulWithSquare : public ArithmeticOptimizerStage {
|
|
~ReplaceMulWithSquare() override = default;
|
|
|
|
bool IsSupported(const NodeDef* node) const override {
|
|
+ if (!node || node->input_size() < 2) {
|
|
+ // Invalid node
|
|
+ return false;
|
|
+ }
|
|
+
|
|
return IsAnyMul(*node) && node->input(0) == node->input(1);
|
|
}
|
|
|
|
diff --git a/tensorflow/core/grappler/optimizers/dependency_optimizer.cc b/tensorflow/core/grappler/optimizers/dependency_optimizer.cc
|
|
index 332d2d4e6d966..1be7f2692e0f7 100644
|
|
--- a/tensorflow/core/grappler/optimizers/dependency_optimizer.cc
|
|
+++ b/tensorflow/core/grappler/optimizers/dependency_optimizer.cc
|
|
@@ -68,6 +68,12 @@ bool DependencyOptimizer::SafeToRemoveIdentity(const NodeDef& node) const {
|
|
// The output values of this node may be needed.
|
|
return false;
|
|
}
|
|
+
|
|
+ if (node.input_size() < 1) {
|
|
+ // Node lacks input, is invalid
|
|
+ return false;
|
|
+ }
|
|
+
|
|
const NodeDef* input = node_map_->GetNode(NodeName(node.input(0)));
|
|
CHECK(input != nullptr) << "node = " << node.name()
|
|
<< " input = " << node.input(0);
|