Skip to content

Commit ff1b74c

Browse files
zacmustintensorflower-gardener
zacmustin
authored and
tensorflower-gardener
committed
Make a PjRtValueType proto and {To, From}Proto functions.
Ideally, these would belong to a `PjRtValueType` class (like we have for e.g. [Shape](https://github.com/openxla/xla/blob/f80e7871e58bef66b0d9a10f10c269935f9251f2/xla/shape.cc#L149)), but to avoid making a `PjRtValueType` class (and updating all the current users) we accept the `PjRtValueType` as an arg for now. PiperOrigin-RevId: 764831009
1 parent 056db11 commit ff1b74c

6 files changed

Lines changed: 232 additions & 0 deletions

File tree

‎third_party/xla/xla/pjrt/BUILD‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -453,13 +453,27 @@ xla_cc_test(
453453

454454
cc_library(
455455
name = "pjrt_common",
456+
srcs = ["pjrt_common.cc"],
456457
hdrs = ["pjrt_common.h"],
457458
visibility = internal_visibility([":friends"]),
458459
deps = [
460+
"//xla/pjrt/proto:pjrt_value_type_proto_cc",
459461
"//xla/tsl/lib/gtl:int_type",
460462
],
461463
)
462464

465+
xla_cc_test(
466+
name = "pjrt_common_test",
467+
srcs = ["pjrt_common_test.cc"],
468+
deps = [
469+
":pjrt_common",
470+
"//xla/pjrt/proto:pjrt_value_type_proto_cc",
471+
"@com_google_googletest//:gtest_main",
472+
"@local_tsl//tsl/platform:status_matchers",
473+
"@local_tsl//tsl/platform:test",
474+
],
475+
)
476+
463477
cc_library(
464478
name = "utils",
465479
srcs = ["utils.cc"],

‎third_party/xla/xla/pjrt/pjrt_common.cc‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
/* Copyright 2025 The OpenXLA Authors.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#include "xla/pjrt/pjrt_common.h"
17+
18+
#include <cstdint>
19+
#include <string>
20+
#include <type_traits>
21+
#include <variant>
22+
#include <vector>
23+
24+
#include "xla/pjrt/proto/pjrt_value_type.pb.h"
25+
26+
namespace xla {
27+
28+
xla::PjRtValueType PjRtValueTypeFromProto(
29+
const xla::PjRtValueTypeProto& value) {
30+
switch (value.value_case()) {
31+
case xla::PjRtValueTypeProto::kStringValue:
32+
return value.string_value();
33+
case xla::PjRtValueTypeProto::kBoolValue:
34+
return value.bool_value();
35+
case xla::PjRtValueTypeProto::kIntValue:
36+
return value.int_value();
37+
case xla::PjRtValueTypeProto::kIntVector:
38+
return std::vector<int64_t>(value.int_vector().values().begin(),
39+
value.int_vector().values().end());
40+
case xla::PjRtValueTypeProto::kFloatValue:
41+
return value.float_value();
42+
default:
43+
return std::string("");
44+
}
45+
}
46+
47+
xla::PjRtValueTypeProto PjRtValueTypeToProto(const xla::PjRtValueType& value) {
48+
xla::PjRtValueTypeProto value_proto;
49+
std::visit(
50+
[&](const auto& v) {
51+
using T = std::decay_t<decltype(v)>;
52+
if constexpr (std::is_same_v<T, float>) {
53+
value_proto.set_float_value(v);
54+
} else if constexpr (std::is_same_v<T, int64_t>) {
55+
value_proto.set_int_value(v);
56+
} else if constexpr (std::is_same_v<T, std::string>) {
57+
value_proto.set_string_value(v);
58+
} else if constexpr (std::is_same_v<T, bool>) {
59+
value_proto.set_bool_value(v);
60+
} else if constexpr (std::is_same_v<T, std::vector<int64_t>>) {
61+
value_proto.mutable_int_vector()->mutable_values()->Reserve(v.size());
62+
value_proto.mutable_int_vector()->mutable_values()->Add(v.begin(),
63+
v.end());
64+
} else {
65+
static_assert(false, "Unhandled type in PjRtValueType variant");
66+
}
67+
},
68+
value);
69+
return value_proto;
70+
}
71+
72+
} // namespace xla

‎third_party/xla/xla/pjrt/pjrt_common.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ limitations under the License.
2121
#include <variant>
2222
#include <vector>
2323

24+
#include "xla/pjrt/proto/pjrt_value_type.pb.h"
2425
#include "xla/tsl/lib/gtl/int_type.h"
2526

2627
namespace xla {
@@ -31,6 +32,10 @@ namespace xla {
3132
using PjRtValueType =
3233
std::variant<std::string, bool, int64_t, std::vector<int64_t>, float>;
3334

35+
xla::PjRtValueTypeProto PjRtValueTypeToProto(const PjRtValueType& value);
36+
37+
PjRtValueType PjRtValueTypeFromProto(const xla::PjRtValueTypeProto& value);
38+
3439
// The strong-typed integer classes to better disambiguate different IDs for
3540
// PJRT devices.
3641
TSL_LIB_GTL_DEFINE_INT_TYPE(PjRtGlobalDeviceId, int32_t);

‎third_party/xla/xla/pjrt/pjrt_common_test.cc‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/* Copyright 2025 The OpenXLA Authors.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#include "xla/pjrt/pjrt_common.h"
17+
18+
#include <cstdint>
19+
#include <string>
20+
#include <variant>
21+
#include <vector>
22+
23+
#include <gmock/gmock.h>
24+
#include <gtest/gtest.h>
25+
#include "xla/pjrt/proto/pjrt_value_type.pb.h"
26+
27+
using testing::ElementsAre;
28+
29+
namespace {
30+
31+
TEST(PjRtCommonTest, PjRtValueTypeFromProtoDefault) {
32+
xla::PjRtValueTypeProto proto;
33+
xla::PjRtValueType value = xla::PjRtValueTypeFromProto(proto);
34+
ASSERT_TRUE(std::holds_alternative<std::string>(value));
35+
EXPECT_EQ(std::get<std::string>(value), "");
36+
}
37+
38+
TEST(PjRtCommonTest, PjRtValueTypeToProtoFloat) {
39+
xla::PjRtValueType value = 1.0f;
40+
xla::PjRtValueTypeProto proto = xla::PjRtValueTypeToProto(value);
41+
ASSERT_TRUE(proto.has_float_value());
42+
EXPECT_EQ(proto.float_value(), 1.0f);
43+
}
44+
45+
TEST(PjRtCommonTest, PjRtValueTypeToProtoInt) {
46+
xla::PjRtValueType value = 1L;
47+
xla::PjRtValueTypeProto proto = xla::PjRtValueTypeToProto(value);
48+
ASSERT_TRUE(proto.has_int_value());
49+
EXPECT_EQ(proto.int_value(), 1L);
50+
}
51+
52+
TEST(PjRtCommonTest, PjRtValueTypeToProtoString) {
53+
xla::PjRtValueType value = std::string("test");
54+
xla::PjRtValueTypeProto proto = xla::PjRtValueTypeToProto(value);
55+
ASSERT_TRUE(proto.has_string_value());
56+
EXPECT_EQ(proto.string_value(), "test");
57+
}
58+
59+
TEST(PjRtCommonTest, PjRtValueTypeToProtoBool) {
60+
xla::PjRtValueType value = true;
61+
xla::PjRtValueTypeProto proto = xla::PjRtValueTypeToProto(value);
62+
ASSERT_TRUE(proto.has_bool_value());
63+
EXPECT_EQ(proto.bool_value(), true);
64+
}
65+
66+
TEST(PjRtCommonTest, PjRtValueTypeToProtoIntVector) {
67+
xla::PjRtValueType value = std::vector<int64_t>{1, 2, 3};
68+
xla::PjRtValueTypeProto proto = xla::PjRtValueTypeToProto(value);
69+
ASSERT_TRUE(proto.has_int_vector());
70+
EXPECT_THAT(proto.int_vector().values(), ElementsAre(1, 2, 3));
71+
}
72+
73+
TEST(PjRtCommonTest, PjRtValueTypeFromProtoFloat) {
74+
xla::PjRtValueTypeProto proto;
75+
proto.set_float_value(1.0f);
76+
xla::PjRtValueType value = xla::PjRtValueTypeFromProto(proto);
77+
ASSERT_TRUE(std::holds_alternative<float>(value));
78+
EXPECT_EQ(std::get<float>(value), 1.0f);
79+
}
80+
81+
TEST(PjRtCommonTest, PjRtValueTypeFromProtoInt) {
82+
xla::PjRtValueTypeProto proto;
83+
proto.set_int_value(1L);
84+
xla::PjRtValueType value = xla::PjRtValueTypeFromProto(proto);
85+
ASSERT_TRUE(std::holds_alternative<int64_t>(value));
86+
EXPECT_EQ(std::get<int64_t>(value), 1L);
87+
}
88+
89+
TEST(PjRtCommonTest, PjRtValueTypeFromProtoString) {
90+
xla::PjRtValueTypeProto proto;
91+
proto.set_string_value("test");
92+
xla::PjRtValueType value = xla::PjRtValueTypeFromProto(proto);
93+
ASSERT_TRUE(std::holds_alternative<std::string>(value));
94+
EXPECT_EQ(std::get<std::string>(value), "test");
95+
}
96+
97+
TEST(PjRtCommonTest, PjRtValueTypeFromProtoBool) {
98+
xla::PjRtValueTypeProto proto;
99+
proto.set_bool_value(true);
100+
xla::PjRtValueType value = xla::PjRtValueTypeFromProto(proto);
101+
ASSERT_TRUE(std::holds_alternative<bool>(value));
102+
EXPECT_EQ(std::get<bool>(value), true);
103+
}
104+
105+
TEST(PjRtCommonTest, PjRtValueTypeFromProtoIntVector) {
106+
xla::PjRtValueTypeProto proto;
107+
proto.mutable_int_vector()->add_values(1);
108+
proto.mutable_int_vector()->add_values(2);
109+
proto.mutable_int_vector()->add_values(3);
110+
xla::PjRtValueType value = xla::PjRtValueTypeFromProto(proto);
111+
ASSERT_TRUE(std::holds_alternative<std::vector<int64_t>>(value));
112+
EXPECT_THAT(std::get<std::vector<int64_t>>(value),
113+
testing::ElementsAre(1, 2, 3));
114+
}
115+
116+
} // namespace

‎third_party/xla/xla/pjrt/proto/BUILD‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,3 +39,9 @@ tf_proto_library(
3939
srcs = ["executable_metadata.proto"],
4040
visibility = ["//visibility:public"],
4141
)
42+
43+
tf_proto_library(
44+
name = "pjrt_value_type_proto",
45+
srcs = ["pjrt_value_type.proto"],
46+
visibility = ["//visibility:public"],
47+
)

‎third_party/xla/xla/pjrt/proto/pjrt_value_type.proto‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
syntax = "proto3";
2+
3+
package xla;
4+
5+
// Mirror of xla::PjRtValueType.
6+
7+
message PjRtValueTypeProto {
8+
message IntVector {
9+
repeated int64 values = 1;
10+
}
11+
12+
oneof value {
13+
string string_value = 1;
14+
bool bool_value = 2;
15+
int64 int_value = 3;
16+
IntVector int_vector = 4;
17+
float float_value = 5;
18+
}
19+
}

0 commit comments

Comments
 (0)