From 079cf71458060f57cc4089eefbe7192402af2a1f Mon Sep 17 00:00:00 2001 From: Matthias Braun Date: Fri, 19 Dec 2025 07:25:43 -0800 Subject: [PATCH] Choose _autovec version of GenerateEmbeddingSpMDMRowWiseSparse on AArch64 (#5247) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2238 `GenerateEmbeddingSpMDMRowWiseSparse_autovec` should be working fine and preferred over the `_Ref` version of the algorithm. Reviewed By: Nicoshev Differential Revision: D88891882 --- src/EmbeddingSpMDM.cc | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/EmbeddingSpMDM.cc b/src/EmbeddingSpMDM.cc index 8dbcc78420..b2ca0df574 100644 --- a/src/EmbeddingSpMDM.cc +++ b/src/EmbeddingSpMDM.cc @@ -1310,8 +1310,11 @@ typename EmbeddingSpMDMKernelSignature:: } #ifdef FBGEMM_AUTOVEC_AVAILABLE - if (!is_autovec_disabled()) { - // There is only the reference implementation for FP8 embedding + if (!cpuinfo_initialize()) { + throw std::runtime_error("Failed to initialize cpuinfo!"); + } + if ((is_autovec_forced() || fbgemmHasArmSve2Support()) && + !is_autovec_disabled()) { return GenerateEmbeddingSpMDMFP8WithStrides_autovec< /*IndexType=*/indxType, /*OffsetType=*/offsetType, @@ -1473,7 +1476,11 @@ GenerateEmbeddingSpMDMRowWiseSparse( #endif #ifdef FBGEMM_AUTOVEC_AVAILABLE - if (is_autovec_forced()) { + if (!cpuinfo_initialize()) { + throw std::runtime_error("Failed to initialize cpuinfo!"); + } + if ((is_autovec_forced() || fbgemmHasArmSve2Support()) && + !is_autovec_disabled()) { return GenerateEmbeddingSpMDMRowWiseSparse_autovec< /*InType=*/inType, /*IndexType=*/indxType,