| 1 | #ifndef CPU_TYPES_HPP |
| 2 | #define CPU_TYPES_HPP |
| 3 | |
| 4 | #if defined(__x86_64__) |
| 5 | // x86 implementation |
| 6 | #include "cpu_types_x86.hpp" |
| 7 | #elif defined(__powerpc__) |
| 8 | // ppc implementation |
| 9 | #include "cpu_types_vsx.hpp" |
| 10 | #elif defined(__s390x__) |
| 11 | // s390 implementation |
| 12 | #include "cpu_types_vxe.hpp" |
| 13 | #elif defined(__aarch64__) |
| 14 | // arm implementation |
| 15 | #include "cpu_types_arm.hpp" |
| 16 | #elif defined(__riscv_v) |
| 17 | // riscv implementation |
| 18 | #include "cpu_types_riscv.hpp" |
| 19 | #else |
| 20 | #warning "unsupported vLLM cpu implementation, vLLM will compile with scalar" |
| 21 | #include "cpu_types_scalar.hpp" |
| 22 | #endif |
| 23 | |
| 24 | #ifdef _OPENMP |
| 25 | #include <omp.h> |
| 26 | #endif |
| 27 | |
| 28 | #include <c10/util/Exception.h> |
| 29 | |
| 30 | namespace cpu_utils { |
| 31 | // Without OpenMP the omp pragmas compile to serial loops, so report 1: kernels |
| 32 | // that barrier on the thread count would otherwise deadlock. |
| 33 | inline int get_max_threads() { |
| 34 | #ifdef _OPENMP |
| 35 | return omp_get_max_threads(); |
| 36 | #else |
| 37 | TORCH_WARN_ONCE( |
| 38 | "vLLM CPU was built without OpenMP; running single-threaded."); |
| 39 | return 1; |
| 40 | #endif |
| 41 | } |
| 42 | } // namespace cpu_utils |
| 43 | |
| 44 | #endif |