Utils.h 803 B

123456789101112131415161718192021222324252627282930313233
  1. #pragma once
  2. #include <cstdint>
  3. #include <c10/macros/Export.h>
  4. namespace at::cpu {
  5. TORCH_API bool is_avx2_supported();
  6. TORCH_API bool is_avx512_supported();
  7. // Detect if CPU support Vector Neural Network Instruction.
  8. TORCH_API bool is_avx512_vnni_supported();
  9. // Detect if CPU supports AVX512_BF16 ISA
  10. TORCH_API bool is_avx512_bf16_supported();
  11. // Detect if CPU support Advanced Matrix Extension.
  12. TORCH_API bool is_amx_tile_supported();
  13. // Detect if CPU support Advanced Matrix Extension for fp16.
  14. TORCH_API bool is_amx_fp16_supported();
  15. // Enable the system to use AMX instructions.
  16. TORCH_API bool init_amx();
  17. // Get the L1 cache size per core in Byte
  18. TORCH_API uint32_t L1d_cache_size();
  19. // Get the L2 cache size per core in Byte
  20. TORCH_API uint32_t L2_cache_size();
  21. } // namespace at::cpu