Skip to content

Commit 0ee38ea

Browse files
beckerhetensorflower-gardener
beckerhe
authored and
tensorflower-gardener
committed
Add proto serialization for xla::gpu::LaunchDimensions
This is adding `ToProto` and `FromProto` methods to `LaunchDimensions` as well as adding missing tests for the non-obvious constructor overloads. PiperOrigin-RevId: 765132837
1 parent 6a34ed4 commit 0ee38ea

5 files changed

Lines changed: 160 additions & 3 deletions

File tree

‎third_party/xla/xla/service/gpu/BUILD‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@ cc_library(
124124
],
125125
compatible_with = get_compatible_with_portable(),
126126
deps = [
127+
":launch_dimensions_proto_cc",
127128
"//xla:shape_util",
128129
"//xla:util",
129130
"//xla/runtime:work_cluster",
@@ -133,11 +134,33 @@ cc_library(
133134
"//xla/service:platform_util",
134135
"//xla/stream_executor:device_description",
135136
"//xla/stream_executor:launch_dim",
137+
"//xla/tsl/platform:statusor",
136138
"@com_google_absl//absl/log:check",
139+
"@com_google_absl//absl/status:statusor",
137140
"@com_google_absl//absl/strings",
138141
],
139142
)
140143

144+
tf_proto_library(
145+
name = "launch_dimensions_proto",
146+
srcs = ["launch_dimensions.proto"],
147+
protodeps = ["//xla/stream_executor:launch_dim_proto"],
148+
)
149+
150+
xla_cc_test(
151+
name = "launch_dimensions_test",
152+
srcs = ["launch_dimensions_test.cc"],
153+
deps = [
154+
":launch_dimensions",
155+
"//xla/stream_executor:launch_dim",
156+
"//xla/tsl/platform:status_matchers",
157+
"//xla/tsl/platform:statusor",
158+
"//xla/tsl/util/proto:proto_matchers",
159+
"@com_google_googletest//:gtest_main",
160+
"@local_tsl//tsl/platform:protobuf",
161+
],
162+
)
163+
141164
xla_test(
142165
name = "custom_call_test",
143166
srcs = ["custom_call_test.cc"],

‎third_party/xla/xla/service/gpu/launch_dimensions.cc‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ limitations under the License.
2020
#include <limits>
2121

2222
#include "absl/log/check.h"
23+
#include "absl/status/statusor.h"
2324
#include "xla/runtime/work_cluster.h"
2425
#include "xla/runtime/work_dimensions.h"
2526
#include "xla/runtime/work_group.h"
@@ -29,6 +30,7 @@ limitations under the License.
2930
#include "xla/shape_util.h"
3031
#include "xla/stream_executor/device_description.h"
3132
#include "xla/stream_executor/launch_dim.h"
33+
#include "xla/tsl/platform/statusor.h"
3234
#include "xla/util.h"
3335

3436
namespace xla {
@@ -84,5 +86,22 @@ LaunchDimensions CalculateLaunchDimensions(
8486
}
8587
}
8688

89+
LaunchDimensionsProto LaunchDimensions::ToProto() const {
90+
LaunchDimensionsProto proto;
91+
*proto.mutable_block_counts() = block_counts_.ToProto();
92+
*proto.mutable_thread_counts_per_block() = thread_counts_per_block_.ToProto();
93+
return proto;
94+
}
95+
96+
absl::StatusOr<LaunchDimensions> LaunchDimensions::FromProto(
97+
const LaunchDimensionsProto& proto) {
98+
TF_ASSIGN_OR_RETURN(
99+
stream_executor::BlockDim block_counts,
100+
stream_executor::BlockDim::FromProto(proto.block_counts()));
101+
TF_ASSIGN_OR_RETURN(
102+
stream_executor::ThreadDim thread_counts_per_block,
103+
stream_executor::ThreadDim::FromProto(proto.thread_counts_per_block()));
104+
return LaunchDimensions{block_counts, thread_counts_per_block};
105+
}
87106
} // namespace gpu
88107
} // namespace xla

‎third_party/xla/xla/service/gpu/launch_dimensions.h‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,12 @@ limitations under the License.
1818

1919
#include <cstdint>
2020
#include <string>
21+
#include <tuple>
2122

