ENH: Enable RVV CPU feature detection

ksco · ksco · commit d1caa616cb96 · 2024-04-07T13:29:05.000+08:00
Enhanced compile-time and runtime CPU feature detection for RISC-V vector
extension, which will aid future RVV support.
diff --git a/meson_cpu/main_config.h.in b/meson_cpu/main_config.h.in
@@ -385,4 +385,8 @@
 #ifdef @P@HAVE_NEON
     #include <arm_neon.h>
 #endif
+
+#ifdef @P@HAVE_RVV
+    #include <riscv_vector.h>
+#endif
 #endif // @P@_CPU_DISPATCHER_CONF_H_
diff --git a/meson_cpu/meson.build b/meson_cpu/meson.build
@@ -75,12 +75,14 @@ subdir('x86')
 subdir('ppc64')
 subdir('s390x')
 subdir('arm')
+subdir('riscv64')
 
 CPU_FEATURES = {}
 CPU_FEATURES += ARM_FEATURES
 CPU_FEATURES += X86_FEATURES
 CPU_FEATURES += PPC64_FEATURES
 CPU_FEATURES += S390X_FEATURES
+CPU_FEATURES += RV64_FEATURES
 
 # Parse the requested baseline (CPU_CONF_BASELINE) and dispatch features
 # (CPU_CONF_DISPATCH).
@@ -93,6 +95,7 @@ min_features = {
   's390x': [],
   'arm': [],
   'aarch64': [ASIMD],
+  'riscv64': [],
   'wasm32': [],
 }.get(cpu_family, [])
 if host_machine.endian() == 'little' and cpu_family == 'ppc64'
@@ -107,6 +110,7 @@ max_features_dict = {
   's390x': S390X_FEATURES,
   'arm': ARM_FEATURES,
   'aarch64': ARM_FEATURES,
+  'riscv64': RV64_FEATURES,
   'wasm32': {},
 }.get(cpu_family, {})
 max_features = []
diff --git a/meson_cpu/riscv64/meson.build b/meson_cpu/riscv64/meson.build
@@ -0,0 +1,8 @@
+source_root = meson.project_source_root()
+mod_features = import('features')
+
+RVV = mod_features.new(
+  'RVV', 1, args: ['-march=rv64gcv'],
+  test_code: files(source_root + '/numpy/distutils/checks/cpu_rvv.c')[0],
+)
+RV64_FEATURES = {'RVV': RVV}
diff --git a/numpy/_core/src/common/npy_cpu_features.c b/numpy/_core/src/common/npy_cpu_features.c
@@ -119,7 +119,8 @@ static struct {
                 {NPY_CPU_FEATURE_ASIMDHP, "ASIMDHP"},
                 {NPY_CPU_FEATURE_ASIMDDP, "ASIMDDP"},
                 {NPY_CPU_FEATURE_ASIMDFHM, "ASIMDFHM"},
-                {NPY_CPU_FEATURE_SVE, "SVE"}};
+                {NPY_CPU_FEATURE_SVE, "SVE"},
+                {NPY_CPU_FEATURE_RVV, "RVV"}};
 
 
 NPY_VISIBILITY_HIDDEN PyObject *
@@ -813,6 +814,28 @@ npy__cpu_init_features(void)
 #endif
 }
 
+/************** RISC-V 64 ***************/
+
+#elif defined(__riscv) && __riscv_xlen == 64
+
+#include <sys/auxv.h>
+
+#ifndef HWCAP_RVV
+    // https://github.com/torvalds/linux/blob/v6.8/arch/riscv/include/uapi/asm/hwcap.h#L24
+    #define COMPAT_HWCAP_ISA_V	(1 << ('V' - 'A'))
+#endif
+
+static void
+npy__cpu_init_features(void)
+{
+    memset(npy__cpu_have, 0, sizeof(npy__cpu_have[0]) * NPY_CPU_FEATURE_MAX);
+
+    unsigned int hwcap = getauxval(AT_HWCAP);
+    if (hwcap & COMPAT_HWCAP_ISA_V) {
+        npy__cpu_have[NPY_CPU_FEATURE_RVV]  = 1;
+    }
+}
+
 /*********** Unsupported ARCH ***********/
 #else
 static void
diff --git a/numpy/_core/src/common/npy_cpu_features.h b/numpy/_core/src/common/npy_cpu_features.h
@@ -98,6 +98,9 @@ enum npy_cpu_features
     // Vector-Enhancements Facility 2
     NPY_CPU_FEATURE_VXE2              = 352,
 
+    // RISC-V
+    NPY_CPU_FEATURE_RVV               = 400,
+
     NPY_CPU_FEATURE_MAX
 };
 
diff --git a/numpy/distutils/checks/cpu_rvv.c b/numpy/distutils/checks/cpu_rvv.c
@@ -0,0 +1,13 @@
+#ifndef __riscv_vector
+  #error RVV not supported
+#endif
+
+#include <riscv_vector.h>
+
+int main(void)
+{
+    size_t vlmax = __riscv_vsetvlmax_e32m1();
+    vuint32m1_t a = __riscv_vmv_v_x_u32m1(0, vlmax);
+    vuint32m1_t b = __riscv_vadd_vv_u32m1(a, a, vlmax);
+    return __riscv_vmv_x_s_u32m1_u32(b);
+}
