Skip to content
Permalink
Browse files Browse the repository at this point in the history
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
  • Loading branch information
dellis23 authored and tensorflower-gardener committed Jul 14, 2021
1 parent f118ff1 commit ee119d4
Showing 1 changed file with 19 additions and 2 deletions.
21 changes: 19 additions & 2 deletions tensorflow/core/common_runtime/shape_refiner.cc
Expand Up @@ -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);
}
}

Expand Down

0 comments on commit ee119d4

Please sign in to comment.