Enable constant folding in L0 optimization level.

Change: 118861866

1 files changed, 21 insertions, 0 deletions

diff --git a/tensorflow/core/common_runtime/memory_types_test.cc b/tensorflow/core/common_runtime/memory_types_test.cc
index a5ac254494..9193ac85c7 100644
--- a/tensorflow/core/common_runtime/memory_types_test.cc
+++ b/tensorflow/core/common_runtime/memory_types_test.cc
@@ -52,4 +52,25 @@ TEST(MemoryTypeChecker, Int32NotOk) {
   delete g;
 }
 
+TEST(MemoryTypeChecker, MemoryTypeForOutput) {
+  Graph* g = new Graph(OpRegistry::Global());
+  Tensor vb(DT_BOOL, {});
+  Tensor vi(DT_INT32, {});
+  Tensor vf(DT_FLOAT, {});
+  auto pred = test::graph::Constant(g, vb);
+  auto si = test::graph::Switch(g, test::graph::Constant(g, vi), pred);
+  auto sf = test::graph::Switch(g, test::graph::Constant(g, vf), pred);
+  MemoryType memory_type;
+
+  TF_EXPECT_OK(MemoryTypeForOutput(DEVICE_CPU, g, sf, 0, &memory_type));
+  // float Switch's output on CPU doesn't have HOST_MEMORY constraint.
+  EXPECT_EQ(memory_type, DEVICE_MEMORY);
+#if GOOGLE_CUDA
+  TF_EXPECT_OK(MemoryTypeForOutput(DEVICE_GPU, g, si, 0, &memory_type));
+  // int Switch's output on GPU has HOST_MEMORY constraint.
+  EXPECT_EQ(memory_type, HOST_MEMORY);
+#endif  // GOOGLE_CUDA
+  delete g;
+}
+
 }  // namespace