Function TEST

tensorflow/core/framework/device_base_test.cc:28–59

namespace tensorflow {

TEST(DeviceBaseTest, CpuDevice) {
  DeviceBase dbase(Env::Default());
  thread::ThreadPool pool(Env::Default(), "test", 16);
  Eigen::ThreadPoolDevice eigen_device(pool.AsEigenThreadPool(),
                                       pool.NumThreads());
  ASSERT_FALSE(dbase.has_eigen_cpu_device());
  dbase.set_eigen_cpu_device(&eigen_device);
  ASSERT_TRUE(dbase.has_eigen_cpu_device());

  {
    auto d = dbase.eigen_cpu_device();
    EXPECT_EQ(d->numThreads(), 16);
  }

  {
    ScopedPerThreadMaxParallelism maxp(4);
    auto d = dbase.eigen_cpu_device();
    EXPECT_EQ(d->numThreads(), 4);
  }

  {
    ScopedPerThreadMaxParallelism maxp(1);
    auto d = dbase.eigen_cpu_device();
    EXPECT_EQ(d->numThreads(), 1);
  }

  {
    ScopedPerThreadMaxParallelism maxp(1000);
    auto d = dbase.eigen_cpu_device();
    EXPECT_EQ(d->numThreads(), 16);
  }
}

}  // namespace tensorflow