Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[pjrt] Removed unused CreateDeviceToHostChannelHandle, CreateChannelHandle and SupportsSendRecvCallbacks #21136

Merged
merged 1 commit into from
Jan 8, 2025
Merged

[pjrt] Removed unused CreateDeviceToHostChannelHandle, CreateChannelHandle and SupportsSendRecvCallbacks #21136

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 0 additions & 7 deletions xla/pjrt/cpu/cpu_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,13 +202,6 @@ class TfrtCpuClient final : public PjRtClient {
std::function<void()> on_delete_callback,
std::optional<std::intptr_t> stream) override;

absl::StatusOr<ChannelHandle> CreateChannelHandle() override {
return Unimplemented("CreateChannelHandle not implemented.");
}
absl::StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() override {
return Unimplemented("CreateDeviceToHostChannelHandle not implemented.");
}

absl::Status Defragment() override {
return Unimplemented("Defragment not implemented.");
}
Expand Down
16 changes: 0 additions & 16 deletions xla/pjrt/pjrt_c_api_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -401,28 +401,12 @@ class PjRtCApiClient : public PjRtClient {
"this feature.");
}

absl::StatusOr<ChannelHandle> CreateChannelHandle() override {
return Unimplemented(
"PJRT C API does not support CreateChannelHandle. Please report an "
"issue at https://github.com/google/jax/issues if you need this "
"feature.");
}

absl::StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() override {
return Unimplemented(
"PJRT C API does not support CreateDeviceToHostChannelHandle. Please "
"report an issue at https://github.com/google/jax/issues if you need "
"this feature.");
}

absl::Status Defragment() override {
return Unimplemented(
"PJRT C API does not support Defragment. Please report an issue at "
"https://github.com/google/jax/issues if you need this feature.");
}

bool SupportsSendRecvCallbacks() const override { return true; }

const PJRT_Api* pjrt_c_api() const;

PJRT_Client* pjrt_c_client() { return c_client_.get(); }
Expand Down
13 changes: 0 additions & 13 deletions xla/pjrt/pjrt_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1070,25 +1070,12 @@ class PjRtClient {
"MakeCrossHostReceiveBuffersForGather is not implemented.");
}

// Create ChannelHandles for XLA send/recv.
virtual absl::StatusOr<ChannelHandle> CreateChannelHandle() {
return Unimplemented("CreateChannelHandle is not implemented.");
}
virtual absl::StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() {
return Unimplemented("CreateDeviceToHostChannelHandle is not implemented.");
}

// TODO(zhangqiaorjc): Experimental API to be removed.
// Defragment device memory.
virtual absl::Status Defragment() {
return Unimplemented("Defragment is not implemented.");
}

// If false, this client does not support send/recv host callbacks, and
// callers should not set the `send_callbacks` and `recv_callbacks` arguments
// in ExecuteOptions.
virtual bool SupportsSendRecvCallbacks() const { return false; }

// Return the PjRtHostMemoryForDeviceManager for this client. It can be
// nullptr if the implementation does not provide one.
virtual PjRtHostMemoryForDeviceManager* GetPjRtHostMemoryForDeviceManager()
Expand Down
7 changes: 0 additions & 7 deletions xla/pjrt/pjrt_stream_executor_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -394,13 +394,6 @@ class PjRtStreamExecutorClient : public PjRtClient {
std::function<void()> on_delete_callback,
std::optional<std::intptr_t> stream) override;

absl::StatusOr<ChannelHandle> CreateChannelHandle() override {
return client()->CreateChannelHandle();
}
absl::StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() override {
return client()->CreateDeviceToHostChannelHandle();
}

// TODO(zhangqiaorjc): Experimental. Will be removed.
absl::Status Defragment() override {
return Unimplemented("Defragment not implemented");
Expand Down
6 changes: 0 additions & 6 deletions xla/pjrt/tf_pjrt_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -340,12 +340,6 @@ class TfPjRtClient : public PjRtClient {
return wrapped_->MakeCrossHostReceiveBuffersForGather(
shapes, std::move(gather_details), device, std::move(notifier));
}
absl::StatusOr<ChannelHandle> CreateChannelHandle() override {
return wrapped_->CreateChannelHandle();
}
absl::StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() override {
return wrapped_->CreateDeviceToHostChannelHandle();
}
absl::StatusOr<const PjRtTopologyDescription*> GetTopologyDescription()
const override {
return wrapped_->GetTopologyDescription();
Expand Down
Loading