23+
#include "absl/status/statusor.h"
2224
#include "absl/strings/str_cat.h"
2325
#include "xla/runtime/work_dimensions.h"
26+
#include "xla/service/gpu/launch_dimensions.pb.h"
2427
#include "xla/shape.h"
2528
#include "xla/stream_executor/device_description.h"
2629
#include "xla/stream_executor/launch_dim.h"
@@ -34,9 +37,7 @@ class LaunchDimensions {
3437
public:
3538
// The default constructor creates a launch dimension that indicate
3639
// single-threaded execution.
37-
constexpr LaunchDimensions()
38-
: block_counts_(se::BlockDim()),
39-
thread_counts_per_block_(se::ThreadDim()) {}
40+
constexpr LaunchDimensions() = default;
4041

4142
constexpr LaunchDimensions(uint64_t block_x_count,
4243
uint64_t thread_x_count_per_block)
@@ -79,6 +80,21 @@ class LaunchDimensions {
7980

8081
WorkDimensions AsWorkDimensions() const;
8182

83+
LaunchDimensionsProto ToProto() const;
84+
static absl::StatusOr<LaunchDimensions> FromProto(
85+
const LaunchDimensionsProto& proto);
86+
87+
friend bool operator==(const LaunchDimensions& lhs,
88+
const LaunchDimensions& rhs) {
89+
return std::tie(lhs.block_counts_, lhs.thread_counts_per_block_) ==
90+
std::tie(rhs.block_counts_, rhs.thread_counts_per_block_);
91+
}
92+
93+
friend bool operator!=(const LaunchDimensions& lhs,
94+
const LaunchDimensions& rhs) {
95+
return !(lhs == rhs);
96+
}
97+
8298
private:
8399
se::BlockDim block_counts_;
84100
se::ThreadDim thread_counts_per_block_;

‎third_party/xla/xla/service/gpu/launch_dimensions.proto‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
syntax = "proto3";
2+
3+
package xla.gpu;
4+
5+
import "xla/stream_executor/launch_dim.proto";
6+
7+
message LaunchDimensionsProto {
8+
stream_executor.BlockDimProto block_counts = 1;
9+
stream_executor.ThreadDimProto thread_counts_per_block = 2;
10+
}

‎third_party/xla/xla/service/gpu/launch_dimensions_test.cc‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
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/service/gpu/launch_dimensions.h"
17+
18+
#include <gmock/gmock.h>
19+
#include <gtest/gtest.h>
20+
#include "xla/stream_executor/launch_dim.h"
21+
#include "xla/tsl/platform/status_matchers.h"
22+
#include "xla/tsl/platform/statusor.h"
23+
#include "xla/tsl/util/proto/proto_matchers.h"
24+
#include "tsl/platform/protobuf.h"
25+
26+
namespace xla::gpu {
27+
namespace {
28+
using ::tsl::proto_testing::EqualsProto;
29+
using ::tsl::testing::IsOkAndHolds;
30+
31+
TEST(LaunchDimensionsTest, DefaultConstruction) {
32+
LaunchDimensions dimensions{};
33+
EXPECT_EQ(dimensions.block_counts().x, 1);
34+
EXPECT_EQ(dimensions.block_counts().y, 1);
35+
EXPECT_EQ(dimensions.block_counts().z, 1);
36+
EXPECT_EQ(dimensions.thread_counts_per_block().x, 1);
37+
EXPECT_EQ(dimensions.thread_counts_per_block().y, 1);
38+
EXPECT_EQ(dimensions.thread_counts_per_block().z, 1);
39+
}
40+
41+
TEST(LaunchDimensionsTest, LinearConstruction) {
42+
constexpr int kNumBlocks = 42;
43+
constexpr int kNumThreads = 123;
44+
LaunchDimensions dimensions{kNumBlocks, kNumThreads};
45+
EXPECT_EQ(dimensions.block_counts().x, kNumBlocks);
46+
EXPECT_EQ(dimensions.block_counts().y, 1);
47+
EXPECT_EQ(dimensions.block_counts().z, 1);
48+
EXPECT_EQ(dimensions.thread_counts_per_block().x, kNumThreads);
49+
EXPECT_EQ(dimensions.thread_counts_per_block().y, 1);
50+
EXPECT_EQ(dimensions.thread_counts_per_block().z, 1);
51+
}
52+
53+
TEST(LaunchDimensionsTest, ToProto) {
54+
se::BlockDim block_dimensions{1, 2, 3};
55+
se::ThreadDim thread_dimensions{4, 5, 6};
56+
LaunchDimensions dimensions{block_dimensions, thread_dimensions};
57+
58+
LaunchDimensionsProto proto = dimensions.ToProto();
59+
EXPECT_THAT(proto, EqualsProto(R"pb(
60+
block_counts { coordinates { x: 1, y: 2, z: 3 } }
61+
thread_counts_per_block { coordinates { x: 4 y: 5 z: 6 } }
62+
)pb"));
63+
}
64+
65+
TEST(LaunchDimensionsTest, ToAndFromProto) {
66+
se::BlockDim block_dimensions{10, 20, 30};
67+
se::ThreadDim thread_dimensions{44, 55, 66};
68+
LaunchDimensions dimensions{block_dimensions, thread_dimensions};
69+
70+
EXPECT_THAT(LaunchDimensions::FromProto(dimensions.ToProto()),
71+
IsOkAndHolds(dimensions));
72+
}
73+
74+
TEST(LaunchDimensionsTest, FromAndToProto) {
75+
LaunchDimensionsProto proto;
76+
ASSERT_TRUE(tsl::protobuf::TextFormat::ParseFromString(
77+
R"pb(
78+
block_counts { coordinates { x: 10, y: 20, z: 30 } }
79+
thread_counts_per_block { coordinates { x: 44 y: 55 z: 66 } }
80+
)pb",
81+
&proto));
82+
83+
TF_ASSERT_OK_AND_ASSIGN(LaunchDimensions launch_dimensions,
84+
LaunchDimensions::FromProto(proto));
85+
EXPECT_THAT(launch_dimensions.ToProto(), EqualsProto(proto));
86+
}
87+
88+
} // namespace
89+
} // namespace xla::gpu

0 commit comments

Comments
 (0)