From f31bc2ea60f685a65885fc3c4c7753e7f1ca5a61 Mon Sep 17 00:00:00 2001 From: velonica0 <47554626+velonica0@users.noreply.github.com> Date: Thu, 11 Jun 2026 12:09:05 +0800 Subject: [PATCH] [CPU][RISC-V] Enable oneDNN W8A8 INT8 to run on RISC-V (#44478) Signed-off-by: velonica0 --- cmake/cpu_extension.cmake | 9 ++++++--- csrc/cpu/cpu_types_riscv_defs.hpp | 4 ++++ csrc/cpu/cpu_types_riscv_impl.hpp | 29 +++++++++++++++++++++++++++++ csrc/cpu/torch_bindings.cpp | 5 +++-- 4 files changed, 42 insertions(+), 5 deletions(-) diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 6f836ff535..e3e9b75030 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -166,6 +166,10 @@ elseif (S390_FOUND) "-mtune=native") elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") message(STATUS "RISC-V detected") + if(DEFINED VLLM_RVV_VLEN AND NOT VLLM_RVV_VLEN GREATER 0) + message(FATAL_ERROR + "VLLM_RVV_VLEN must be a positive integer; got '${VLLM_RVV_VLEN}'") + endif() # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256. if(NOT DEFINED VLLM_RVV_VLEN) @@ -189,8 +193,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") "RISC-V RVV is available but VLEN could not be auto-detected. " "Please specify VLEN explicitly:\n" " -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n" - " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)\n" - " -DVLLM_RVV_VLEN=0 (force scalar, no RVV)") + " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)") endif() endif() if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0) @@ -219,7 +222,7 @@ endif() # Build oneDNN for GEMM kernels -if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) +if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND OR RVV_FP16_FOUND OR RVV_BF16_FOUND) # Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64 # TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "") diff --git a/csrc/cpu/cpu_types_riscv_defs.hpp b/csrc/cpu/cpu_types_riscv_defs.hpp index 8871617f05..650dc5bcc7 100644 --- a/csrc/cpu/cpu_types_riscv_defs.hpp +++ b/csrc/cpu/cpu_types_riscv_defs.hpp @@ -57,6 +57,10 @@ typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t __attribute__((riscv_rvv_vector_bits(1024))); +// int8 +typedef RVVTYPE(vint8, LMUL_128, _t) fixed_i8x16_t + __attribute__((riscv_rvv_vector_bits(128))); + // int32 typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t __attribute__((riscv_rvv_vector_bits(256))); diff --git a/csrc/cpu/cpu_types_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp index 06a38c780a..a8c178db4c 100644 --- a/csrc/cpu/cpu_types_riscv_impl.hpp +++ b/csrc/cpu/cpu_types_riscv_impl.hpp @@ -734,10 +734,18 @@ struct FP32Vec16 : public Vec { return FP32Vec16( RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); } + FP32Vec16 max(const FP32Vec16& b, const int elem_num) const { + return FP32Vec16( + RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, elem_num)); + } FP32Vec16 min(const FP32Vec16& b) const { return FP32Vec16( RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); } + FP32Vec16 min(const FP32Vec16& b, const int elem_num) const { + return FP32Vec16( + RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, elem_num)); + } FP32Vec16 abs() const { return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM)); } @@ -867,6 +875,27 @@ struct FP32Vec16 : public Vec { } }; +struct INT8Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + fixed_i8x16_t reg; + + explicit INT8Vec16(const FP32Vec16& vec) { + auto i32_vec = + RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(vec.reg, VEC_ELEM_NUM); + auto i16_vec = RVVI(__riscv_vnclip_wx_i16, LMUL_256)( + i32_vec, 0, __RISCV_VXRM_RNU, VEC_ELEM_NUM); + reg = RVVI(__riscv_vnclip_wx_i8, LMUL_128)(i16_vec, 0, __RISCV_VXRM_RNU, + VEC_ELEM_NUM); + } + + void save(int8_t* ptr) const { + RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, VEC_ELEM_NUM); + } + void save(int8_t* ptr, int elem_num) const { + RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, elem_num); + } +}; + // ============================================================================ // Type Traits & Global Helpers // ============================================================================ diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 7a8188b8c8..c5ce7c46bb 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -329,8 +329,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding); // Quantization -#if defined(__AVX512F__) || defined(__AVX2__) || \ - (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) +#if defined(__AVX512F__) || defined(__AVX2__) || \ + (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \ + defined(__riscv_v) // Helper function to release oneDNN handlers ops.def("release_dnnl_matmul_handler(int handler) -> ()", &release_dnnl_matmul_handler);