RISC-V Vector Extension (RVV) support for ATen

.github/workflows/ci-riscv64.yml

@@ -0,0 +1,213 @@
+# Note: this runner is provided externally, so we minimize its access to
+# secrets.
+
+name: CI (riscv64)
+
+on:
+  push:
+    branches: [riscv]
+
+  pull_request_target:
+    types: [opened, synchronize, reopened]
+
+permissions:
+  contents: read
+  # No permissions to secrets.
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+# FIXME: Drop this
+env:
+  RUSTFLAGS: -D warnings
+  CARGO_TERM_COLOR: always
+
+jobs:
+  core-ci:
+    name: pytorch-riscv64-core-ci
+    runs-on: [self-hosted, linux, amd64]
+
+    outputs:
+      base_commit: ${{ steps.meta.outputs.base_commit }}
+      head_commit: ${{ steps.meta.outputs.head_commit }}
+      patch_file: ${{ steps.patch.outputs.patch_file }}
+      ci_result_base_url: ${{ steps.jenkins.outputs.ci_result_base_url }}
+      ci_stat_url: ${{ steps.jenkins.outputs.ci_stat_url }}
+
+    # This is in its own separate environment.
+    environment: riscv64
+    steps:
+      - name: Checkout repo
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 3000   # shadow clone?
+          ref: ${{ github.sha }} # including latest sha
+
+      - name: Extract PR info
+        run: |
+          echo "BASE_SHA=${{ github.event.pull_request.base.sha }}" >> $GITHUB_ENV
+          echo "HEAD_SHA=${{ github.event.pull_request.head.sha }}" >> $GITHUB_ENV
+
+      - name: Diff base and head
+        id: meta
+        run: |
+          if [[ "${{ github.event_name }}" = "pull_request" || "${{ github.event_name }}" == "pull_request_target" ]]; then
+            echo "Push PR build"
+            BASE_REF="${{ github.base_ref }}"
+            HEAD_REF="${{ github.head_ref }}"
+
+            echo "Base ref: $BASE_REF"
+            echo "Head ref: $HEAD_REF"
+
+            # must based on riscv
+            if [ "$BASE_REF" != "riscv" ]; then
+              echo "ERROR: PR must target 'riscv' branch, got '$BASE_REF'"
+              exit 1
+            fi
+
+            # need to get contents of the PR
+            git fetch --quiet origin pull/${{ github.event.pull_request.number }}/head:pr-head
+            git fetch --quiet origin main
+            BASE=$(git merge-base pr-head origin/main)
+            HEAD=$(git rev-parse pr-head)
+          else
+            echo "Push to riscv"
+            # 统一用 riscv 作为 baseline
+            git fetch --quiet origin main
+            #git fetch origin riscv
+
+            BASE=$(git merge-base ${{ github.sha }} origin/main) # The latest commit
+            HEAD=${{ github.sha }}
+
+          fi
+
+          echo "BASE_COMMIT=$BASE" >> $GITHUB_ENV
+          echo "HEAD_COMMIT=$HEAD" >> $GITHUB_ENV
+
+          echo "base_commit=$BASE" >> "$GITHUB_OUTPUT"
+          echo "head_commit=$HEAD" >> "$GITHUB_OUTPUT"
+
+          echo "Base: $BASE"
+          echo "Head: $HEAD"
+
+      - name: Generate patch
+        id: patch
+        run: |
+          echo "Generating patch..."
+
+          SHORT_HEAD=${HEAD_COMMIT:0:7}
+          PATCH_NAME="patch_${SHORT_HEAD}.patch"
+
+          git diff $BASE_COMMIT $HEAD_COMMIT > $PATCH_NAME
+
+          echo "Patch size:"
+          wc -l $PATCH_NAME
+
+          cp $PATCH_NAME /home/jenkins/patch/
+          cat /home/jenkins/patch/$PATCH_NAME
+
+          echo "PATCH_FILE=$PATCH_NAME" >> "$GITHUB_ENV"
+          echo "patch_file=$PATCH_NAME" >> "$GITHUB_OUTPUT"
+
+      - name: Trigger Jenkins Job and get the CI results
+        id: jenkins
+        run: |
+          set -euo pipefail
+
+          # Fail early if the runner does not provide the required tools.
+          for tool in curl jq; do
+            if ! command -v "$tool" >/dev/null 2>&1; then
+              echo "::error::Required command is unavailable: $tool"
+              exit 1
+            fi
+          done
+
+          BASE=${{ steps.meta.outputs.base_commit }}
+          PATCH=${{ steps.patch.outputs.patch_file }}
+          CORE_RESULT_FILE="$RUNNER_TEMP/pytorch-core-ci-stat.json"
+
+
+          set +e
+          bash /home/jenkins/scripts/jenkins-run.sh $BASE  $PATCH | tee jenkins.log
+          JENKINS_RC=${PIPESTATUS[0]}
+          set -e
+
+          CI_STAT_URL=$(grep -oE 'https://[^ ]+/pytorch-ci-stat\.json' jenkins.log | tail -n1 || true)
+
+          if [[ -z "$CI_STAT_URL" ]]; then
+            echo "ERROR: cannot find pytorch-ci-stat.json URL from Jenkins log"
+            echo "jenkins-run.sh rc=$JENKINS_RC"
+            exit 1
+          fi
+
+          CI_RESULT_BASE_URL="${CI_STAT_URL%/pytorch-ci-stat.json}"
+          CI_CORE_RESULT_URL="${CI_STAT_URL%/*}/pytorch-core-ci-stat.json"
+
+          echo "ci_stat_url=$CI_STAT_URL" >> "$GITHUB_OUTPUT"
+          echo "ci_result_base_url=$CI_RESULT_BASE_URL" >> "$GITHUB_OUTPUT"
+
+          echo "CI_STAT_URL=$CI_STAT_URL"
+          echo "CI_RESULT_BASE_URL=$CI_RESULT_BASE_URL"
+
+
+
+          if ! curl \
+            --fail \
+            --silent \
+            --show-error \
+            --location \
+            --retry 5 \
+           --retry-delay 5 \
+           --retry-connrefused \
+           --output "$CORE_RESULT_FILE" \
+           "$CI_CORE_RESULT_URL"; then
+           echo "::error::Failed to download $CI_CORE_RESULT_URL"
+           exit 1
+          fi
+
+          if jq -e '.failed == []' "$CORE_RESULT_FILE" >/dev/null; then
+            echo "Core CI succeeded: no failed tests"
+            exit 0
+          fi
+
+          echo "::error::Core CI reported failures"
+          echo "Failed tests:"
+          jq '.failed' "$CORE_RESULT_FILE" 2>/dev/null || cat "$CORE_RESULT_FILE"
+          exit 1
+
+  full-ci:
+    name: pytorch-riscv64-full-ci
+    runs-on: [self-hosted, linux, amd64]
+    needs: core-ci
+    if: always()
+    continue-on-error: true
+
+    steps:
+      - name: Query existing full test result
+        shell: bash
+        run: |
+          set -euo pipefail
+
+          BASE_URL="${{ needs.core-ci.outputs.ci_result_base_url }}"
+          STAT_URL="${BASE_URL}/pytorch-ci-stat.json"
+
+          echo "STAT_URL=$STAT_URL"
+
+          curl -fsSL "$STAT_URL" -o pytorch-ci-stat.json
+
+          echo "==== FULL TEST STAT ===="
+          cat pytorch-ci-stat.json
+          echo
+
+          FAILED=$(jq '.failed | length' pytorch-ci-stat.json)
+
+          if [[ "$FAILED" != "0" ]]; then
+            echo "==== FULL TEST FAILED ===="
+            echo "failed cases: $FAILED"
+            exit 1
+          fi
+
+          echo "==== FULL TEST PASSED ===="
+          echo "full test no failures"
+

