xref: /aosp_15_r20/external/pytorch/aten/src/ATen/cpu/Utils.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <cstdint>
4 
5 #include <c10/macros/Export.h>
6 
7 namespace at::cpu {
8 
9 TORCH_API bool is_avx2_supported();
10 TORCH_API bool is_avx512_supported();
11 
12 // Detect if CPU support Vector Neural Network Instruction.
13 TORCH_API bool is_avx512_vnni_supported();
14 
15 // Detect if CPU supports AVX512_BF16 ISA
16 TORCH_API bool is_avx512_bf16_supported();
17 
18 // Detect if CPU support Advanced Matrix Extension.
19 TORCH_API bool is_amx_tile_supported();
20 
21 // Enable the system to use AMX instructions.
22 TORCH_API bool init_amx();
23 
24 // Get the L1 cache size per core in Byte
25 TORCH_API uint32_t L1d_cache_size();
26 
27 // Get the L2 cache size per core in Byte
28 TORCH_API uint32_t L2_cache_size();
29 
30 } // namespace at::cpu
31