contraction_benchmarks_cpu.cc 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839
  1. #define EIGEN_USE_THREADS
  2. #include <string>
  3. #include "tensor_benchmarks.h"
  4. #define CREATE_THREAD_POOL(threads) \
  5. Eigen::ThreadPool pool(threads); \
  6. Eigen::ThreadPoolDevice device(&pool, threads);
  7. // Contractions for number of threads ranging from 1 to 32
  8. // Dimensions are Rows, Cols, Depth
  9. #define BM_ContractionCPU(D1, D2, D3) \
  10. static void BM_##Contraction##_##D1##x##D2##x##D3(int iters, int Threads) { \
  11. StopBenchmarkTiming(); \
  12. CREATE_THREAD_POOL(Threads); \
  13. BenchmarkSuite<Eigen::ThreadPoolDevice, float> suite(device, D1, D2, D3); \
  14. suite.contraction(iters); \
  15. } \
  16. BENCHMARK_RANGE(BM_##Contraction##_##D1##x##D2##x##D3, 1, 32);
  17. // Vector Matrix and Matrix Vector products
  18. BM_ContractionCPU(1, 2000, 500);
  19. BM_ContractionCPU(2000, 1, 500);
  20. // Various skinny matrices
  21. BM_ContractionCPU(250, 3, 512);
  22. BM_ContractionCPU(1500, 3, 512);
  23. BM_ContractionCPU(512, 800, 4);
  24. BM_ContractionCPU(512, 80, 800);
  25. BM_ContractionCPU(512, 80, 13522);
  26. BM_ContractionCPU(1, 80, 13522);
  27. BM_ContractionCPU(3200, 512, 4);
  28. BM_ContractionCPU(3200, 512, 80);
  29. BM_ContractionCPU(3200, 80, 512);