.gitmodules

@@ -41,7 +41,7 @@
 [submodule "third_party/cpuinfo"]
     ignore = dirty
     path = third_party/cpuinfo
-    url = https://github.com/pytorch/cpuinfo.git
+    url = https://github.com/RuyiAI-Stack/cpuinfo.git
 [submodule "third_party/python-peachpy"]
     ignore = dirty
     path = third_party/python-peachpy

aten/src/ATen/CMakeLists.txt

@@ -74,7 +74,7 @@ if(NOT BUILD_LITE_INTERPRETER)
 endif()
 EXCLUDE(ATen_CORE_SRCS "${ATen_CORE_SRCS}" ${ATen_CORE_TEST_SRCS})
 
-file(GLOB base_h CONFIGURE_DEPENDS "*.h" "detail/*.h" "cpu/*.h" "cpu/vec/vec512/*.h" "cpu/vec/vec128/*.h" "cpu/vec/vec256/*.h" "cpu/vec/vec256/vsx/*.h" "cpu/vec/vec256/zarch/*.h" "cpu/vec/sve/*.h" "cpu/vec/*.h" "quantized/*.h" "functorch/*.h" "accelerator/*.h")
+file(GLOB base_h CONFIGURE_DEPENDS "*.h" "detail/*.h" "cpu/*.h" "cpu/vec/vec512/*.h" "cpu/vec/vec128/*.h" "cpu/vec/vec256/*.h" "cpu/vec/vec256/vsx/*.h" "cpu/vec/vec256/zarch/*.h" "cpu/vec/sve/*.h" "cpu/vec/rvv/*.h" "cpu/vec/*.h" "quantized/*.h" "functorch/*.h" "accelerator/*.h")
 file(GLOB base_cpp CONFIGURE_DEPENDS "*.cpp" "detail/*.cpp" "cpu/*.cpp" "functorch/*.cpp" "accelerator/*.cpp")
 file(GLOB cuda_h CONFIGURE_DEPENDS "cuda/*.h" "cuda/detail/*.h" "cuda/*.cuh" "cuda/detail/*.cuh" "cuda/tunable/*.cuh" "cuda/tunable/*.h")
 file(GLOB cuda_cpp CONFIGURE_DEPENDS "cuda/*.cpp" "cuda/detail/*.cpp" "cuda/tunable/*.cpp")

