57 lines
2.6 KiB
Diff
57 lines
2.6 KiB
Diff
|
From ee119d4a498979525046fba1c3dd3f13a039fbb1 Mon Sep 17 00:00:00 2001
|
||
|
|
From: Daniel Ellis <danielellis@google.com>
|
||
|
|
Date: Wed, 14 Jul 2021 12:43:17 -0700
|
||
|
|
Subject: [PATCH] Fix segmentation fault in shape inference logic.
|
||
|
|
|
||
|
|
When running shape functions, some functions (such as `MutableHashTableShape`)
|
||
|
|
produce extra output information in the form of a `ShapeAndType` struct. The
|
||
|
|
shapes embedded in this struct are owned by an inference context that is
|
||
|
|
cleaned up almost immediately; if the upstream code attempts to access this
|
||
|
|
shape information, it can trigger a segfault.
|
||
|
|
|
||
|
|
`ShapeRefiner` is mitigating this for normal output shapes by cloning them
|
||
|
|
(and thus putting the newly created shape under ownership of an inference
|
||
|
|
context that will not die), but we were not doing the same for shapes and
|
||
|
|
types. This commit fixes that by doing similar logic on output shapes and
|
||
|
|
types.
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 384761124
|
||
|
|
Change-Id: I07c0c42d29dfbb55bfa13ec1f09ef825fb0a1a1d
|
||
|
|
---
|
||
|
|
.../core/common_runtime/shape_refiner.cc | 21 +++++++++++++++++--
|
||
|
|
1 file changed, 19 insertions(+), 2 deletions(-)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/common_runtime/shape_refiner.cc b/tensorflow/core/common_runtime/shape_refiner.cc
|
||
|
|
index 375f809b31b36..2e29ef48189a5 100644
|
||
|
|
--- a/tensorflow/core/common_runtime/shape_refiner.cc
|
||
|
|
+++ b/tensorflow/core/common_runtime/shape_refiner.cc
|
||
|
|
@@ -120,9 +120,26 @@ Status ShapeRefiner::InferShapesForFunctionSubNode(
|
||
|
|
TF_RETURN_IF_ERROR(outer_context->MakeShapeFromShapeProto(proto, &handle));
|
||
|
|
outer_context->set_output(index, handle);
|
||
|
|
|
||
|
|
- auto* resource = node_context->input_handle_shapes_and_types(0);
|
||
|
|
+ const std::vector<ShapeAndType>* resource =
|
||
|
|
+ node_context->input_handle_shapes_and_types(0);
|
||
|
|
if (resource) {
|
||
|
|
- outer_context->set_output_handle_shapes_and_types(index, *resource);
|
||
|
|
+ // `ShapesAndType`s contain `ShapeHandle`s. These `ShapeHandle`s point
|
||
|
|
+ // to `Shape`s that are owned by a different inference context too. We
|
||
|
|
+ // need to copy them to the outer context to prevent them from being
|
||
|
|
+ // destroyed before they are used.
|
||
|
|
+ std::vector<ShapeAndType> copied_shapes_and_types;
|
||
|
|
+ for (auto& shape_and_type : *resource) {
|
||
|
|
+ ShapeHandle handle;
|
||
|
|
+ TensorShapeProto proto;
|
||
|
|
+ node_context->ShapeHandleToProto(shape_and_type.shape, &proto);
|
||
|
|
+ TF_RETURN_IF_ERROR(
|
||
|
|
+ outer_context->MakeShapeFromShapeProto(proto, &handle));
|
||
|
|
+ copied_shapes_and_types.push_back(
|
||
|
|
+ ShapeAndType(handle, shape_and_type.dtype, shape_and_type.type));
|
||
|
|
+ }
|
||
|
|
+
|
||
|
|
+ outer_context->set_output_handle_shapes_and_types(
|
||
|
|
+ index, copied_shapes_and_types);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|