Fix use-after-free in BM_FeedInputFetchOutput().
We were holding on to Node* in a graph whose ownership was transferred to
test::Benchmark(). This was not previously an issue, because the executor
implementation would hold on to a copy of the graph until the executor itself
was deleted, but it became an issue when the executor now that the executor no
longer stores a copy of the graph.
PiperOrigin-RevId: 273515081
diff --git a/tensorflow/core/common_runtime/executor_test.cc b/tensorflow/core/common_runtime/executor_test.cc
index 3167c91..e994512 100644
--- a/tensorflow/core/common_runtime/executor_test.cc
+++ b/tensorflow/core/common_runtime/executor_test.cc
@@ -479,12 +479,18 @@
Node* y = test::graph::Recv(g, "y", "float", ALICE, 1, BOB);
Node* sum = test::graph::Add(g, x, y);
Node* z = test::graph::Send(g, sum, "z", BOB, 1, ALICE);
+
+ string x_key = test::GetRendezvousKey(x);
+ string y_key = test::GetRendezvousKey(y);
+ string z_key = test::GetRendezvousKey(z);
+
Tensor val(DT_FLOAT, TensorShape({}));
val.scalar<float>()() = 3.14;
#ifdef PLATFORM_GOOGLE
SetBenchmarkItemsProcessed(static_cast<int64>(iters));
#endif // PLATFORM_GOOGLE
- test::Benchmark("cpu", g).RunWithArgs({{x, val}, {y, val}}, {z}, iters);
+ test::Benchmark("cpu", g).RunWithRendezvousArgs({{x_key, val}, {y_key, val}},
+ {z_key}, iters);
}
BENCHMARK(BM_FeedInputFetchOutput);
diff --git a/tensorflow/core/common_runtime/kernel_benchmark_testlib.cc b/tensorflow/core/common_runtime/kernel_benchmark_testlib.cc
index fe34083..fe70305 100644
--- a/tensorflow/core/common_runtime/kernel_benchmark_testlib.cc
+++ b/tensorflow/core/common_runtime/kernel_benchmark_testlib.cc
@@ -117,7 +117,7 @@
}
}
-void Benchmark::Run(int iters) { RunWithArgs({}, {}, iters); }
+void Benchmark::Run(int iters) { RunWithRendezvousArgs({}, {}, iters); }
string GetRendezvousKey(const Node* node) {
string send_device;
@@ -133,23 +133,12 @@
recv_device, tensor_name, FrameAndIter(0, 0));
}
-void Benchmark::RunWithArgs(
- const std::vector<std::pair<const Node*, Tensor>>& inputs,
- const std::vector<const Node*>& outputs, int iters) {
+void Benchmark::RunWithRendezvousArgs(
+ const std::vector<std::pair<string, Tensor>>& inputs,
+ const std::vector<string>& outputs, int iters) {
if (!device_ || iters == 0) {
return;
}
- // Gets inputs' and outputs' rendezvous keys.
- std::vector<std::pair<string, Tensor>> in;
- in.reserve(inputs.size());
- for (const auto& p : inputs) {
- in.push_back({GetRendezvousKey(p.first), p.second});
- }
- std::vector<string> out;
- out.reserve(outputs.size());
- for (const auto& n : outputs) {
- out.push_back(GetRendezvousKey(n));
- }
Tensor unused; // In benchmark, we don't care the return value.
bool is_dead;
@@ -161,13 +150,13 @@
};
static const int kWarmupRuns = 3;
for (int i = 0; i < kWarmupRuns; ++i) {
- for (const auto& p : in) {
+ for (const auto& p : inputs) {
Rendezvous::ParsedKey parsed;
TF_CHECK_OK(Rendezvous::ParseKey(p.first, &parsed));
TF_CHECK_OK(rendez_->Send(parsed, Rendezvous::Args(), p.second, false));
}
TF_CHECK_OK(exec_->Run(args));
- for (const string& key : out) {
+ for (const string& key : outputs) {
Rendezvous::ParsedKey parsed;
TF_CHECK_OK(Rendezvous::ParseKey(key, &parsed));
TF_CHECK_OK(rendez_->Recv(parsed, Rendezvous::Args(), &unused, &is_dead));
@@ -178,13 +167,13 @@
testing::StartTiming();
while (iters-- > 0) {
- for (const auto& p : in) {
+ for (const auto& p : inputs) {
Rendezvous::ParsedKey parsed;
TF_CHECK_OK(Rendezvous::ParseKey(p.first, &parsed));
TF_CHECK_OK(rendez_->Send(parsed, Rendezvous::Args(), p.second, false));
}
TF_CHECK_OK(exec_->Run(args));
- for (const string& key : out) {
+ for (const string& key : outputs) {
Rendezvous::ParsedKey parsed;
TF_CHECK_OK(Rendezvous::ParseKey(key, &parsed));
TF_CHECK_OK(rendez_->Recv(parsed, Rendezvous::Args(), &unused, &is_dead));
diff --git a/tensorflow/core/common_runtime/kernel_benchmark_testlib.h b/tensorflow/core/common_runtime/kernel_benchmark_testlib.h
index b1557c5..742f40d 100644
--- a/tensorflow/core/common_runtime/kernel_benchmark_testlib.h
+++ b/tensorflow/core/common_runtime/kernel_benchmark_testlib.h
@@ -46,12 +46,13 @@
void Run(int iters);
// If "g" contains send/recv nodes, before each execution, we send
- // inputs to the corresponding recv nodes in the graph, after each
- // execution, we recv outputs from the corresponding send nodes in
+ // inputs to the corresponding recv keys in the graph, after each
+ // execution, we recv outputs from the corresponding send keys in
// the graph. In the benchmark, we throw away values returned by the
// graph.
- void RunWithArgs(const std::vector<std::pair<const Node*, Tensor>>& inputs,
- const std::vector<const Node*>& outputs, int iters);
+ void RunWithRendezvousArgs(
+ const std::vector<std::pair<string, Tensor>>& inputs,
+ const std::vector<string>& outputs, int iters);
private:
thread::ThreadPool* pool_ = nullptr;
@@ -62,6 +63,9 @@
TF_DISALLOW_COPY_AND_ASSIGN(Benchmark);
};
+// Returns the rendezvous key associated with the given Send/Recv node.
+string GetRendezvousKey(const Node* node);
+
} // end namespace test
} // end namespace tensorflow
