return architecture levels for micromamba (#2921)

* return architecture levels for micromamba * fix formatting * use function multiversioning only on x86 * fix formatting * update test * use __builtin_cpu_supports and check for more features Co-authored-by: Marcel Bargull <mbargull@users.noreply.github.com> --------- Co-authored-by: Marcel Bargull <mbargull@users.noreply.github.com>
mamba-org · Oct 23, 2023 · 99278ff · 99278ff
1 parent 24f26a0
commit 99278ff
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 2 deletions.
diff --git a/libmamba/src/core/virtual_packages.cpp b/libmamba/src/core/virtual_packages.cpp
@@ -163,6 +163,33 @@ namespace mamba
             return res;
         }
 
+        std::string get_archspec_x86_64()
+        {
+#if (defined(__GNUC__) || defined(__clang__)) && __x86_64__
+            /* if (__builtin_cpu_supports ("x86-64-v4")) */
+            if (__builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512bw")
+                && __builtin_cpu_supports("avx512cd") && __builtin_cpu_supports("avx512dq")
+                && __builtin_cpu_supports("avx512vl"))
+            {
+                return "x86_64-v4";
+            }
+            /* if (__builtin_cpu_supports ("x86-64-v3")) */
+            if (__builtin_cpu_supports("avx") && __builtin_cpu_supports("avx2")
+                && __builtin_cpu_supports("bmi") && __builtin_cpu_supports("bmi2")
+                && __builtin_cpu_supports("fma"))
+            {
+                return "x86_64-v3";
+            }
+            /* if (__builtin_cpu_supports ("x86-64-v2")) */
+            if (__builtin_cpu_supports("popcnt") && __builtin_cpu_supports("sse3")
+                && __builtin_cpu_supports("ssse3") && __builtin_cpu_supports("sse4.1")
+                && __builtin_cpu_supports("sse4.2"))
+            {
+                return "x86_64-v2";
+            }
+#endif
+            return "x86_64";
+        }
         std::vector<PackageInfo> dist_packages(const Context& context)
         {
             LOG_DEBUG << "Loading distribution virtual packages";
@@ -222,7 +249,7 @@ namespace mamba
 
             if (arch == "64")
             {
-                arch = "x86_64";
+                arch = get_archspec_x86_64();
             }
             else if (arch == "32")
             {

diff --git a/libmamba/tests/src/core/test_virtual_packages.cpp b/libmamba/tests/src/core/test_virtual_packages.cpp
@@ -73,7 +73,7 @@ namespace mamba
                 }
 #if __x86_64__ || defined(_WIN64)
                 CHECK_EQ(pkgs.back().name, "__archspec");
-                CHECK_EQ(pkgs.back().build_string, "x86_64");
+                CHECK_EQ(pkgs.back().build_string.find("x86_64"), 0);
 #endif
 
                 // This is bad design, tests should not interfer