aten/src/ATen/Version.cpp

@@ -107,6 +107,9 @@ std::string get_cpu_capability() {
       return "SVE128";
     case native::CPUCapability::SVE256:
       return "SVE256";
+#elif defined(HAVE_RVV_CPU_DEFINITION)
+    case native::CPUCapability::RVV:
+      return "RVV";
 #else
     case native::CPUCapability::AVX2:
       return "AVX2";

aten/src/ATen/cpu/vec/rvv/rvv_helper.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <ATen/cpu/vec/intrinsics.h>
+#include <ATen/cpu/vec/vec_base.h>
+
+#if defined(CPU_CAPABILITY_RVV)
+
+typedef int8_t fixed_vint8m2_t[CONFIG_VLMAX / sizeof(int8_t)];
+typedef int16_t fixed_vint16m2_t[CONFIG_VLMAX / sizeof(int16_t)];
+typedef int32_t fixed_vint32m2_t[CONFIG_VLMAX / sizeof(int32_t)];
+typedef int64_t fixed_vint64m2_t[CONFIG_VLMAX / sizeof(int64_t)];
+
+typedef uint8_t fixed_vuint8m2_t[CONFIG_VLMAX / sizeof(uint8_t)];
+typedef uint16_t fixed_vuint16m2_t[CONFIG_VLMAX / sizeof(uint16_t)];
+typedef uint32_t fixed_vuint32m2_t[CONFIG_VLMAX / sizeof(uint32_t)];
+typedef uint64_t fixed_vuint64m2_t[CONFIG_VLMAX / sizeof(uint64_t)];
+
+typedef float fixed_vfloat32m2_t[CONFIG_VLMAX / sizeof(float)];
+typedef double fixed_vfloat64m2_t[CONFIG_VLMAX / sizeof(double)];
+
+#define VFLOAT32_VL (CONFIG_VLMAX_BITS / 32)
+#define VQINT8_VL (CONFIG_VLMAX_BITS / 8)
+#define VQUINT8_VL (CONFIG_VLMAX_BITS / 8)
+#define VQINT32_VL (CONFIG_VLMAX_BITS / 32)
+
+#endif // defined(CPU_CAPABILITY_RVV)

aten/src/ATen/cpu/vec/rvv/vec_bfloat16.h

@@ -0,0 +1,126 @@
+#pragma once
+
+#include <ATen/cpu/vec/intrinsics.h>
+#include <ATen/cpu/vec/vec_base.h>
+#include <c10/util/irange.h>
+
+namespace at {
+namespace vec {
+// See Note [CPU_CAPABILITY namespace]
+inline namespace CPU_CAPABILITY {
+
+// Generic implementation
+//  TO-DO : Add Zvfbfmin for rvv
+inline std::tuple<Vectorized<float>, Vectorized<float>> convert_bfloat16_float(
+    const Vectorized<BFloat16>& a) {
+  constexpr int64_t K = Vectorized<BFloat16>::size();
+  __at_align__ float arr[K];
+  __at_align__ BFloat16 arr2[K];
+  a.store(arr2);
+  convert(arr2, arr, K);
+  return std::make_tuple(
+      Vectorized<float>::loadu(arr),
+      Vectorized<float>::loadu(arr + Vectorized<float>::size()));
+}
+
+inline Vectorized<BFloat16> convert_float_bfloat16(
+    const Vectorized<float>& a,
+    const Vectorized<float>& b) {
+  constexpr int64_t K = Vectorized<BFloat16>::size();
+  __at_align__ float arr[K];
+  __at_align__ BFloat16 arr2[K];
+  a.store(arr);
+  b.store(arr + Vectorized<float>::size());
+  convert(arr, arr2, K);
+  return Vectorized<BFloat16>::loadu(arr2);
+}
+
+// Generic implementation
+// TO-DO : Add Zvfhmin for rvv
+inline std::tuple<Vectorized<float>, Vectorized<float>> convert_half_float(
+    const Vectorized<Half>& a) {
+  constexpr int64_t K = Vectorized<Half>::size();
+  __at_align__ float arr[K];
+  __at_align__ Half arr2[K];
+  a.store(arr2);
+  convert(arr2, arr, K);
+  return std::make_tuple(
+      Vectorized<float>::loadu(arr),
+      Vectorized<float>::loadu(arr + Vectorized<float>::size()));
+}
+
+inline Vectorized<Half> convert_float_half(
+    const Vectorized<float>& a,
+    const Vectorized<float>& b) {
+  constexpr int64_t K = Vectorized<Half>::size();
+  __at_align__ float arr[K];
+  __at_align__ Half arr2[K];
+  a.store(arr);
+  b.store(arr + Vectorized<float>::size());
+  convert(arr, arr2, K);
+  return Vectorized<Half>::loadu(arr2);
+};
+
+template <>
+inline Vectorized<Half> fmadd(
+    const Vectorized<Half>& a,
+    const Vectorized<Half>& b,
+    const Vectorized<Half>& c) {
+  const auto [a0, a1] = convert_half_float(a);
+  const auto [b0, b1] = convert_half_float(b);
+  const auto [c0, c1] = convert_half_float(c);
+  return convert_float_half(
+      at::vec::fmadd(a0, b0, c0), at::vec::fmadd(a1, b1, c1));
+}
+
+template <>
+inline Vectorized<Half> fmsub(
+    const Vectorized<Half>& a,
+    const Vectorized<Half>& b,
+    const Vectorized<Half>& c) {
+  const auto [a0, a1] = convert_half_float(a);
+  const auto [b0, b1] = convert_half_float(b);
+  const auto [c0, c1] = convert_half_float(c);
+  return convert_float_half(
+      at::vec::fmsub(a0, b0, c0), at::vec::fmsub(a1, b1, c1));
+}
+
+inline void load_fp32_from_bf16(
+    const c10::BFloat16* data,
+    Vectorized<float>& out) {
+  __at_align__ float values[Vectorized<float>::size()];
+  for (const auto k : c10::irange(Vectorized<float>::size())) {
+    values[k] = data[k];
+  }
+  out = Vectorized<float>::loadu(values);
+}
+
+inline void load_fp32_from_bf16(
+    const c10::BFloat16* data,
+    Vectorized<float>& out1,
+    Vectorized<float>& out2) {
+  load_fp32_from_bf16(data, out1);
+  data += Vectorized<float>::size();
+  load_fp32_from_bf16(data, out2);
+}
+
+inline void load_fp32_from_fp16(const c10::Half* data, Vectorized<float>& out) {
+  __at_align__ float values[Vectorized<float>::size()];
+  for (const auto k : c10::irange(Vectorized<float>::size())) {
+    values[k] = data[k];
+  }
+  out = Vectorized<float>::loadu(values);
+}
+
+inline void load_fp32_from_fp16(
+    const c10::Half* data,
+    Vectorized<float>& out1,
+    Vectorized<float>& out2) {
+  load_fp32_from_fp16(data, out1);
+  data += Vectorized<float>::size();
+  load_fp32_from_fp16(data, out2);
+}
+
+} // namespace CPU_CAPABILITY
+} // namespace vec
+} // namespace at