torch.cpu.get_capabilities

torch.cpu.get_capabilities() [исходный код]

Возвращает неизменяемое отображение возможностей процессора, обнаруженных во время выполнения.

Эта функция с помощью cpuinfo запрашивает у процессора поддерживаемые наборы инструкций и функции. Для повышения эффективности результат кэшируется после первого вызова.

Возвращаемое отображение содержит возможности, специфичные для архитектуры:

Для x86/x86_64:

• Семейство SSE: sse, sse2, sse3, ssse3, sse4_1, sse4_2, sse4a
• Семейство AVX: avx, avx2, avx_vnni
• Семейство AVX-512: avx512_f, avx512_cd, avx512_dq, avx512_bw, avx512_vl, avx512_ifma, avx512_vbmi, avx512_vbmi2, avx512_bitalg, avx512_vpopcntdq, avx512_vnni, avx512_bf16, avx512_fp16, avx512_vp2intersect, avx512_4vnniw, avx512_4fmaps
• Семейство AVX10: avx10_1, avx10_2
• AVX-VNNI-INT: avx_vnni_int8, avx_vnni_int16, avx_ne_convert
• AMX: amx_bf16, amx_tile, amx_int8, amx_fp16
• FMA: fma3, fma4
• Прочее: f16c, bmi, bmi2, popcnt, lzcnt, aes, sha, clflush, clflushopt, clwb

Для ARM64:

• SIMD: neon, fp16_arith, bf16, i8mm, dot
• SVE: sve, sve2, sve_bf16, sve_max_length (если поддерживается)
• SME: sme, sme2, sme_max_length (если поддерживается)
• Прочее: atomics, fhm, rdm, crc32, aes, sha1, sha2, pmull

Общие для всех архитектур:

• architecture: строка, идентифицирующая архитектуру процессора

Возвращает:

Неизменяемое отображение, в котором ключи — это названия возможностей (например, ‘avx2’, ‘sve’), а значения — логические значения, указывающие на поддержку, либо целые числа для таких свойств, как длина векторов.

Тип возвращаемого значения:

MappingProxyType

Пример

>>> caps = torch.cpu.get_capabilities()
>>> if caps.get("avx2", False):
...     print("AVX2 is supported")
>>> print(f"Architecture: {caps['architecture']}")