PR #102607: Add support for Linux RISC-V 64

Imported from GitHub PR tensorflow/tensorflow#102607

This commit fixes #102159 which prevented TensorFlow from compiling on RISC-V 64 due to missing codepath in the build system.

RFC: Currently this patch overrides timeout of rules_python for all platforms which is obviously not ideal. How do we set pip package timeout here in Tensorflow?

See also:
- bazelbuild/bazel#25236
- bazelbuild/bazel#25699
- bazel-contrib/rules_python#3350
- #32812
- tensorflow/runtime#135
Copybara import of the project:

--
13b7db883e6ea0aa631490fb5f806c55ceea7cb5 by gns <[email protected]>:

[XLA:CPU] Add support for riscv64

Co-authored-by: Levi Zim <[email protected]>

--
a00c3a3ba94664798960cbc2fb0a6746fe0c3ac1 by gns <[email protected]>:

Add missing llvm:: namespace qualifiers for DTensor

Add the `llvm::` namespace prefix to `cast` and `isa`
where it was missing in the DTensor MLIR code.

--
febc5abc56ab302f5f39fe20176ed00c0c95cb7e by gns <[email protected]>:

runtime: add missing stdint header

Co-authored-by: Levi Zim <[email protected]>

--
6bf63fec56a6d961b06d932fe69b76eae89d8cb7 by gns <[email protected]>:

Add riscv64 support patch for rules_python

bazel-contrib/rules_python#3350

--
68c44fa9613d5577f36c51a192c932795e13ed83 by gns <[email protected]>:

tools: py: add riscv64 to pip and manylinux compilance test

--
3f2fe1e0256476449f64658219d997727c9cc534 by gns <[email protected]>:

Lift timeout of rules_python for BFS wheels

--
252a9b7ad2336b7abbb94e269b4a976b45d5d6e7 by gns <[email protected]>:

Revert llvm toolchain patch for riscv64

As already upstreamed.

Co-authored-by: Levi Zim <[email protected]>

--
51047a86131e991ffea74fccb8f55e3521237a64 by gns <[email protected]>:

Lift timeout for BFS wheels during pip init

--
3b8276420031969423c9a1f150341a0e41fa372e by gns <[email protected]>:

Refresh `rules_python` riscv64 patch

Co-authored-by: Levi Zim <[email protected]>

Merging this change closes #102607

FUTURE_COPYBARA_INTEGRATE_REVIEW=tensorflow/tensorflow#102607 from infiWang:riscv64 3b8276420031969423c9a1f150341a0e41fa372e
PiperOrigin-RevId: 825073762

Author: infiWang

third_party/py/python_init_pip.bzl

@@ -39,4 +39,5 @@ cc_library(
             "numpy": ["numpy_headers"],
         },
         requirements_lock = REQUIREMENTS_WITH_LOCAL_WHEELS,
+        timeout = 3600,
     )

third_party/py/python_init_rules.bzl

@@ -44,5 +44,6 @@ def python_init_rules(extra_patches = []):
             Label("//third_party/py:rules_python_pip_version.patch"),
             Label("//third_party/py:rules_python_freethreaded.patch"),
             Label("//third_party/py:rules_python_versions.patch"),
+            Label("//third_party/py:rules_python_riscv64_pypi.patch"),
         ] + extra_patches,
     )

third_party/py/rules_python_riscv64_pypi.patch

