PR #32812: [XLA:CPU] Add support for riscv64

Imported from GitHub PR openxla/xla#32812

Co-author: @kxxt

📝 Summary of Changes:

This pull request adds support for RISC-V 64 architecture across the build system, code generation, and Python packaging infrastructure.

🎯 Justification:

The changes ensure that riscv64 is recognized as a valid target in Bazel build configurations, LLVM toolchain selection, Python manylinux compliance checks, and related tests and patches. This allows the project to build and test components for riscv64 alongside other supported architectures.

🚀 Kind of Contribution: ✨ New Feature

Copybara import of the project:

--
0d02393a6335fb43d67678d0cd15d671e77dc089 by gns <[email protected]>:

[XLA:CPU] Add support for riscv64

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

--
5d95fb479e45524299ff4193b99bb4db0d74483b by gns <[email protected]>:

Refresh `rules_python` riscv64 patch

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

Merging this change closes #32812

PiperOrigin-RevId: 828379922

Author: infiWang

tensorflow/core/platform/build_config_root.bzl

@@ -7,6 +7,7 @@ load(
     _if_llvm_arm_available = "if_llvm_arm_available",
     _if_llvm_hexagon_available = "if_llvm_hexagon_available",
     _if_llvm_powerpc_available = "if_llvm_powerpc_available",
+    _if_llvm_riscv_available = "if_llvm_riscv_available",
     _if_llvm_system_z_available = "if_llvm_system_z_available",
     _if_llvm_x86_available = "if_llvm_x86_available",
     _if_pywrap = "if_pywrap",
@@ -32,6 +33,7 @@ if_llvm_aarch64_available = _if_llvm_aarch64_available
 if_llvm_arm_available = _if_llvm_arm_available
 if_llvm_hexagon_available = _if_llvm_hexagon_available
 if_llvm_powerpc_available = _if_llvm_powerpc_available
+if_llvm_riscv_available = _if_llvm_riscv_available
 if_llvm_system_z_available = _if_llvm_system_z_available
 if_llvm_x86_available = _if_llvm_x86_available
 if_dynamic_kernels = _if_dynamic_kernels

third_party/xla/third_party/hwloc/hwloc.BUILD

@@ -272,6 +272,10 @@ cc_library(
             "hwloc/topology-linux.c",
             "include/hwloc/linux.h",
         ],
+        "@local_xla//xla/tsl:linux_riscv64": [
+            "hwloc/topology-linux.c",
+            "include/hwloc/linux.h",
+        ],
         "@local_xla//xla/tsl:linux_s390x": [
             "hwloc/topology-linux.c",
             "include/hwloc/linux.h",

third_party/xla/third_party/mkl_dnn/mkldnn_v1.BUILD

@@ -220,6 +220,7 @@ cc_library(
         "@local_xla//xla/tsl:linux_aarch64": ["-lrt"],
         "@local_xla//xla/tsl:linux_x86_64": ["-lrt"],
         "@local_xla//xla/tsl:linux_ppc64le": ["-lrt"],
+        "@local_xla//xla/tsl:linux_riscv64": ["-lrt"],
         "//conditions:default": [],
     }),
     textual_hdrs = _TEXTUAL_HDRS_LIST,

third_party/xla/xla/backends/cpu/codegen/BUILD

@@ -5,6 +5,7 @@ load(
     "//xla/tsl/platform:build_config_root.bzl",
     "if_llvm_aarch64_available",
     "if_llvm_powerpc_available",
+    "if_llvm_riscv_available",
     "if_llvm_system_z_available",
     "if_llvm_x86_available",
 )
@@ -164,6 +165,8 @@ xla_cc_test(
         "@llvm-project//llvm:AArch64CodeGen",  # fixdeps: keep
     ]) + if_llvm_powerpc_available([
         "@llvm-project//llvm:PowerPCCodeGen",  # fixdeps: keep
+    ]) + if_llvm_riscv_available([
+        "@llvm-project//llvm:RISCVCodeGen",  # fixdeps: keep
     ]) + if_llvm_system_z_available([
         "@llvm-project//llvm:SystemZCodeGen",  # fixdeps: keep
     ]) + if_llvm_x86_available([
@@ -286,6 +289,8 @@ cc_library(
         "@llvm-project//llvm:AArch64CodeGen",  # fixdeps: keep
     ]) + if_llvm_powerpc_available([
         "@llvm-project//llvm:PowerPCCodeGen",  # fixdeps: keep
+    ]) + if_llvm_riscv_available([
+        "@llvm-project//llvm:RISCVCodeGen",  # fixdeps: keep
     ]) + if_llvm_system_z_available([
         "@llvm-project//llvm:SystemZCodeGen",  # fixdeps: keep
     ]) + if_llvm_x86_available([

third_party/xla/xla/backends/cpu/codegen/tools/BUILD

@@ -4,6 +4,7 @@ load(
     "//xla/tsl/platform:build_config_root.bzl",
     "if_llvm_aarch64_available",
     "if_llvm_powerpc_available",
+    "if_llvm_riscv_available",
     "if_llvm_system_z_available",
     "if_llvm_x86_available",
 )
@@ -38,6 +39,8 @@ cc_library(
         "@llvm-project//llvm:AArch64CodeGen",  # fixdeps: keep
     ]) + if_llvm_powerpc_available([
         "@llvm-project//llvm:PowerPCCodeGen",  # fixdeps: keep
+    ]) + if_llvm_riscv_available([
+        "@llvm-project//llvm:RISCVCodeGen",  # fixdeps: keep
     ]) + if_llvm_system_z_available([
         "@llvm-project//llvm:SystemZCodeGen",  # fixdeps: keep
     ]) + if_llvm_x86_available([

third_party/xla/xla/codegen/intrinsic/BUILD

@@ -3,6 +3,7 @@ load(
     "//xla/tsl/platform:build_config_root.bzl",
     "if_llvm_aarch64_available",
     "if_llvm_powerpc_available",
+    "if_llvm_riscv_available",
     "if_llvm_system_z_available",
     "if_llvm_x86_available",
 )
@@ -149,6 +150,9 @@ cc_library(
     ]) + if_llvm_powerpc_available([
         "@llvm-project//llvm:PowerPCAsmParser",  # fixdeps: keep
         "@llvm-project//llvm:PowerPCCodeGen",  # fixdeps: keep
+    ]) + if_llvm_riscv_available([
+        "@llvm-project//llvm:RISCVAsmParser",  # fixdeps: keep
+        "@llvm-project//llvm:RISCVCodeGen",  # fixdeps: keep
     ]) + if_llvm_system_z_available([
         "@llvm-project//llvm:SystemZAsmParser",  # fixdeps: keep
         "@llvm-project//llvm:SystemZCodeGen",  # fixdeps: keep

third_party/xla/xla/service/cpu/BUILD

@@ -19,6 +19,7 @@ load(
     "//xla/tsl/platform:build_config_root.bzl",
     "if_llvm_aarch64_available",
     "if_llvm_powerpc_available",
+    "if_llvm_riscv_available",
     "if_llvm_system_z_available",
     "if_llvm_x86_available",
 )
@@ -392,6 +393,8 @@ cc_library(
         "@llvm-project//llvm:AArch64CodeGen",  # fixdeps: keep
     ]) + if_llvm_powerpc_available([
         "@llvm-project//llvm:PowerPCCodeGen",  # fixdeps: keep
+    ]) + if_llvm_riscv_available([
+        "@llvm-project//llvm:RISCVCodeGen",  # fixdeps: keep
     ]) + if_llvm_system_z_available([
         "@llvm-project//llvm:SystemZCodeGen",  # fixdeps: keep
     ]) + if_llvm_x86_available([
@@ -2120,6 +2123,8 @@ cc_library(
         "@llvm-project//llvm:AArch64CodeGen",  # fixdeps: keep
     ]) + if_llvm_powerpc_available([
         "@llvm-project//llvm:PowerPCCodeGen",  # fixdeps: keep
+    ]) + if_llvm_riscv_available([
+        "@llvm-project//llvm:RISCVCodeGen",  # fixdeps: keep
     ]) + if_llvm_system_z_available([
         "@llvm-project//llvm:SystemZCodeGen",  # fixdeps: keep
     ]) + if_llvm_x86_available([

third_party/xla/xla/service/cpu/test_target_triple_helper.h

@@ -23,6 +23,9 @@ static const char kTargetTripleForHost[] = "aarch64-unknown-linux-gnu";
        defined(__ppc__) && (__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__))
 static const char kTargetCpuForHost[] = "ppc";
 static const char kTargetTripleForHost[] = "ppc64le-ibm-linux-gnu";
+#elif defined(__riscv) && (__riscv_xlen == 64)
+static const char kTargetCpuForHost[] = "";
+static const char kTargetTripleForHost[] = "riscv64-unknown-linux-gnu";
 #elif defined(__s390x__)
 static const char kTargetCpuForHost[] = "s390x";
 static const char kTargetTripleForHost[] = "systemz-none-linux-gnu";

third_party/xla/xla/tsl/BUILD

@@ -316,6 +316,16 @@ config_setting(
     visibility = ["//visibility:public"],
 )
 
+config_setting(
+    name = "linux_riscv64",
+    constraint_values =
+        [
+            "@platforms//cpu:riscv64",
+            "@platforms//os:linux",
+        ],
+    visibility = ["//visibility:public"],
+)
+
 config_setting(
     name = "linux_s390x",
     constraint_values =
@@ -382,6 +392,15 @@ selects.config_setting_group(
     visibility = ["//visibility:public"],
 )
 
+selects.config_setting_group(
+    name = "riscv64_or_cross",
+    match_any = [
+        ":linux_riscv64",
+        ":with_cross_compiler_support",
+    ],
+    visibility = ["//visibility:public"],
+)
+
 selects.config_setting_group(
     name = "s390x_or_cross",
     match_any = [
@@ -453,6 +472,7 @@ selects.config_setting_group(
         ":linux_aarch64",
         ":linux_armhf",
         ":linux_ppc64le",
+        ":linux_riscv64",
         ":linux_s390x",
         ":linux_x86_64",
     ],

third_party/xla/xla/tsl/framework/contraction/BUILD

@@ -113,6 +113,7 @@ cc_library(
         "//xla/tsl:fuchsia_x86_64": [],
         "//xla/tsl:ios": [],
         "//xla/tsl:linux_ppc64le": [],
+        "//xla/tsl:linux_riscv64": [],
         "//xla/tsl:linux_s390x": [],
         "//xla/tsl:macos_arm64": [],
         "//conditions:default": [
@@ -132,6 +133,7 @@ cc_library(
         "//xla/tsl:fuchsia_x86_64": [],
         "//xla/tsl:ios": [],
         "//xla/tsl:linux_ppc64le": [],
+        "//xla/tsl:linux_riscv64": [],
         "//xla/tsl:linux_s390x": [],
         "//xla/tsl:macos_arm64": [],
         "//conditions:default": ["//xla/tsl/mkl:onednn"],