1 #include <cassert> 2 #include <ATen/cpu/Utils.h> 3 #if !defined(__s390x__ ) && !defined(__powerpc__) 4 #include <cpuinfo.h> 5 #endif 6 #if defined(__linux__) 7 #include <sys/syscall.h> 8 #include <unistd.h> 9 #endif 10 11 namespace at::cpu { is_avx2_supported()12bool is_avx2_supported() { 13 #if !defined(__s390x__) && !defined(__powerpc__) 14 return cpuinfo_initialize() && cpuinfo_has_x86_avx2(); 15 #else 16 return false; 17 #endif 18 } 19 is_avx512_supported()20bool is_avx512_supported() { 21 #if !defined(__s390x__) && !defined(__powerpc__) 22 return cpuinfo_initialize() && cpuinfo_has_x86_avx512f() && cpuinfo_has_x86_avx512vl() && cpuinfo_has_x86_avx512bw() && cpuinfo_has_x86_avx512dq(); 23 #else 24 return false; 25 #endif 26 } 27 is_avx512_vnni_supported()28bool is_avx512_vnni_supported() { 29 #if !defined(__s390x__) && !defined(__powerpc__) 30 return cpuinfo_initialize() && cpuinfo_has_x86_avx512vnni(); 31 #else 32 return false; 33 #endif 34 } 35 is_avx512_bf16_supported()36bool is_avx512_bf16_supported() { 37 #if !defined(__s390x__) && !defined(__powerpc__) 38 return cpuinfo_initialize() && cpuinfo_has_x86_avx512bf16(); 39 #else 40 return false; 41 #endif 42 } 43 is_amx_tile_supported()44bool is_amx_tile_supported() { 45 #if !defined(__s390x__) && !defined(__powerpc__) 46 return cpuinfo_initialize() && cpuinfo_has_x86_amx_tile(); 47 #else 48 return false; 49 #endif 50 } 51 init_amx()52bool init_amx() { 53 if (!is_amx_tile_supported()) { 54 return false; 55 } 56 57 #if defined(__linux__) && !defined(__ANDROID__) && defined(__x86_64__) 58 #define XFEATURE_XTILECFG 17 59 #define XFEATURE_XTILEDATA 18 60 #define XFEATURE_MASK_XTILECFG (1 << XFEATURE_XTILECFG) 61 #define XFEATURE_MASK_XTILEDATA (1 << XFEATURE_XTILEDATA) 62 #define XFEATURE_MASK_XTILE (XFEATURE_MASK_XTILECFG | XFEATURE_MASK_XTILEDATA) 63 64 #define ARCH_GET_XCOMP_PERM 0x1022 65 #define ARCH_REQ_XCOMP_PERM 0x1023 66 67 unsigned long bitmask = 0; 68 // Request permission to use AMX instructions 69 long rc = syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA); 70 if (rc) { 71 return false; 72 } 73 // Check if the system supports AMX instructions 74 rc = syscall(SYS_arch_prctl, ARCH_GET_XCOMP_PERM, &bitmask); 75 if (rc) { 76 return false; 77 } 78 if (bitmask & XFEATURE_MASK_XTILE) { 79 return true; 80 } 81 return false; 82 #else 83 return true; 84 #endif 85 } 86 get_cache_size(int level)87static uint32_t get_cache_size(int level) { 88 #if !defined(__s390x__) && !defined(__powerpc__) 89 if (!cpuinfo_initialize()) { 90 return 0; 91 } 92 const struct cpuinfo_processor* processors = cpuinfo_get_processors(); 93 if (!processors) { 94 return 0; 95 } 96 const struct cpuinfo_cache* cache = nullptr; 97 switch (level) { 98 case 1: 99 cache = processors[0].cache.l1d; 100 break; 101 case 2: 102 cache = processors[0].cache.l2; 103 break; 104 default: 105 assert(false && "Unsupported cache level"); 106 } 107 108 if (!cache) { 109 return 0; 110 } 111 return cache->size; 112 #else 113 return 0; 114 #endif 115 } 116 L1d_cache_size()117uint32_t L1d_cache_size() { 118 return get_cache_size(1); 119 } 120 L2_cache_size()121uint32_t L2_cache_size() { 122 return get_cache_size(2); 123 } 124 125 } // namespace at::cpu 126