@@ -0,0 +1,131 @@
+From: Levi Zim <[email protected]>
+Date: Sun, 26 Oct 2025 12:35:53 +0800
+Subject: [PATCH] fix: Add linux_riscv64 to _pip_repository_impl (#3350)
+
+Add `linux_riscv64` support for pulling pip dependencies. This is not
+adding any hermetic toolchain support - user has to provide a working
+toolchain.
+
+Fix #2729
+
+---------
+
+Co-authored-by: Ignas Anikevicius <[email protected]>
+---
+ python/private/pypi/pip_repository.bzl                    | 1 +
+ python/private/pypi/whl_installer/platform.py             | 3 +++
+ python/private/pypi/whl_target_platforms.bzl              | 1 +
+ tests/pypi/whl_installer/platform_test.py                 | 6 +++---
+ .../whl_target_platforms/whl_target_platforms_tests.bzl   | 8 ++++++++
+ 5 files changed, 16 insertions(+), 3 deletions(-)
+
+diff --git a/python/private/pypi/pip_repository.bzl b/python/private/pypi/pip_repository.bzl
+index e9a4c44da3..d635651039 100644
+--- a/python/private/pypi/pip_repository.bzl
++++ b/python/private/pypi/pip_repository.bzl
+@@ -96,6 +96,7 @@ def _pip_repository_impl(rctx):
+                 "linux_aarch64",
+                 "linux_arm",
+                 "linux_ppc",
++                "linux_riscv64",
+                 "linux_s390x",
+                 "linux_x86_64",
+                 "osx_aarch64",
+diff --git a/python/private/pypi/whl_installer/platform.py b/python/private/pypi/whl_installer/platform.py
+index ff267fe4aa..0757d86990 100644
+--- a/python/private/pypi/whl_installer/platform.py
++++ b/python/private/pypi/whl_installer/platform.py
+@@ -45,6 +45,7 @@ class Arch(Enum):
+     ppc64le = 5
+     s390x = 6
+     arm = 7
++    riscv64 = 8
+     amd64 = x86_64
+     arm64 = aarch64
+     i386 = x86_32
+@@ -269,6 +270,8 @@ def platform_machine(self) -> str:
+             return "ppc"
+         elif self.arch == Arch.ppc64le:
+             return "ppc64le"
++        elif self.arch == Arch.riscv64:
++            return "riscv64"
+         elif self.arch == Arch.s390x:
+             return "s390x"
+         else:
+diff --git a/python/private/pypi/whl_target_platforms.bzl b/python/private/pypi/whl_target_platforms.bzl
+index 6c3dd5da83..28547c679c 100644
+--- a/python/private/pypi/whl_target_platforms.bzl
++++ b/python/private/pypi/whl_target_platforms.bzl
+@@ -30,6 +30,7 @@ _CPU_ALIASES = {
+     "ppc": "ppc",
+     "ppc64": "ppc",
+     "ppc64le": "ppc64le",
++    "riscv64": "riscv64",
+     "s390x": "s390x",
+     "arm": "arm",
+     "armv6l": "arm",
+diff --git a/tests/pypi/whl_installer/platform_test.py b/tests/pypi/whl_installer/platform_test.py
+index ad65650779..0d944bb196 100644
+--- a/tests/pypi/whl_installer/platform_test.py
++++ b/tests/pypi/whl_installer/platform_test.py
+@@ -38,17 +38,17 @@ def test_can_get_specific_from_string(self):
+ 
+     def test_can_get_all_for_py_version(self):
+         cp39 = Platform.all(minor_version=9, micro_version=0)
+-        self.assertEqual(21, len(cp39), f"Got {cp39}")
++        self.assertEqual(24, len(cp39), f"Got {cp39}")
+         self.assertEqual(cp39, Platform.from_string("cp39.0_*"))
+ 
+     def test_can_get_all_for_os(self):
+         linuxes = Platform.all(OS.linux, minor_version=9)
+-        self.assertEqual(7, len(linuxes))
++        self.assertEqual(8, len(linuxes))
+         self.assertEqual(linuxes, Platform.from_string("cp39_linux_*"))
+ 
+     def test_can_get_all_for_os_for_host_python(self):
+         linuxes = Platform.all(OS.linux)
+-        self.assertEqual(7, len(linuxes))
++        self.assertEqual(8, len(linuxes))
+         self.assertEqual(linuxes, Platform.from_string("linux_*"))
+ 
+     def test_platform_sort(self):
+diff --git a/tests/pypi/whl_target_platforms/whl_target_platforms_tests.bzl b/tests/pypi/whl_target_platforms/whl_target_platforms_tests.bzl
+index a976a0cf95..6bec26c10c 100644
+--- a/tests/pypi/whl_target_platforms/whl_target_platforms_tests.bzl
++++ b/tests/pypi/whl_target_platforms/whl_target_platforms_tests.bzl
+@@ -34,6 +34,9 @@ def _test_simple(env):
+         "musllinux_1_1_ppc64le": [
+             struct(os = "linux", cpu = "ppc64le", abi = None, target_platform = "linux_ppc64le", version = (1, 1)),
+         ],
++        "musllinux_1_2_riscv64": [
++            struct(os = "linux", cpu = "riscv64", abi = None, target_platform = "linux_riscv64", version = (1, 2)),
++        ],
+         "win_amd64": [
+             struct(os = "windows", cpu = "x86_64", abi = None, target_platform = "windows_x86_64", version = (0, 0)),
+         ],
+@@ -66,6 +69,9 @@ def _test_with_abi(env):
+         "musllinux_1_1_ppc64le": [
+             struct(os = "linux", cpu = "ppc64le", abi = "cp311", target_platform = "cp311_linux_ppc64le", version = (1, 1)),
+         ],
++        "musllinux_1_2_riscv64": [
++            struct(os = "linux", cpu = "riscv64", abi = "cp311", target_platform = "cp311_linux_riscv64", version = (1, 2)),
++        ],
+         "win_amd64": [
+             struct(os = "windows", cpu = "x86_64", abi = "cp311", target_platform = "cp311_windows_x86_64", version = (0, 0)),
+         ],
+@@ -103,6 +109,7 @@ def _can_parse_existing_tags(env):
+         "manylinux_11_12_i686": 1,
+         "manylinux_11_12_ppc64": 1,
+         "manylinux_11_12_ppc64le": 1,
++        "manylinux_11_12_riscv64": 1,
+         "manylinux_11_12_s390x": 1,
+         "manylinux_11_12_x86_64": 1,
+         "manylinux_1_2_aarch64": 1,
+@@ -111,6 +118,7 @@ def _can_parse_existing_tags(env):
+         "musllinux_11_12_armv7l": 1,
+         "musllinux_11_12_i686": 1,
+         "musllinux_11_12_ppc64le": 1,
++        "musllinux_11_12_riscv64": 1,
+         "musllinux_11_12_s390x": 1,
+         "musllinux_11_12_x86_64": 1,
+         "win32": 1,

xla/backends/gpu/runtime/BUILD

@@ -344,15 +344,12 @@ xla_test(
         "//xla:xla_data_proto_cc",
         "//xla/service:buffer_assignment",
         "//xla/stream_executor:gpu_solver_context",
-        "//xla/stream_executor:platform_manager",
         "//xla/stream_executor/platform:platform_object_registry",
         "//xla/tests:hlo_test_base",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/util/proto:proto_matchers",
-        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_googletest//:gtest_main",
-        "@tsl//tsl/platform:protobuf",
     ],
 )
 
@@ -725,7 +722,6 @@ cc_library(
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log",
-        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
@@ -2901,16 +2897,18 @@ cc_library(
     srcs = ["thunk_buffer_debug_pass.cc"],
     hdrs = ["thunk_buffer_debug_pass.h"],
     deps = [
+        ":buffer_debug_log_entry_metadata_store",
+        ":buffer_debug_log_structs",
         ":buffers_checksum_thunk",
         ":buffers_nan_count_thunk",
         ":custom_call_thunk",
         ":sequential_thunk",
         ":thunk",
-        ":thunk_buffer_id",
         ":thunk_pass_pipeline",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
         "//xla/ffi",
+        "//xla/ffi:attribute_map",
         "//xla/ffi/api:c_api",
         "//xla/hlo/ir:hlo",
         "//xla/runtime:buffer_use",
@@ -2922,6 +2920,7 @@ cc_library(
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/functional:bind_front",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -2940,7 +2939,6 @@ xla_cc_test(
         ":sequential_thunk",
         ":thunk",
         ":thunk_buffer_debug_pass",
-        ":thunk_buffer_id",
         ":thunk_id",
         ":thunk_pass_pipeline",
         "//xla:literal_util",
@@ -3006,8 +3004,10 @@ cc_library(
     srcs = ["buffers_checksum_thunk.cc"],
     hdrs = ["buffers_checksum_thunk.h"],
     deps = [
+        ":buffer_debug_log_entry_metadata_store",
+        ":buffer_debug_log_structs",
         ":thunk",
-        ":thunk_buffer_id",
+        ":thunk_id",
         "//xla/service:buffer_assignment",
         "//xla/stream_executor:device_memory",
         "//xla/stream_executor:launch_dim",
@@ -3035,10 +3035,10 @@ xla_test(
         "gpu",
     ],
     deps = [
+        ":buffer_debug_log_entry_metadata_store",
         ":buffer_debug_log_structs",
         ":buffers_checksum_thunk",
         ":thunk",
-        ":thunk_buffer_id",
         ":thunk_id",
         "//xla/service:buffer_assignment",
         "//xla/service:executable",
@@ -3052,6 +3052,7 @@ xla_test(
         "//xla/stream_executor/gpu:buffer_debug_log",
         "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/status",
         "@com_google_googletest//:gtest_main",
     ],
 )
@@ -3061,8 +3062,10 @@ cc_library(
     srcs = ["buffers_nan_count_thunk.cc"],
     hdrs = ["buffers_nan_count_thunk.h"],
     deps = [
+        ":buffer_debug_log_entry_metadata_store",
+        ":buffer_debug_log_structs",
         ":thunk",
-        ":thunk_buffer_id",
+        ":thunk_id",
         "//xla:types",
         "//xla/service:buffer_assignment",
         "//xla/stream_executor:device_memory",
@@ -3093,10 +3096,10 @@ xla_test(
         "gpu",
     ],
     deps = [
+        ":buffer_debug_log_entry_metadata_store",
         ":buffer_debug_log_structs",
         ":buffers_nan_count_thunk",
         ":thunk",
-        ":thunk_buffer_id",
         ":thunk_id",
         "//xla:types",
         "//xla/service:buffer_assignment",
@@ -3125,38 +3128,14 @@ xla_py_proto_library(
     deps = [":buffer_debug_log_proto"],
 )
 
-cc_library(
-    name = "thunk_buffer_id",
-    hdrs = ["thunk_buffer_id.h"],
-    compatible_with = get_compatible_with_portable(),
-    deps = [
-        ":thunk_id",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:str_format",
-    ],
-)
-
-xla_cc_test(
-    name = "thunk_buffer_id_test",
-    srcs = ["thunk_buffer_id_test.cc"],
-    deps = [
-        ":thunk_buffer_id",
-        ":thunk_id",
-        "//xla/tsl/platform:statusor",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_matchers",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
 cc_library(
     name = "buffer_debug_log_structs",
     hdrs = ["buffer_debug_log_structs.h"],
     compatible_with = get_compatible_with_portable(),
     deps = [
-        ":thunk_buffer_id",
+        ":buffer_debug_log_proto_cc",
         "//xla/tsl/lib/gtl:int_type",
+        "@com_google_absl//absl/strings:str_format",
     ],
 )
 
@@ -3182,7 +3161,6 @@ xla_cc_test(
     deps = [
         ":buffer_debug_log_entry_metadata_store",
         ":buffer_debug_log_structs",
-        ":thunk_buffer_id",
         ":thunk_id",
         "//xla/tsl/util/proto:proto_matchers",
         "@com_google_googletest//:gtest_main",

xla/backends/gpu/runtime/buffer_debug_log_entry_metadata_store.cc

@@ -63,9 +63,7 @@ BufferDebugLogProto BufferDebugLogEntryMetadataStore::EntriesToProto(
 
   BufferDebugLogProto proto;
   for (const BufferDebugLogEntry& entry : entries) {
-    // TODO: b/447080910 - simplify once entry_id is a BufferDebugLogEntryId.
-    std::optional<Metadata> metadata =
-        GetEntryMetadataLocked(BufferDebugLogEntryId{entry.entry_id.value()});
+    std::optional<Metadata> metadata = GetEntryMetadataLocked(entry.entry_id);
     if (!metadata.has_value()) {
       continue;
     }

xla/backends/gpu/runtime/buffer_debug_log_entry_metadata_store_test.cc

@@ -21,7 +21,6 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "xla/backends/gpu/runtime/buffer_debug_log_structs.h"
-#include "xla/backends/gpu/runtime/thunk_buffer_id.h"
 #include "xla/backends/gpu/runtime/thunk_id.h"
 #include "xla/tsl/util/proto/proto_matchers.h"
 
@@ -76,13 +75,11 @@ TEST(BufferDebugLogEntryMetadataStoreTest, EntriesToProto) {
   });
   std::vector<BufferDebugLogEntry> entries = {
       {
-          // TODO: b/447080910 - use BufferDebugLogEntryId directly.
-          /*entry_id=*/reinterpret_cast<const ThunkBufferId&>(entry_id1),
+          /*entry_id=*/entry_id1,
           /*checksum=*/12341234,
       },
       {
-          // TODO: b/447080910 - use BufferDebugLogEntryId directly.
-          /*entry_id=*/reinterpret_cast<const ThunkBufferId&>(entry_id2),
+          /*entry_id=*/entry_id2,
           /*checksum=*/56785678,
       },
   };

xla/backends/gpu/runtime/buffer_debug_log_structs.h

@@ -20,23 +20,22 @@ limitations under the License.
 #include <cstdint>
 #include <tuple>
 
-#include "xla/backends/gpu/runtime/thunk_buffer_id.h"
+#include "absl/strings/str_format.h"
+#include "xla/backends/gpu/runtime/buffer_debug_log.pb.h"
 #include "xla/tsl/lib/gtl/int_type.h"
 
 namespace xla::gpu {
 
-// TODO: b/447080910 - use this instead of ThunkBufferId.
 TSL_LIB_GTL_DEFINE_INT_TYPE(BufferDebugLogEntryId, uint32_t)
 
 struct BufferDebugLogEntry {
-  // An ID that uniquely identifies a thunk and its specific input or output
-  // buffer.
-  ThunkBufferId entry_id;
+  // An ID that uniquely identifies a log entry within a HLO module execution.
+  BufferDebugLogEntryId entry_id;
   uint32_t value;
 
   template <typename Sink>
   friend void AbslStringify(Sink& sink, const BufferDebugLogEntry& entry) {
-    absl::Format(&sink, "{entry_id: %v, value: %u}", entry.entry_id,
+    absl::Format(&sink, "{entry_id: %v, value: %u}", entry.entry_id.value(),
                  entry.value);
   }