Add RISC-V Vector (RVV) CPU capability

.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/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/native/DispatchStub.cpp

@@ -62,6 +62,10 @@ static CPUCapability compute_cpu_capability() {
         return CPUCapability::DEFAULT;
       }
     }
+#elif defined(HAVE_RVV_CPU_DEFINITION)
+    if (envar == "rvv") {
+      return CPUCapability::RVV;
+    }
 #else
 #ifdef HAVE_AVX512_CPU_DEFINITION
     if (envar == "avx512") {
@@ -116,6 +120,11 @@ static CPUCapability compute_cpu_capability() {
       return CPUCapability::SVE128;
   }
 #endif
+#if defined(__linux__) && defined(HAVE_RVV_CPU_DEFINITION)
+  if (cpuinfo_initialize() && cpuinfo_has_riscv_v()) {
+    return CPUCapability::RVV;
+  }
+#endif
 #ifdef HAVE_VSX_CPU_DEFINITION
   return CPUCapability::VSX;
 #else
@@ -147,6 +156,9 @@ DispatchResult DispatchStubImpl::try_get_call_ptr(
   , void *SVE128
   , void *SVE256
 #endif
+#ifdef HAVE_RVV_CPU_DEFINITION
+  , void *RVV
+#endif
 ) {
   constexpr auto supported_devices = std::to_array<c10::DeviceType>(
         {c10::DeviceType::CPU,
@@ -185,6 +197,9 @@ DispatchResult DispatchStubImpl::try_get_call_ptr(
 #ifdef HAVE_SVE_CPU_DEFINITION
           , SVE128
           , SVE256
+#endif
+#ifdef HAVE_RVV_CPU_DEFINITION
+          , RVV
 #endif
         );
         if (!std::holds_alternative<ErrorType>(result)) {
@@ -243,6 +258,9 @@ void* DispatchStubImpl::get_call_ptr(
   , void *SVE128
   , void *SVE256
 #endif
+#ifdef HAVE_RVV_CPU_DEFINITION
+  , void *RVV
+#endif
 ) {
 
   auto result = try_get_call_ptr(
@@ -269,6 +287,10 @@ void* DispatchStubImpl::get_call_ptr(
       SVE128
       ,
       SVE256
+#endif
+#ifdef HAVE_RVV_CPU_DEFINITION
+      ,
+      RVV
 #endif
   );
   if (std::holds_alternative<ErrorType>(result)) {
@@ -303,6 +325,9 @@ DispatchResult DispatchStubImpl::try_choose_cpu_impl(
 #ifdef HAVE_SVE_CPU_DEFINITION
     , void *SVE128
     , void *SVE256
+#endif
+#ifdef HAVE_RVV_CPU_DEFINITION
+    , void *RVV
 #endif
   ){
 
@@ -349,6 +374,11 @@ DispatchResult DispatchStubImpl::try_choose_cpu_impl(
     }
     return DispatchResult(SVE256);
   }
+#endif
+#ifdef HAVE_RVV_CPU_DEFINITION
+  if (capability >= static_cast<int>(CPUCapability::RVV)) {
+    return RVV != nullptr ? DispatchResult(RVV) : ErrorType::MissingDeviceKernel;
+  }
 #endif
   return DEFAULT != nullptr ? DispatchResult(DEFAULT) : ErrorType::MissingDeviceKernel;
 }
@@ -371,6 +401,9 @@ void* DispatchStubImpl::choose_cpu_impl(
   , void *SVE128
   , void *SVE256
 #endif
+#ifdef HAVE_RVV_CPU_DEFINITION
+  , void *RVV
+#endif
 ) {
   auto capability = static_cast<int>(get_cpu_capability());
   (void)capability;
@@ -421,6 +454,12 @@ void* DispatchStubImpl::choose_cpu_impl(
     }
     return SVE256;
   }
+#endif
+#ifdef HAVE_RVV_CPU_DEFINITION
+  if (capability >= static_cast<int>(CPUCapability::RVV)) {
+    TORCH_INTERNAL_ASSERT(RVV, "DispatchStub: missing RVV kernel");
+    return RVV;
+  }
 #endif
   TORCH_INTERNAL_ASSERT(DEFAULT, "DispatchStub: missing default kernel");
   return DEFAULT;