1 #include <ATen/cpu/Utils.h> 2 #include <torch/csrc/cpu/Module.h> 3 #include <torch/csrc/utils/pybind.h> 4 5 namespace torch::cpu { 6 initModule(PyObject * module)7void initModule(PyObject* module) { 8 auto m = py::handle(module).cast<py::module>(); 9 10 auto cpu = m.def_submodule("_cpu", "cpu related pybind."); 11 cpu.def("_is_avx2_supported", at::cpu::is_avx2_supported); 12 cpu.def("_is_avx512_supported", at::cpu::is_avx512_supported); 13 cpu.def("_is_avx512_vnni_supported", at::cpu::is_avx512_vnni_supported); 14 cpu.def("_is_avx512_bf16_supported", at::cpu::is_avx512_bf16_supported); 15 cpu.def("_is_amx_tile_supported", at::cpu::is_amx_tile_supported); 16 cpu.def("_init_amx", at::cpu::init_amx); 17 cpu.def("_L1d_cache_size", at::cpu::L1d_cache_size); 18 cpu.def("_L2_cache_size", at::cpu::L2_cache_size); 19 } 20 21 } // namespace torch::cpu 22