Skip to content

Commit e8ad0eb

Browse files
infiWangGoogle-ML-Automation
infiWang
authored and
Google-ML-Automation
committed
PR #32812: [XLA:CPU] Add support for riscv64
Imported from GitHub PR #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: -- 0d02393 by gns <[email protected]>: [XLA:CPU] Add support for riscv64 Co-authored-by: Levi Zim <[email protected]> -- 5d95fb4 by gns <[email protected]>: Refresh `rules_python` riscv64 patch Co-authored-by: Levi Zim <[email protected]> Merging this change closes #32812 FUTURE_COPYBARA_INTEGRATE_REVIEW=#32812 from infiWang:riscv64 5d95fb4 PiperOrigin-RevId: 826410020
1 parent bf76a75 commit e8ad0eb

File tree

14 files changed

+123
-0
lines changed

14 files changed

+123
-0
lines changed

‎third_party/hwloc/hwloc.BUILD‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,10 @@ cc_library(
272272
"hwloc/topology-linux.c",
273273
"include/hwloc/linux.h",
274274
],
275+
"@xla//xla/tsl:linux_riscv64": [
276+
"hwloc/topology-linux.c",
277+
"include/hwloc/linux.h",
278+
],
275279
"@xla//xla/tsl:linux_s390x": [
276280
"hwloc/topology-linux.c",
277281
"include/hwloc/linux.h",

‎third_party/mkl_dnn/mkldnn_v1.BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,7 @@ cc_library(
220220
"@xla//xla/tsl:linux_aarch64": ["-lrt"],
221221
"@xla//xla/tsl:linux_x86_64": ["-lrt"],
222222
"@xla//xla/tsl:linux_ppc64le": ["-lrt"],
223+
"@xla//xla/tsl:linux_riscv64": ["-lrt"],
223224
"//conditions:default": [],
224225
}),
225226
textual_hdrs = _TEXTUAL_HDRS_LIST,

