Refactor compiler version checks and arch flags

2025-08-27 09:26:02 +08:00
parent 3a93f4eb28
commit 0e49c3353b
2 changed files with 18 additions and 17 deletions
--- a/csrc/jit/device_runtime.hpp
+++ b/csrc/jit/device_runtime.hpp
@@ -25,11 +25,15 @@ public:
        return {prop->major, prop->minor};
    }

-    std::string get_arch() {
+    std::string get_arch(const bool& number_only = false,
+                         const bool& support_arch_family = false) {
        const auto& [major, minor] = get_arch_pair();
-        if (major == 10 and minor != 1)
-            return "100f";
-        return std::to_string(major * 10 + minor) + "a";
+        if (major == 10 and minor != 1) {
+            if (number_only)
+                return "100";
+            return support_arch_family ? "100f" : "100a";
+        }
+        return std::to_string(major * 10 + minor) + (number_only ? "" : "a");
    }

    int get_arch_major() {