Fix resource misclassification for unsorted _read_only_resource_inputs - #110665
Conversation
Replace pointer-based linear scan with set-based lookup in _get_read_only_resource_input_indices_op and get_read_write_resource_inputs. The previous implementation assumed the _read_only_resource_inputs attribute was always sorted, causing read-only resources to be misclassified as read-write when indices were unsorted, injecting spurious control dependencies. Fixes tensorflow#110516
dmiltr3
left a comment
There was a problem hiding this comment.
Thanks for the contribution! The logic fix looks correct and will resolve issues with unsorted read-only resource indices.
Could you please add a unit test in tensorflow/python/framework/auto_control_deps_test.py to prevent regressions for this edge case? You can use the test snippet below, which verifies the classification behavior for unsorted, sorted, empty, and mixed index configurations:
def testReadOnlyResourceInputsUnsorted(self):
from tensorflow.python.framework import auto_control_deps_utils as acd_utils
from unittest import mock
# 1. Unsorted inputs
input0 = mock.MagicMock()
input0.dtype = dtypes.resource
input1 = mock.MagicMock()
input1.dtype = dtypes.int32
input2 = mock.MagicMock()
input2.dtype = dtypes.resource
op = mock.MagicMock()
op.type = "SomeOp"
op.inputs = [input0, input1, input2]
op.get_attr.side_effect = lambda name: [2, 0] if name == acd_utils.READ_ONLY_RESOURCE_INPUTS_ATTR else ValueError()
# _get_read_only_resource_input_indices_op should return both indices, sorted
self.assertEqual(
acd_utils._get_read_only_resource_input_indices_op(op), [0, 2]
)
# get_read_write_resource_inputs should categorize both as reads and zero as writes
reads, writes = acd_utils.get_read_write_resource_inputs(op)
self.assertIn(input0, reads)
self.assertIn(input2, reads)
self.assertEqual(len(writes), 0)
# 2. Empty indices
op2 = mock.MagicMock()
op2.type = "SomeOp"
op2.inputs = [input0, input1, input2]
op2.get_attr.side_effect = lambda name: [] if name == acd_utils.READ_ONLY_RESOURCE_INPUTS_ATTR else ValueError()
self.assertEqual(
acd_utils._get_read_only_resource_input_indices_op(op2), []
)
reads2, writes2 = acd_utils.get_read_write_resource_inputs(op2)
self.assertEqual(len(reads2), 0)
self.assertIn(input0, writes2)
self.assertIn(input2, writes2)
# 3. Mixed reads and writes (index 0 is read, index 2 is write)
op3 = mock.MagicMock()
op3.type = "SomeOp"
op3.inputs = [input0, input1, input2]
op3.get_attr.side_effect = lambda name: [0] if name == acd_utils.READ_ONLY_RESOURCE_INPUTS_ATTR else ValueError()
self.assertEqual(
acd_utils._get_read_only_resource_input_indices_op(op3), [0]
)
reads3, writes3 = acd_utils.get_read_write_resource_inputs(op3)
self.assertIn(input0, reads3)
self.assertNotIn(input2, reads3)
self.assertIn(input2, writes3)
self.assertNotIn(input0, writes3)Covers unsorted, empty, and mixed read/write resource index configurations for _get_read_only_resource_input_indices_op and get_read_write_resource_inputs, per review feedback.
|
@dmiltr3 thanks for the contribution and the detailed test snippet! Added it as Two small adjustments to fit the file's style, with identical coverage/behavior to your snippet:
PTAL 🙏 |
Summary
auto_control_deps_utils.pywhere_get_read_only_resource_input_indices_opandget_read_write_resource_inputsused a pointer-based linear scan that assumed_read_only_resource_inputswas always sorted.[2, 0]), read-only resources were misclassified as read-write, injecting spurious control dependencies and serializing operations that should run in parallel.Fixes #110516
Test plan
[2, 0]now correctly detect both indices[0, 2]get_read_write_resource_inputscorrectly classifies reads vs writes for unsorted indices