From ee119d4a498979525046fba1c3dd3f13a039fbb1 Mon Sep 17 00:00:00 2001 From: Daniel Ellis 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* 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 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); } }