‎xla/backends/cpu/codegen/BUILD‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ load(
55
"//xla/tsl/platform:build_config_root.bzl",
66
"if_llvm_aarch64_available",
77
"if_llvm_powerpc_available",
8+
"if_llvm_riscv_available",
89
"if_llvm_system_z_available",
910
"if_llvm_x86_available",
1011
)
@@ -129,6 +130,8 @@ xla_cc_test(
129130
"@llvm-project//llvm:AArch64CodeGen", # fixdeps: keep
130131
]) + if_llvm_powerpc_available([
131132
"@llvm-project//llvm:PowerPCCodeGen", # fixdeps: keep
133+
]) + if_llvm_riscv_available([
134+
"@llvm-project//llvm:RISCVCodeGen", # fixdeps: keep
132135
]) + if_llvm_system_z_available([
133136
"@llvm-project//llvm:SystemZCodeGen", # fixdeps: keep
134137
]) + if_llvm_x86_available([
@@ -251,6 +254,8 @@ cc_library(
251254
"@llvm-project//llvm:AArch64CodeGen", # fixdeps: keep
252255
]) + if_llvm_powerpc_available([
253256
"@llvm-project//llvm:PowerPCCodeGen", # fixdeps: keep
257+
]) + if_llvm_riscv_available([
258+
"@llvm-project//llvm:RISCVCodeGen", # fixdeps: keep
254259
]) + if_llvm_system_z_available([
255260
"@llvm-project//llvm:SystemZCodeGen", # fixdeps: keep
256261
]) + if_llvm_x86_available([

‎xla/backends/cpu/codegen/tools/BUILD‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ load(
44
"//xla/tsl/platform:build_config_root.bzl",
55
"if_llvm_aarch64_available",
66
"if_llvm_powerpc_available",
7+
"if_llvm_riscv_available",
78
"if_llvm_system_z_available",
89
"if_llvm_x86_available",
910
)
@@ -38,6 +39,8 @@ cc_library(
3839
"@llvm-project//llvm:AArch64CodeGen", # fixdeps: keep
3940
]) + if_llvm_powerpc_available([
4041
"@llvm-project//llvm:PowerPCCodeGen", # fixdeps: keep
42+
]) + if_llvm_riscv_available([
43+
"@llvm-project//llvm:RISCVCodeGen", # fixdeps: keep
4144
]) + if_llvm_system_z_available([
4245
"@llvm-project//llvm:SystemZCodeGen", # fixdeps: keep
4346
]) + if_llvm_x86_available([

‎xla/codegen/intrinsic/BUILD‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ load(
33
"//xla/tsl/platform:build_config_root.bzl",
44
"if_llvm_aarch64_available",
55
"if_llvm_powerpc_available",
6+
"if_llvm_riscv_available",
67
"if_llvm_system_z_available",
78
"if_llvm_x86_available",
89
)
@@ -149,6 +150,9 @@ cc_library(
149150
]) + if_llvm_powerpc_available([
150151
"@llvm-project//llvm:PowerPCAsmParser", # fixdeps: keep
151152
"@llvm-project//llvm:PowerPCCodeGen", # fixdeps: keep
153+
]) + if_llvm_riscv_available([
154+
"@llvm-project//llvm:RISCVAsmParser", # fixdeps: keep
155+
"@llvm-project//llvm:RISCVCodeGen", # fixdeps: keep
152156
]) + if_llvm_system_z_available([
153157
"@llvm-project//llvm:SystemZAsmParser", # fixdeps: keep
154158
"@llvm-project//llvm:SystemZCodeGen", # fixdeps: keep

‎xla/ffi/api/api.h‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1356,10 +1356,42 @@ class DictionaryBase {
13561356
public:
13571357
explicit DictionaryBase(const XLA_FFI_Attrs* attrs) : attrs_(attrs) {}
13581358

1359+
// Iterator for iterating over dictionary attribute names.
1360+
class Iterator {
1361+
public:
1362+
using iterator_category = std::forward_iterator_tag;
1363+
using difference_type = ptrdiff_t;
1364+
using value_type = std::string_view;
1365+
1366+
bool operator==(const Iterator& it) const { return idx_ == it.idx_; }
1367+
bool operator!=(const Iterator& it) const { return idx_ != it.idx_; }
1368+
1369+
std::string_view operator*() const {
1370+
return std::string_view{attrs_->names[idx_]->ptr,
1371+
attrs_->names[idx_]->len};
1372+
}
1373+
1374+
Iterator& operator++() {
1375+
++idx_;
1376+
return *this;
1377+
}
1378+
1379+
private:
1380+
friend class DictionaryBase;
1381+
Iterator(const XLA_FFI_Attrs* attrs, size_t idx)
1382+
: attrs_(attrs), idx_(idx) {}
1383+
1384+
const XLA_FFI_Attrs* attrs_;
1385+
size_t idx_ = 0;
1386+
};
1387+
13591388
size_t size() const { return attrs_->size; }
13601389

13611390
bool contains(std::string_view name) const { return Find(name).has_value(); }
13621391

1392+
Iterator begin() const { return Iterator(attrs_, 0); }
1393+
Iterator end() const { return Iterator(attrs_, size()); }
1394+
13631395
template <typename T>
13641396
bool contains(std::string_view name) const {
13651397
std::optional<size_t> idx = Find(name);
@@ -1372,10 +1404,25 @@ class DictionaryBase {
13721404
return AttrDecoding<T>::Isa(attr_type, attr);
13731405
}
13741406

1407+
template <typename T>
1408+
bool isa(const Iterator& it) const {
1409+
XLA_FFI_AttrType attr_type = attrs_->types[it.idx_];
1410+
void* attr = attrs_->attrs[it.idx_];
1411+
return AttrDecoding<T>::Isa(attr_type, attr);
1412+
}
1413+
13751414
protected:
13761415
template <typename T, typename... Ts>
13771416
friend struct DecodeDictionaryAttr;
13781417

1418+
template <typename T>
1419+
std::optional<size_t> get(const Iterator& it,
1420+
DiagnosticEngine& diagnostic) const {
1421+
XLA_FFI_AttrType attr_type = attrs_->types[it.idx_];
1422+
void* attr = attrs_->attrs[it.idx_];
1423+
return AttrDecoding<T>::Decode(attr_type, attr, diagnostic);
1424+
}
1425+
13791426
template <typename T>
13801427
std::optional<T> get(std::string_view name,
13811428
DiagnosticEngine& diagnostic) const {

‎xla/ffi/api/ffi.h‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1080,6 +1080,16 @@ class Dictionary : public internal::DictionaryBase {
10801080
public:
10811081
using internal::DictionaryBase::DictionaryBase;
10821082

1083+
template <typename T>
1084+
ErrorOr<T> get(const Iterator& it) const {
1085+
DiagnosticEngine diagnostic;
1086+
auto value = internal::DictionaryBase::get<T>(it, diagnostic);
1087+
if (XLA_FFI_PREDICT_FALSE(!value.has_value())) {
1088+
return Unexpected(Error::Internal(diagnostic.Result()));
1089+
}
1090+
return *value;
1091+
}
1092+
10831093
template <typename T>
10841094
ErrorOr<T> get(std::string_view name) const {
10851095
DiagnosticEngine diagnostic;

‎xla/ffi/api/ffi_test.cc‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -912,6 +912,17 @@ TEST(FfiTest, AutoBindingDictionary) {
912912
EXPECT_EQ(*attrs.get<int32_t>("i32"), 42);
913913
EXPECT_EQ(*attrs.get<float>("f32"), 42.0f);
914914

915+
auto it = attrs.begin();
916+
EXPECT_EQ(*it, "f32");
917+
EXPECT_TRUE(attrs.isa<float>(it));
918+
EXPECT_EQ(*attrs.get<float>(it), 42.0f);
919+
920+
EXPECT_EQ(*++it, "i32");
921+
EXPECT_TRUE(attrs.isa<int32_t>(it));
922+
EXPECT_EQ(*attrs.get<int32_t>(it), 42);
923+
924+
EXPECT_EQ(++it, attrs.end());
925+
915926
return Error::Success();
916927
});
917928

‎xla/service/cpu/BUILD‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ load(
1919
"//xla/tsl/platform:build_config_root.bzl",
2020
"if_llvm_aarch64_available",
2121
"if_llvm_powerpc_available",
22+
"if_llvm_riscv_available",
2223
"if_llvm_system_z_available",
2324
"if_llvm_x86_available",
2425
)
@@ -396,6 +397,8 @@ cc_library(
396397
"@llvm-project//llvm:AArch64CodeGen", # fixdeps: keep
397398
]) + if_llvm_powerpc_available([
398399
"@llvm-project//llvm:PowerPCCodeGen", # fixdeps: keep
400+
]) + if_llvm_riscv_available([
401+
"@llvm-project//llvm:RISCVCodeGen", # fixdeps: keep
399402
]) + if_llvm_system_z_available([
400403
"@llvm-project//llvm:SystemZCodeGen", # fixdeps: keep
401404
]) + if_llvm_x86_available([
@@ -2168,6 +2171,8 @@ cc_library(
21682171
"@llvm-project//llvm:AArch64CodeGen", # fixdeps: keep
21692172
]) + if_llvm_powerpc_available([
21702173
"@llvm-project//llvm:PowerPCCodeGen", # fixdeps: keep
2174+
]) + if_llvm_riscv_available([
2175+
"@llvm-project//llvm:RISCVCodeGen", # fixdeps: keep
21712176
]) + if_llvm_system_z_available([
21722177
"@llvm-project//llvm:SystemZCodeGen", # fixdeps: keep
21732178
]) + if_llvm_x86_available([

‎xla/service/cpu/test_target_triple_helper.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@ static const char kTargetTripleForHost[] = "aarch64-unknown-linux-gnu";
2323
defined(__ppc__) && (__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__))
2424
static const char kTargetCpuForHost[] = "ppc";
2525
static const char kTargetTripleForHost[] = "ppc64le-ibm-linux-gnu";
26+
#elif defined(__riscv) && (__riscv_xlen == 64)
27+
static const char kTargetCpuForHost[] = "";
28+
static const char kTargetTripleForHost[] = "riscv64-unknown-linux-gnu";
2629
#elif defined(__s390x__)
2730
static const char kTargetCpuForHost[] = "s390x";
2831
static const char kTargetTripleForHost[] = "systemz-none-linux-gnu";

0 commit comments

Comments
 (0)