diff --git a/xla/service/cpu/runtime/BUILD b/xla/backends/cpu/runtime/BUILD similarity index 100% rename from xla/service/cpu/runtime/BUILD rename to xla/backends/cpu/runtime/BUILD diff --git a/xla/service/cpu/runtime/all_gather_thunk.cc b/xla/backends/cpu/runtime/all_gather_thunk.cc similarity index 95% rename from xla/service/cpu/runtime/all_gather_thunk.cc rename to xla/backends/cpu/runtime/all_gather_thunk.cc index 3bb705ebf9fcd..fa55bbc48dbff 100644 --- a/xla/service/cpu/runtime/all_gather_thunk.cc +++ b/xla/backends/cpu/runtime/all_gather_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/all_gather_thunk.h" +#include "xla/backends/cpu/runtime/all_gather_thunk.h" #include #include @@ -24,11 +24,11 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" #include "xla/service/collective_ops_utils.h" #include "xla/service/cpu/collectives_interface.h" -#include "xla/service/cpu/runtime/collective_thunk.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/xla/service/cpu/runtime/all_gather_thunk.h b/xla/backends/cpu/runtime/all_gather_thunk.h similarity index 85% rename from xla/service/cpu/runtime/all_gather_thunk.h rename to xla/backends/cpu/runtime/all_gather_thunk.h index 28ba6c6ace84a..2d2dca9a7eac9 100644 --- a/xla/service/cpu/runtime/all_gather_thunk.h +++ b/xla/backends/cpu/runtime/all_gather_thunk.h @@ -13,13 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_ALL_GATHER_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_ALL_GATHER_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_ALL_GATHER_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_ALL_GATHER_THUNK_H_ #include #include "absl/status/statusor.h" -#include "xla/service/cpu/runtime/collective_thunk.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/xla_data.pb.h" @@ -40,4 +40,4 @@ class AllGatherThunk final : public CollectiveThunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_ALL_GATHER_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_ALL_GATHER_THUNK_H_ diff --git a/xla/service/cpu/runtime/all_reduce_thunk.cc b/xla/backends/cpu/runtime/all_reduce_thunk.cc similarity index 96% rename from xla/service/cpu/runtime/all_reduce_thunk.cc rename to xla/backends/cpu/runtime/all_reduce_thunk.cc index 923d03ce7fd46..a5d9d283867c2 100644 --- a/xla/service/cpu/runtime/all_reduce_thunk.cc +++ b/xla/backends/cpu/runtime/all_reduce_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/all_reduce_thunk.h" +#include "xla/backends/cpu/runtime/all_reduce_thunk.h" #include #include @@ -26,12 +26,12 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/primitive_util.h" #include "xla/service/buffer_assignment.h" #include "xla/service/collective_ops_utils.h" #include "xla/service/cpu/collectives_interface.h" -#include "xla/service/cpu/runtime/collective_thunk.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/xla/service/cpu/runtime/all_reduce_thunk.h b/xla/backends/cpu/runtime/all_reduce_thunk.h similarity index 87% rename from xla/service/cpu/runtime/all_reduce_thunk.h rename to xla/backends/cpu/runtime/all_reduce_thunk.h index f4580b0f63be4..77866382353e0 100644 --- a/xla/service/cpu/runtime/all_reduce_thunk.h +++ b/xla/backends/cpu/runtime/all_reduce_thunk.h @@ -13,14 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_ALL_REDUCE_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_ALL_REDUCE_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_ALL_REDUCE_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_ALL_REDUCE_THUNK_H_ #include #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" #include "xla/service/collective_ops_utils.h" -#include "xla/service/cpu/runtime/collective_thunk.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/xla_data.pb.h" @@ -45,4 +45,4 @@ class AllReduceThunk final : public CollectiveThunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_ALL_REDUCE_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_ALL_REDUCE_THUNK_H_ diff --git a/xla/service/cpu/runtime/all_to_all_thunk.cc b/xla/backends/cpu/runtime/all_to_all_thunk.cc similarity index 96% rename from xla/service/cpu/runtime/all_to_all_thunk.cc rename to xla/backends/cpu/runtime/all_to_all_thunk.cc index d55486602d654..8badd0c4e7e23 100644 --- a/xla/service/cpu/runtime/all_to_all_thunk.cc +++ b/xla/backends/cpu/runtime/all_to_all_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/all_to_all_thunk.h" +#include "xla/backends/cpu/runtime/all_to_all_thunk.h" #include #include @@ -23,11 +23,11 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" #include "xla/service/collective_ops_utils.h" #include "xla/service/cpu/collectives_interface.h" -#include "xla/service/cpu/runtime/collective_thunk.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/xla/service/cpu/runtime/all_to_all_thunk.h b/xla/backends/cpu/runtime/all_to_all_thunk.h similarity index 85% rename from xla/service/cpu/runtime/all_to_all_thunk.h rename to xla/backends/cpu/runtime/all_to_all_thunk.h index 0c24627354829..b58afe9439457 100644 --- a/xla/service/cpu/runtime/all_to_all_thunk.h +++ b/xla/backends/cpu/runtime/all_to_all_thunk.h @@ -13,13 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_ALL_TO_ALL_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_ALL_TO_ALL_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_ALL_TO_ALL_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_ALL_TO_ALL_THUNK_H_ #include #include "absl/status/statusor.h" -#include "xla/service/cpu/runtime/collective_thunk.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/xla_data.pb.h" @@ -40,4 +40,4 @@ class AllToAllThunk final : public CollectiveThunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_ALL_TO_ALL_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_ALL_TO_ALL_THUNK_H_ diff --git a/xla/service/cpu/runtime/buffer_allocations.h b/xla/backends/cpu/runtime/buffer_allocations.h similarity index 96% rename from xla/service/cpu/runtime/buffer_allocations.h rename to xla/backends/cpu/runtime/buffer_allocations.h index 4d757261c5a39..44d71712a9c19 100644 --- a/xla/service/cpu/runtime/buffer_allocations.h +++ b/xla/backends/cpu/runtime/buffer_allocations.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_BUFFER_ALLOCATIONS_H_ -#define XLA_SERVICE_CPU_RUNTIME_BUFFER_ALLOCATIONS_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_BUFFER_ALLOCATIONS_H_ +#define XLA_BACKENDS_CPU_RUNTIME_BUFFER_ALLOCATIONS_H_ #include #include @@ -141,4 +141,4 @@ BufferAllocations::GetDeviceAddressUnchecked( } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_BUFFER_ALLOCATIONS_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_BUFFER_ALLOCATIONS_H_ diff --git a/xla/service/cpu/runtime/buffer_allocations_test.cc b/xla/backends/cpu/runtime/buffer_allocations_test.cc similarity index 97% rename from xla/service/cpu/runtime/buffer_allocations_test.cc rename to xla/backends/cpu/runtime/buffer_allocations_test.cc index 9fd7d447825de..c92be6205ac91 100644 --- a/xla/service/cpu/runtime/buffer_allocations_test.cc +++ b/xla/backends/cpu/runtime/buffer_allocations_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/buffer_allocations.h" #include #include diff --git a/xla/service/cpu/runtime/call_thunk.cc b/xla/backends/cpu/runtime/call_thunk.cc similarity index 92% rename from xla/service/cpu/runtime/call_thunk.cc rename to xla/backends/cpu/runtime/call_thunk.cc index a0a4d2bf5c967..0473ad78e40f4 100644 --- a/xla/service/cpu/runtime/call_thunk.cc +++ b/xla/backends/cpu/runtime/call_thunk.cc @@ -13,15 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/call_thunk.h" +#include "xla/backends/cpu/runtime/call_thunk.h" #include #include #include "absl/memory/memory.h" #include "absl/status/statusor.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_executor.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "tsl/platform/statusor.h" #include "tsl/profiler/lib/traceme.h" diff --git a/xla/service/cpu/runtime/call_thunk.h b/xla/backends/cpu/runtime/call_thunk.h similarity index 84% rename from xla/service/cpu/runtime/call_thunk.h rename to xla/backends/cpu/runtime/call_thunk.h index e6c9ecbd3544e..b7addf7297c39 100644 --- a/xla/service/cpu/runtime/call_thunk.h +++ b/xla/backends/cpu/runtime/call_thunk.h @@ -13,14 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_CALL_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_CALL_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_CALL_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_CALL_THUNK_H_ #include #include "absl/status/statusor.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_executor.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include "xla/tsl/concurrency/async_value_ref.h" namespace xla::cpu { @@ -45,4 +45,4 @@ class CallThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_CALL_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_CALL_THUNK_H_ diff --git a/xla/service/cpu/runtime/collective_permute_thunk.cc b/xla/backends/cpu/runtime/collective_permute_thunk.cc similarity index 96% rename from xla/service/cpu/runtime/collective_permute_thunk.cc rename to xla/backends/cpu/runtime/collective_permute_thunk.cc index 1908c3ff66e40..a830c0f7fd4ea 100644 --- a/xla/service/cpu/runtime/collective_permute_thunk.cc +++ b/xla/backends/cpu/runtime/collective_permute_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/collective_permute_thunk.h" +#include "xla/backends/cpu/runtime/collective_permute_thunk.h" #include #include @@ -28,12 +28,12 @@ limitations under the License. #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" #include "xla/service/collective_ops_utils.h" #include "xla/service/computation_placer.h" #include "xla/service/cpu/collectives_interface.h" -#include "xla/service/cpu/runtime/collective_thunk.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/status_macros.h" diff --git a/xla/service/cpu/runtime/collective_permute_thunk.h b/xla/backends/cpu/runtime/collective_permute_thunk.h similarity index 86% rename from xla/service/cpu/runtime/collective_permute_thunk.h rename to xla/backends/cpu/runtime/collective_permute_thunk.h index 6478ced6f1e93..702b2f2b15f3d 100644 --- a/xla/service/cpu/runtime/collective_permute_thunk.h +++ b/xla/backends/cpu/runtime/collective_permute_thunk.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_COLLECTIVE_PERMUTE_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_COLLECTIVE_PERMUTE_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_COLLECTIVE_PERMUTE_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_COLLECTIVE_PERMUTE_THUNK_H_ #include #include @@ -23,7 +23,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/types/span.h" -#include "xla/service/cpu/runtime/collective_thunk.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/xla_data.pb.h" @@ -51,4 +51,4 @@ class CollectivePermuteThunk final : public CollectiveThunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_COLLECTIVE_PERMUTE_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_COLLECTIVE_PERMUTE_THUNK_H_ diff --git a/xla/service/cpu/runtime/collective_thunk.cc b/xla/backends/cpu/runtime/collective_thunk.cc similarity index 98% rename from xla/service/cpu/runtime/collective_thunk.cc rename to xla/backends/cpu/runtime/collective_thunk.cc index 32a452a6bcdd0..4bebdd09cd31c 100644 --- a/xla/service/cpu/runtime/collective_thunk.cc +++ b/xla/backends/cpu/runtime/collective_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/collective_thunk.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" #include #include @@ -32,13 +32,13 @@ limitations under the License. #include "absl/strings/str_join.h" #include "absl/time/time.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" #include "xla/service/collective_ops_utils.h" #include "xla/service/computation_placer.h" #include "xla/service/cpu/collectives_interface.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/global_device_id.h" #include "xla/shape.h" #include "xla/status_macros.h" diff --git a/xla/service/cpu/runtime/collective_thunk.h b/xla/backends/cpu/runtime/collective_thunk.h similarity index 94% rename from xla/service/cpu/runtime/collective_thunk.h rename to xla/backends/cpu/runtime/collective_thunk.h index 5ae9c98844f88..8efc767838806 100644 --- a/xla/service/cpu/runtime/collective_thunk.h +++ b/xla/backends/cpu/runtime/collective_thunk.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_COLLECTIVE_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_COLLECTIVE_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_COLLECTIVE_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_COLLECTIVE_THUNK_H_ #include #include @@ -27,11 +27,11 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/time/time.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" #include "xla/service/collective_ops_utils.h" #include "xla/service/cpu/collectives_interface.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/global_device_id.h" #include "xla/shape.h" #include "xla/stream_executor/device_memory.h" @@ -122,4 +122,4 @@ class CollectiveThunk : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_COLLECTIVE_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_COLLECTIVE_THUNK_H_ diff --git a/xla/service/cpu/runtime/conditional_thunk.cc b/xla/backends/cpu/runtime/conditional_thunk.cc similarity index 96% rename from xla/service/cpu/runtime/conditional_thunk.cc rename to xla/backends/cpu/runtime/conditional_thunk.cc index 4ee46a975e621..42246dd1d3df5 100644 --- a/xla/service/cpu/runtime/conditional_thunk.cc +++ b/xla/backends/cpu/runtime/conditional_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/conditional_thunk.h" +#include "xla/backends/cpu/runtime/conditional_thunk.h" #include #include @@ -22,10 +22,10 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/strings/str_format.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_executor.h" #include "xla/stream_executor/device_memory.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/util.h" diff --git a/xla/service/cpu/runtime/conditional_thunk.h b/xla/backends/cpu/runtime/conditional_thunk.h similarity index 85% rename from xla/service/cpu/runtime/conditional_thunk.h rename to xla/backends/cpu/runtime/conditional_thunk.h index 6185b6dad9b27..0b01d8517a6ff 100644 --- a/xla/service/cpu/runtime/conditional_thunk.h +++ b/xla/backends/cpu/runtime/conditional_thunk.h @@ -13,16 +13,16 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_CONDITIONAL_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_CONDITIONAL_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_CONDITIONAL_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_CONDITIONAL_THUNK_H_ #include #include #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_executor.h" #include "xla/tsl/concurrency/async_value_ref.h" namespace xla::cpu { @@ -48,4 +48,4 @@ class ConditionalThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_CONDITIONAL_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_CONDITIONAL_THUNK_H_ diff --git a/xla/service/cpu/runtime/conditional_thunk_test.cc b/xla/backends/cpu/runtime/conditional_thunk_test.cc similarity index 92% rename from xla/service/cpu/runtime/conditional_thunk_test.cc rename to xla/backends/cpu/runtime/conditional_thunk_test.cc index d24a58dec3edc..a5222a8de6bb3 100644 --- a/xla/service/cpu/runtime/conditional_thunk_test.cc +++ b/xla/backends/cpu/runtime/conditional_thunk_test.cc @@ -13,18 +13,18 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/conditional_thunk.h" +#include "xla/backends/cpu/runtime/conditional_thunk.h" #include #include #include #include +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_testlib.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_testlib.h" #include "tsl/platform/statusor.h" #include "tsl/platform/test.h" diff --git a/xla/service/cpu/runtime/conv_impl.cc b/xla/backends/cpu/runtime/conv_impl.cc similarity index 98% rename from xla/service/cpu/runtime/conv_impl.cc rename to xla/backends/cpu/runtime/conv_impl.cc index 199a97919fa53..1d5c6dddf4a02 100644 --- a/xla/service/cpu/runtime/conv_impl.cc +++ b/xla/backends/cpu/runtime/conv_impl.cc @@ -14,7 +14,7 @@ limitations under the License. ==============================================================================*/ #define EIGEN_USE_THREADS -#include "xla/service/cpu/runtime/conv_impl.h" +#include "xla/backends/cpu/runtime/conv_impl.h" namespace tensorflow::xla { diff --git a/xla/service/cpu/runtime/conv_impl.h b/xla/backends/cpu/runtime/conv_impl.h similarity index 98% rename from xla/service/cpu/runtime/conv_impl.h rename to xla/backends/cpu/runtime/conv_impl.h index b97bc85a4edc7..4441e46b53740 100644 --- a/xla/service/cpu/runtime/conv_impl.h +++ b/xla/backends/cpu/runtime/conv_impl.h @@ -12,8 +12,8 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_CONV_IMPL_H_ -#define XLA_SERVICE_CPU_RUNTIME_CONV_IMPL_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_CONV_IMPL_H_ +#define XLA_BACKENDS_CPU_RUNTIME_CONV_IMPL_H_ #include #include @@ -261,4 +261,4 @@ CONV3D_EXTERN_TEMPLATE(Eigen::ThreadPoolDevice, float); } // namespace xla } // namespace tensorflow -#endif // XLA_SERVICE_CPU_RUNTIME_CONV_IMPL_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_CONV_IMPL_H_ diff --git a/xla/service/cpu/runtime/convolution_thunk.cc b/xla/backends/cpu/runtime/convolution_thunk.cc similarity index 99% rename from xla/service/cpu/runtime/convolution_thunk.cc rename to xla/backends/cpu/runtime/convolution_thunk.cc index c7bdd0a2ccf18..601fbeabea887 100644 --- a/xla/service/cpu/runtime/convolution_thunk.cc +++ b/xla/backends/cpu/runtime/convolution_thunk.cc @@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/convolution_thunk.h" +#include "xla/backends/cpu/runtime/convolution_thunk.h" #define EIGEN_USE_THREADS @@ -31,10 +31,10 @@ limitations under the License. #include "absl/types/span.h" #include "Eigen/Core" #include "unsupported/Eigen/CXX11/Tensor" +#include "xla/backends/cpu/runtime/conv_impl.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/executable_run_options.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/conv_impl.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/cpu/runtime_conv2d_acl.h" #include "xla/shape.h" #include "xla/status_macros.h" diff --git a/xla/service/cpu/runtime/convolution_thunk.h b/xla/backends/cpu/runtime/convolution_thunk.h similarity index 95% rename from xla/service/cpu/runtime/convolution_thunk.h rename to xla/backends/cpu/runtime/convolution_thunk.h index d3ba117336982..9f682191ae882 100644 --- a/xla/service/cpu/runtime/convolution_thunk.h +++ b/xla/backends/cpu/runtime/convolution_thunk.h @@ -13,17 +13,17 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_CONVOLUTION_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_CONVOLUTION_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_CONVOLUTION_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_CONVOLUTION_THUNK_H_ #include #include #include "absl/container/inlined_vector.h" #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/tsl/concurrency/async_value_ref.h" @@ -123,4 +123,4 @@ class ConvolutionThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_CONVOLUTION_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_CONVOLUTION_THUNK_H_ diff --git a/xla/service/cpu/runtime/convolution_thunk_test.cc b/xla/backends/cpu/runtime/convolution_thunk_test.cc similarity index 99% rename from xla/service/cpu/runtime/convolution_thunk_test.cc rename to xla/backends/cpu/runtime/convolution_thunk_test.cc index 3671431333d59..20a75d1f97ebc 100644 --- a/xla/service/cpu/runtime/convolution_thunk_test.cc +++ b/xla/backends/cpu/runtime/convolution_thunk_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/convolution_thunk.h" +#include "xla/backends/cpu/runtime/convolution_thunk.h" #include #include @@ -25,10 +25,10 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/status/status.h" #include "Eigen/Core" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/primitive_util.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/maybe_owning_device_memory.h" #include "xla/shape.h" #include "xla/shape_util.h" diff --git a/xla/service/cpu/runtime/copy_thunk.cc b/xla/backends/cpu/runtime/copy_thunk.cc similarity index 98% rename from xla/service/cpu/runtime/copy_thunk.cc rename to xla/backends/cpu/runtime/copy_thunk.cc index 1ea16dbdbf4d5..67b4d55725695 100644 --- a/xla/service/cpu/runtime/copy_thunk.cc +++ b/xla/backends/cpu/runtime/copy_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/copy_thunk.h" +#include "xla/backends/cpu/runtime/copy_thunk.h" #define EIGEN_USE_THREADS @@ -34,10 +34,10 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "unsupported/Eigen/CXX11/Tensor" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/pjrt/transpose.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/device_memory.h" diff --git a/xla/service/cpu/runtime/copy_thunk.h b/xla/backends/cpu/runtime/copy_thunk.h similarity index 92% rename from xla/service/cpu/runtime/copy_thunk.h rename to xla/backends/cpu/runtime/copy_thunk.h index a65425c7f5427..ed2cd68df5137 100644 --- a/xla/service/cpu/runtime/copy_thunk.h +++ b/xla/backends/cpu/runtime/copy_thunk.h @@ -13,17 +13,17 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_COPY_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_COPY_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_COPY_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_COPY_THUNK_H_ #include #include #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/pjrt/transpose.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/tsl/concurrency/async_value_ref.h" @@ -69,4 +69,4 @@ class CopyThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_COPY_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_COPY_THUNK_H_ diff --git a/xla/service/cpu/runtime/copy_thunk_test.cc b/xla/backends/cpu/runtime/copy_thunk_test.cc similarity index 97% rename from xla/service/cpu/runtime/copy_thunk_test.cc rename to xla/backends/cpu/runtime/copy_thunk_test.cc index 406d6b1a8aa7d..8a8e4fb4debd2 100644 --- a/xla/service/cpu/runtime/copy_thunk_test.cc +++ b/xla/backends/cpu/runtime/copy_thunk_test.cc @@ -13,15 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/copy_thunk.h" +#include "xla/backends/cpu/runtime/copy_thunk.h" #include #include +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/layout_util.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/maybe_owning_device_memory.h" #include "xla/shape.h" #include "xla/shape_util.h" diff --git a/xla/service/cpu/runtime/custom_call_thunk.cc b/xla/backends/cpu/runtime/custom_call_thunk.cc similarity index 99% rename from xla/service/cpu/runtime/custom_call_thunk.cc rename to xla/backends/cpu/runtime/custom_call_thunk.cc index 8c6deca2d2406..966da1c85d476 100644 --- a/xla/service/cpu/runtime/custom_call_thunk.cc +++ b/xla/backends/cpu/runtime/custom_call_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/custom_call_thunk.h" +#include "xla/backends/cpu/runtime/custom_call_thunk.h" #include #include @@ -35,12 +35,12 @@ limitations under the License. #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/MLIRContext.h" #include "mlir/Support/LLVM.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/ffi/attribute_map.h" #include "xla/ffi/call_frame.h" #include "xla/ffi/ffi_api.h" #include "xla/primitive_util.h" #include "xla/runtime/buffer_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/custom_call_status.h" #include "xla/service/custom_call_status_internal.h" #include "xla/service/custom_call_target_registry.h" diff --git a/xla/service/cpu/runtime/custom_call_thunk.h b/xla/backends/cpu/runtime/custom_call_thunk.h similarity index 92% rename from xla/service/cpu/runtime/custom_call_thunk.h rename to xla/backends/cpu/runtime/custom_call_thunk.h index 901545fa9f5d1..bfea5368f7cb9 100644 --- a/xla/service/cpu/runtime/custom_call_thunk.h +++ b/xla/backends/cpu/runtime/custom_call_thunk.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_CUSTOM_CALL_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_CUSTOM_CALL_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_CUSTOM_CALL_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_CUSTOM_CALL_THUNK_H_ #include #include @@ -25,9 +25,9 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/string_view.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/ffi/call_frame.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/custom_call_status.h" #include "xla/shape.h" #include "xla/tsl/concurrency/async_value_ref.h" @@ -80,4 +80,4 @@ class CustomCallThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_CUSTOM_CALL_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_CUSTOM_CALL_THUNK_H_ diff --git a/xla/service/cpu/runtime/dot_thunk.cc b/xla/backends/cpu/runtime/dot_thunk.cc similarity index 99% rename from xla/service/cpu/runtime/dot_thunk.cc rename to xla/backends/cpu/runtime/dot_thunk.cc index c92307c52f064..418ed65ce1cbb 100644 --- a/xla/service/cpu/runtime/dot_thunk.cc +++ b/xla/backends/cpu/runtime/dot_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/dot_thunk.h" +#include "xla/backends/cpu/runtime/dot_thunk.h" #include #include @@ -30,10 +30,10 @@ limitations under the License. #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/layout_util.h" #include "xla/primitive_util.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/device_memory.h" diff --git a/xla/service/cpu/runtime/dot_thunk.h b/xla/backends/cpu/runtime/dot_thunk.h similarity index 97% rename from xla/service/cpu/runtime/dot_thunk.h rename to xla/backends/cpu/runtime/dot_thunk.h index acaa94d5bf777..61bcb8194e115 100644 --- a/xla/service/cpu/runtime/dot_thunk.h +++ b/xla/backends/cpu/runtime/dot_thunk.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_DOT_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_DOT_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_DOT_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_DOT_THUNK_H_ #define EIGEN_USE_THREADS @@ -29,9 +29,9 @@ limitations under the License. #include "absl/status/statusor.h" #include "Eigen/Core" #include "unsupported/Eigen/CXX11/Tensor" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/xla_data.pb.h" @@ -175,4 +175,4 @@ DOT_THUNK_EXTERN_MATMUL_TEMPLATE(std::complex); } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_DOT_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_DOT_THUNK_H_ diff --git a/xla/service/cpu/runtime/dot_thunk_c128.cc b/xla/backends/cpu/runtime/dot_thunk_c128.cc similarity index 92% rename from xla/service/cpu/runtime/dot_thunk_c128.cc rename to xla/backends/cpu/runtime/dot_thunk_c128.cc index cd2852e26aa98..1c791bd6fac78 100644 --- a/xla/service/cpu/runtime/dot_thunk_c128.cc +++ b/xla/backends/cpu/runtime/dot_thunk_c128.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/dot_thunk.h" +#include "xla/backends/cpu/runtime/dot_thunk.h" // NOLINT IWYU pragma: keep template void ::xla::cpu::DotThunk::TypedMatMul>( const Eigen::ThreadPoolDevice* device, void* out, void* lhs, void* rhs, diff --git a/xla/service/cpu/runtime/dot_thunk_c64.cc b/xla/backends/cpu/runtime/dot_thunk_c64.cc similarity index 92% rename from xla/service/cpu/runtime/dot_thunk_c64.cc rename to xla/backends/cpu/runtime/dot_thunk_c64.cc index 55f21cceb344b..957e2d6d85563 100644 --- a/xla/service/cpu/runtime/dot_thunk_c64.cc +++ b/xla/backends/cpu/runtime/dot_thunk_c64.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/dot_thunk.h" +#include "xla/backends/cpu/runtime/dot_thunk.h" // NOLINT IWYU pragma: keep template void ::xla::cpu::DotThunk::TypedMatMul>( const Eigen::ThreadPoolDevice* device, void* out, void* lhs, void* rhs, diff --git a/xla/service/cpu/runtime/dot_thunk_f16.cc b/xla/backends/cpu/runtime/dot_thunk_f16.cc similarity index 92% rename from xla/service/cpu/runtime/dot_thunk_f16.cc rename to xla/backends/cpu/runtime/dot_thunk_f16.cc index df04b0d1272a1..35d85c8915418 100644 --- a/xla/service/cpu/runtime/dot_thunk_f16.cc +++ b/xla/backends/cpu/runtime/dot_thunk_f16.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/dot_thunk.h" +#include "xla/backends/cpu/runtime/dot_thunk.h" // NOLINT IWYU pragma: keep template void ::xla::cpu::DotThunk::TypedMatMul( const Eigen::ThreadPoolDevice* device, void* out, void* lhs, void* rhs, diff --git a/xla/service/cpu/runtime/dot_thunk_f32.cc b/xla/backends/cpu/runtime/dot_thunk_f32.cc similarity index 93% rename from xla/service/cpu/runtime/dot_thunk_f32.cc rename to xla/backends/cpu/runtime/dot_thunk_f32.cc index d98c5d940ed3b..f3aee5501ac41 100644 --- a/xla/service/cpu/runtime/dot_thunk_f32.cc +++ b/xla/backends/cpu/runtime/dot_thunk_f32.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/dot_thunk.h" +#include "xla/backends/cpu/runtime/dot_thunk.h" // NOLINT IWYU pragma: keep #if defined(TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL) #include "xla/tsl/framework/contraction/eigen_contraction_kernel.h" // IWYU pragma: keep diff --git a/xla/service/cpu/runtime/dot_thunk_f64.cc b/xla/backends/cpu/runtime/dot_thunk_f64.cc similarity index 91% rename from xla/service/cpu/runtime/dot_thunk_f64.cc rename to xla/backends/cpu/runtime/dot_thunk_f64.cc index f782cc7045ff7..bcb8bd676af8d 100644 --- a/xla/service/cpu/runtime/dot_thunk_f64.cc +++ b/xla/backends/cpu/runtime/dot_thunk_f64.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/dot_thunk.h" +#include "xla/backends/cpu/runtime/dot_thunk.h" // NOLINT IWYU pragma: keep template void ::xla::cpu::DotThunk::TypedMatMul( const Eigen::ThreadPoolDevice* device, void* out, void* lhs, void* rhs, diff --git a/xla/service/cpu/runtime/dot_thunk_s32.cc b/xla/backends/cpu/runtime/dot_thunk_s32.cc similarity index 91% rename from xla/service/cpu/runtime/dot_thunk_s32.cc rename to xla/backends/cpu/runtime/dot_thunk_s32.cc index 59186ec8a5669..0851e01b539c0 100644 --- a/xla/service/cpu/runtime/dot_thunk_s32.cc +++ b/xla/backends/cpu/runtime/dot_thunk_s32.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/dot_thunk.h" +#include "xla/backends/cpu/runtime/dot_thunk.h" // NOLINT IWYU pragma: keep template void ::xla::cpu::DotThunk::TypedMatMul( const Eigen::ThreadPoolDevice* device, void* out, void* lhs, void* rhs, diff --git a/xla/service/cpu/runtime/fft_thunk.cc b/xla/backends/cpu/runtime/fft_thunk.cc similarity index 98% rename from xla/service/cpu/runtime/fft_thunk.cc rename to xla/backends/cpu/runtime/fft_thunk.cc index 5d792c2fc8c16..b7c898b26d177 100644 --- a/xla/service/cpu/runtime/fft_thunk.cc +++ b/xla/backends/cpu/runtime/fft_thunk.cc @@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/fft_thunk.h" +#include "xla/backends/cpu/runtime/fft_thunk.h" #include #include @@ -21,10 +21,10 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/status/statusor.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/layout_util.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/cpu/runtime_fft.h" #include "xla/service/cpu/runtime_single_threaded_fft.h" #include "xla/shape.h" diff --git a/xla/service/cpu/runtime/fft_thunk.h b/xla/backends/cpu/runtime/fft_thunk.h similarity index 92% rename from xla/service/cpu/runtime/fft_thunk.h rename to xla/backends/cpu/runtime/fft_thunk.h index b63ed5e9b744e..64d4063d828cf 100644 --- a/xla/service/cpu/runtime/fft_thunk.h +++ b/xla/backends/cpu/runtime/fft_thunk.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_FFT_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_FFT_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_FFT_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_FFT_THUNK_H_ #include #include @@ -22,8 +22,8 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/stream_executor/stream_executor.h" #include "xla/tsl/concurrency/async_value_ref.h" @@ -68,4 +68,4 @@ class FftThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_FFT_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_FFT_THUNK_H_ diff --git a/xla/service/cpu/runtime/infeed_thunk.cc b/xla/backends/cpu/runtime/infeed_thunk.cc similarity index 96% rename from xla/service/cpu/runtime/infeed_thunk.cc rename to xla/backends/cpu/runtime/infeed_thunk.cc index 9e8acff4ecb27..e1a601565c69d 100644 --- a/xla/service/cpu/runtime/infeed_thunk.cc +++ b/xla/backends/cpu/runtime/infeed_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/infeed_thunk.h" +#include "xla/backends/cpu/runtime/infeed_thunk.h" #include #include @@ -24,10 +24,10 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/cpu/xfeed_manager.h" #include "xla/stream_executor/device_memory.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/xla/service/cpu/runtime/infeed_thunk.h b/xla/backends/cpu/runtime/infeed_thunk.h similarity index 87% rename from xla/service/cpu/runtime/infeed_thunk.h rename to xla/backends/cpu/runtime/infeed_thunk.h index 622046f2e3785..1d4225d1ddd00 100644 --- a/xla/service/cpu/runtime/infeed_thunk.h +++ b/xla/backends/cpu/runtime/infeed_thunk.h @@ -13,17 +13,17 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_INFEED_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_INFEED_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_INFEED_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_INFEED_THUNK_H_ #include #include #include "absl/status/statusor.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/tsl/concurrency/async_value_ref.h" @@ -61,4 +61,4 @@ class InfeedThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_INFEED_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_INFEED_THUNK_H_ diff --git a/xla/service/cpu/runtime/infeed_thunk_test.cc b/xla/backends/cpu/runtime/infeed_thunk_test.cc similarity index 92% rename from xla/service/cpu/runtime/infeed_thunk_test.cc rename to xla/backends/cpu/runtime/infeed_thunk_test.cc index 53394e242c56a..3bbb4272f2283 100644 --- a/xla/service/cpu/runtime/infeed_thunk_test.cc +++ b/xla/backends/cpu/runtime/infeed_thunk_test.cc @@ -13,14 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/infeed_thunk.h" +#include "xla/backends/cpu/runtime/infeed_thunk.h" #include +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape_util.h" #include "tsl/platform/statusor.h" #include "tsl/platform/test.h" diff --git a/xla/service/cpu/runtime/kernel_thunk.cc b/xla/backends/cpu/runtime/kernel_thunk.cc similarity index 98% rename from xla/service/cpu/runtime/kernel_thunk.cc rename to xla/backends/cpu/runtime/kernel_thunk.cc index 5ab801cc42aab..6b50f1427b837 100644 --- a/xla/service/cpu/runtime/kernel_thunk.cc +++ b/xla/backends/cpu/runtime/kernel_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/kernel_thunk.h" +#include "xla/backends/cpu/runtime/kernel_thunk.h" #define EIGEN_USE_THREADS @@ -35,10 +35,10 @@ limitations under the License. #include "absl/synchronization/mutex.h" #include "absl/types/span.h" #include "unsupported/Eigen/CXX11/Tensor" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/stream_executor/device_memory.h" #include "xla/stream_executor/host/host_kernel.h" #include "xla/stream_executor/host/host_kernel_c_api.h" diff --git a/xla/service/cpu/runtime/kernel_thunk.h b/xla/backends/cpu/runtime/kernel_thunk.h similarity index 96% rename from xla/service/cpu/runtime/kernel_thunk.h rename to xla/backends/cpu/runtime/kernel_thunk.h index 134602f99537b..3d5a872b692f4 100644 --- a/xla/service/cpu/runtime/kernel_thunk.h +++ b/xla/backends/cpu/runtime/kernel_thunk.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_KERNEL_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_KERNEL_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_KERNEL_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_KERNEL_THUNK_H_ #include #include @@ -32,8 +32,8 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/stream_executor/host/host_kernel.h" #include "xla/stream_executor/host/host_kernel_c_api.h" #include "xla/stream_executor/launch_dim.h" @@ -157,4 +157,4 @@ class KernelThunk final : public internal::KernelThunk<> { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_KERNEL_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_KERNEL_THUNK_H_ diff --git a/xla/service/cpu/runtime/kernel_thunk_test.cc b/xla/backends/cpu/runtime/kernel_thunk_test.cc similarity index 95% rename from xla/service/cpu/runtime/kernel_thunk_test.cc rename to xla/backends/cpu/runtime/kernel_thunk_test.cc index 63696c0e83278..eddd9836923de 100644 --- a/xla/service/cpu/runtime/kernel_thunk_test.cc +++ b/xla/backends/cpu/runtime/kernel_thunk_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/kernel_thunk.h" +#include "xla/backends/cpu/runtime/kernel_thunk.h" #include #include @@ -22,9 +22,9 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/match.h" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/maybe_owning_device_memory.h" #include "xla/stream_executor/device_memory.h" #include "xla/stream_executor/host/host_kernel_c_api.h" diff --git a/xla/service/cpu/runtime/logical_id_thunk.cc b/xla/backends/cpu/runtime/logical_id_thunk.cc similarity index 97% rename from xla/service/cpu/runtime/logical_id_thunk.cc rename to xla/backends/cpu/runtime/logical_id_thunk.cc index 61c8f4c801db3..ace52302dc953 100644 --- a/xla/service/cpu/runtime/logical_id_thunk.cc +++ b/xla/backends/cpu/runtime/logical_id_thunk.cc @@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/logical_id_thunk.h" +#include "xla/backends/cpu/runtime/logical_id_thunk.h" #include #include @@ -22,10 +22,10 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" #include "xla/service/computation_placer.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/global_device_id.h" #include "xla/status_macros.h" #include "xla/stream_executor/device_memory.h" diff --git a/xla/service/cpu/runtime/logical_id_thunk.h b/xla/backends/cpu/runtime/logical_id_thunk.h similarity index 90% rename from xla/service/cpu/runtime/logical_id_thunk.h rename to xla/backends/cpu/runtime/logical_id_thunk.h index bb4d2fd12840f..6a42fe69963d1 100644 --- a/xla/service/cpu/runtime/logical_id_thunk.h +++ b/xla/backends/cpu/runtime/logical_id_thunk.h @@ -13,16 +13,16 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_LOGICAL_ID_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_LOGICAL_ID_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_LOGICAL_ID_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_LOGICAL_ID_THUNK_H_ #include #include #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" #include "xla/service/computation_placer.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/global_device_id.h" #include "xla/tsl/concurrency/async_value_ref.h" @@ -68,4 +68,4 @@ class PartitionIdThunk final } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_LOGICAL_ID_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_LOGICAL_ID_THUNK_H_ diff --git a/xla/service/cpu/runtime/logical_id_thunk_test.cc b/xla/backends/cpu/runtime/logical_id_thunk_test.cc similarity index 96% rename from xla/service/cpu/runtime/logical_id_thunk_test.cc rename to xla/backends/cpu/runtime/logical_id_thunk_test.cc index 72ce59f85dad5..c8dd0a60782fe 100644 --- a/xla/service/cpu/runtime/logical_id_thunk_test.cc +++ b/xla/backends/cpu/runtime/logical_id_thunk_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/logical_id_thunk.h" +#include "xla/backends/cpu/runtime/logical_id_thunk.h" #include #include @@ -22,10 +22,10 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/executable_run_options.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/maybe_owning_device_memory.h" #include "xla/stream_executor/device_memory.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/xla/service/cpu/runtime/outfeed_thunk.cc b/xla/backends/cpu/runtime/outfeed_thunk.cc similarity index 96% rename from xla/service/cpu/runtime/outfeed_thunk.cc rename to xla/backends/cpu/runtime/outfeed_thunk.cc index a56ae0c437a7e..b541953a403de 100644 --- a/xla/service/cpu/runtime/outfeed_thunk.cc +++ b/xla/backends/cpu/runtime/outfeed_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/outfeed_thunk.h" +#include "xla/backends/cpu/runtime/outfeed_thunk.h" #include #include @@ -23,10 +23,10 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/strings/str_format.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/cpu/xfeed_manager.h" #include "xla/stream_executor/device_memory.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/xla/service/cpu/runtime/outfeed_thunk.h b/xla/backends/cpu/runtime/outfeed_thunk.h similarity index 87% rename from xla/service/cpu/runtime/outfeed_thunk.h rename to xla/backends/cpu/runtime/outfeed_thunk.h index ff05339002ffc..74920899255d4 100644 --- a/xla/service/cpu/runtime/outfeed_thunk.h +++ b/xla/backends/cpu/runtime/outfeed_thunk.h @@ -13,17 +13,17 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_OUTFEED_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_OUTFEED_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_OUTFEED_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_OUTFEED_THUNK_H_ #include #include #include "absl/status/statusor.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/tsl/concurrency/async_value_ref.h" @@ -60,4 +60,4 @@ class OutfeedThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_OUTFEED_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_OUTFEED_THUNK_H_ diff --git a/xla/service/cpu/runtime/outfeed_thunk_test.cc b/xla/backends/cpu/runtime/outfeed_thunk_test.cc similarity index 92% rename from xla/service/cpu/runtime/outfeed_thunk_test.cc rename to xla/backends/cpu/runtime/outfeed_thunk_test.cc index 2c6b9b9a91123..0139a95f777e4 100644 --- a/xla/service/cpu/runtime/outfeed_thunk_test.cc +++ b/xla/backends/cpu/runtime/outfeed_thunk_test.cc @@ -13,14 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/outfeed_thunk.h" +#include "xla/backends/cpu/runtime/outfeed_thunk.h" #include +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape_util.h" #include "tsl/platform/statusor.h" #include "tsl/platform/test.h" diff --git a/xla/service/cpu/runtime/reduce_scatter_thunk.cc b/xla/backends/cpu/runtime/reduce_scatter_thunk.cc similarity index 96% rename from xla/service/cpu/runtime/reduce_scatter_thunk.cc rename to xla/backends/cpu/runtime/reduce_scatter_thunk.cc index 701ac3243ebd9..920aa3dc545b1 100644 --- a/xla/service/cpu/runtime/reduce_scatter_thunk.cc +++ b/xla/backends/cpu/runtime/reduce_scatter_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/reduce_scatter_thunk.h" +#include "xla/backends/cpu/runtime/reduce_scatter_thunk.h" #include #include @@ -24,12 +24,12 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/primitive_util.h" #include "xla/service/buffer_assignment.h" #include "xla/service/collective_ops_utils.h" #include "xla/service/cpu/collectives_interface.h" -#include "xla/service/cpu/runtime/collective_thunk.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/xla/service/cpu/runtime/reduce_scatter_thunk.h b/xla/backends/cpu/runtime/reduce_scatter_thunk.h similarity index 86% rename from xla/service/cpu/runtime/reduce_scatter_thunk.h rename to xla/backends/cpu/runtime/reduce_scatter_thunk.h index d37e1b22db556..104d6c354dfa8 100644 --- a/xla/service/cpu/runtime/reduce_scatter_thunk.h +++ b/xla/backends/cpu/runtime/reduce_scatter_thunk.h @@ -13,14 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_REDUCE_SCATTER_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_REDUCE_SCATTER_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_REDUCE_SCATTER_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_REDUCE_SCATTER_THUNK_H_ #include #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" #include "xla/service/collective_ops_utils.h" -#include "xla/service/cpu/runtime/collective_thunk.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/xla_data.pb.h" @@ -44,4 +44,4 @@ class ReduceScatterThunk final : public CollectiveThunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_REDUCE_SCATTER_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_REDUCE_SCATTER_THUNK_H_ diff --git a/xla/service/cpu/runtime/resource_use.cc b/xla/backends/cpu/runtime/resource_use.cc similarity index 98% rename from xla/service/cpu/runtime/resource_use.cc rename to xla/backends/cpu/runtime/resource_use.cc index 3e5ceabb9ac53..a3c03849b5178 100644 --- a/xla/service/cpu/runtime/resource_use.cc +++ b/xla/backends/cpu/runtime/resource_use.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/resource_use.h" #include diff --git a/xla/service/cpu/runtime/resource_use.h b/xla/backends/cpu/runtime/resource_use.h similarity index 96% rename from xla/service/cpu/runtime/resource_use.h rename to xla/backends/cpu/runtime/resource_use.h index 6ee1f1bfd6ac9..1442a2895a02b 100644 --- a/xla/service/cpu/runtime/resource_use.h +++ b/xla/backends/cpu/runtime/resource_use.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_RESOURCE_USE_H_ -#define XLA_SERVICE_CPU_RUNTIME_RESOURCE_USE_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_RESOURCE_USE_H_ +#define XLA_BACKENDS_CPU_RUNTIME_RESOURCE_USE_H_ #include #include @@ -111,4 +111,4 @@ class ResourceUse { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_RESOURCE_USE_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_RESOURCE_USE_H_ diff --git a/xla/service/cpu/runtime/resource_use_test.cc b/xla/backends/cpu/runtime/resource_use_test.cc similarity index 97% rename from xla/service/cpu/runtime/resource_use_test.cc rename to xla/backends/cpu/runtime/resource_use_test.cc index 4d3c9bbaf4cec..dd5115bcaf2ae 100644 --- a/xla/service/cpu/runtime/resource_use_test.cc +++ b/xla/backends/cpu/runtime/resource_use_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/resource_use.h" #include "tsl/platform/test.h" diff --git a/xla/service/cpu/runtime/rng_state_thunk.cc b/xla/backends/cpu/runtime/rng_state_thunk.cc similarity index 96% rename from xla/service/cpu/runtime/rng_state_thunk.cc rename to xla/backends/cpu/runtime/rng_state_thunk.cc index df611bd5fe169..39a3de9b9429d 100644 --- a/xla/service/cpu/runtime/rng_state_thunk.cc +++ b/xla/backends/cpu/runtime/rng_state_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/rng_state_thunk.h" +#include "xla/backends/cpu/runtime/rng_state_thunk.h" #include #include @@ -26,8 +26,8 @@ limitations under the License. #include "absl/status/status.h" #include "absl/strings/str_format.h" #include "absl/synchronization/mutex.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/stream_executor/device_memory.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/util.h" diff --git a/xla/service/cpu/runtime/rng_state_thunk.h b/xla/backends/cpu/runtime/rng_state_thunk.h similarity index 89% rename from xla/service/cpu/runtime/rng_state_thunk.h rename to xla/backends/cpu/runtime/rng_state_thunk.h index 9798ed7c105f4..d00bf4523e5de 100644 --- a/xla/service/cpu/runtime/rng_state_thunk.h +++ b/xla/backends/cpu/runtime/rng_state_thunk.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_RNG_STATE_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_RNG_STATE_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_RNG_STATE_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_RNG_STATE_THUNK_H_ #include #include @@ -23,9 +23,9 @@ limitations under the License. #include "absl/numeric/int128.h" #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" namespace xla::cpu { @@ -56,4 +56,4 @@ class RngGetAndUpdateStateThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_RNG_STATE_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_RNG_STATE_THUNK_H_ diff --git a/xla/service/cpu/runtime/sort_thunk.cc b/xla/backends/cpu/runtime/sort_thunk.cc similarity index 99% rename from xla/service/cpu/runtime/sort_thunk.cc rename to xla/backends/cpu/runtime/sort_thunk.cc index 041bf030d52ab..8d2df6f298cbc 100644 --- a/xla/service/cpu/runtime/sort_thunk.cc +++ b/xla/backends/cpu/runtime/sort_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/sort_thunk.h" +#include "xla/backends/cpu/runtime/sort_thunk.h" #include #include @@ -38,11 +38,11 @@ limitations under the License. #include "absl/strings/str_join.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/layout_util.h" #include "xla/primitive_util.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/device_memory.h" diff --git a/xla/service/cpu/runtime/sort_thunk.h b/xla/backends/cpu/runtime/sort_thunk.h similarity index 93% rename from xla/service/cpu/runtime/sort_thunk.h rename to xla/backends/cpu/runtime/sort_thunk.h index 049fa062cff60..a1c2b5eda242e 100644 --- a/xla/service/cpu/runtime/sort_thunk.h +++ b/xla/backends/cpu/runtime/sort_thunk.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_SORT_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_SORT_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_SORT_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_SORT_THUNK_H_ #include #include @@ -28,8 +28,8 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/shape.h" #include "xla/tsl/concurrency/async_value_ref.h" @@ -81,4 +81,4 @@ class SortThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_SORT_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_SORT_THUNK_H_ diff --git a/xla/service/cpu/runtime/sort_thunk_test.cc b/xla/backends/cpu/runtime/sort_thunk_test.cc similarity index 98% rename from xla/service/cpu/runtime/sort_thunk_test.cc rename to xla/backends/cpu/runtime/sort_thunk_test.cc index 4c7b2514a1c70..1f450f77548d7 100644 --- a/xla/service/cpu/runtime/sort_thunk_test.cc +++ b/xla/backends/cpu/runtime/sort_thunk_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/sort_thunk.h" +#include "xla/backends/cpu/runtime/sort_thunk.h" #include #include @@ -21,11 +21,11 @@ limitations under the License. #include #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/layout.h" #include "xla/layout_util.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/maybe_owning_device_memory.h" #include "xla/shape.h" #include "xla/shape_util.h" diff --git a/xla/service/cpu/runtime/thunk.cc b/xla/backends/cpu/runtime/thunk.cc similarity index 99% rename from xla/service/cpu/runtime/thunk.cc rename to xla/backends/cpu/runtime/thunk.cc index 9228de3d5f156..41a02a5ca3a41 100644 --- a/xla/service/cpu/runtime/thunk.cc +++ b/xla/backends/cpu/runtime/thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk.h" #include #include diff --git a/xla/service/cpu/runtime/thunk.h b/xla/backends/cpu/runtime/thunk.h similarity index 98% rename from xla/service/cpu/runtime/thunk.h rename to xla/backends/cpu/runtime/thunk.h index 9141da7462869..cfc60597e6ac6 100644 --- a/xla/service/cpu/runtime/thunk.h +++ b/xla/backends/cpu/runtime/thunk.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_THUNK_H_ #include #include @@ -31,12 +31,12 @@ limitations under the License. #include "absl/container/inlined_vector.h" #include "absl/functional/any_invocable.h" #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/resource_use.h" #include "xla/executable_run_options.h" #include "xla/ffi/execution_context.h" #include "xla/runtime/buffer_use.h" #include "xla/service/cpu/collectives_interface.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/resource_use.h" #include "xla/service/cpu/xfeed_manager.h" #include "xla/service/global_device_id.h" #include "xla/stream_executor/host/host_kernel_c_api.h" @@ -379,4 +379,4 @@ class ThunkSequence : public std::vector> { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_THUNK_H_ diff --git a/xla/service/cpu/runtime/thunk_executor.cc b/xla/backends/cpu/runtime/thunk_executor.cc similarity index 99% rename from xla/service/cpu/runtime/thunk_executor.cc rename to xla/backends/cpu/runtime/thunk_executor.cc index 9b4c735703fdc..eb32b508b3a1b 100644 --- a/xla/service/cpu/runtime/thunk_executor.cc +++ b/xla/backends/cpu/runtime/thunk_executor.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/thunk_executor.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include #include @@ -32,9 +32,9 @@ limitations under the License. #include "absl/strings/str_join.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "tsl/platform/logging.h" #include "tsl/profiler/lib/traceme.h" diff --git a/xla/service/cpu/runtime/thunk_executor.h b/xla/backends/cpu/runtime/thunk_executor.h similarity index 97% rename from xla/service/cpu/runtime/thunk_executor.h rename to xla/backends/cpu/runtime/thunk_executor.h index f0df6cfafb3d8..5ba15b0432b50 100644 --- a/xla/service/cpu/runtime/thunk_executor.h +++ b/xla/backends/cpu/runtime/thunk_executor.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_THUNK_EXECUTOR_H_ -#define XLA_SERVICE_CPU_RUNTIME_THUNK_EXECUTOR_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_THUNK_EXECUTOR_H_ +#define XLA_BACKENDS_CPU_RUNTIME_THUNK_EXECUTOR_H_ #include #include @@ -33,7 +33,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" -#include "xla/service/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/tsl/concurrency/async_value_ref.h" namespace xla::cpu { @@ -259,4 +259,4 @@ class ThunkExecutor { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_THUNK_EXECUTOR_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_THUNK_EXECUTOR_H_ diff --git a/xla/service/cpu/runtime/thunk_executor_test.cc b/xla/backends/cpu/runtime/thunk_executor_test.cc similarity index 99% rename from xla/service/cpu/runtime/thunk_executor_test.cc rename to xla/backends/cpu/runtime/thunk_executor_test.cc index 60996ebd7ed61..ebe98304b9f6f 100644 --- a/xla/service/cpu/runtime/thunk_executor_test.cc +++ b/xla/backends/cpu/runtime/thunk_executor_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/thunk_executor.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include #include @@ -30,11 +30,11 @@ limitations under the License. #include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/maybe_owning_device_memory.h" #include "xla/stream_executor/device_memory.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/xla/service/cpu/runtime/thunk_test.cc b/xla/backends/cpu/runtime/thunk_test.cc similarity index 99% rename from xla/service/cpu/runtime/thunk_test.cc rename to xla/backends/cpu/runtime/thunk_test.cc index 3b975750be6f1..1b20de023d91f 100644 --- a/xla/service/cpu/runtime/thunk_test.cc +++ b/xla/backends/cpu/runtime/thunk_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk.h" #include #include diff --git a/xla/service/cpu/runtime/thunk_testlib.h b/xla/backends/cpu/runtime/thunk_testlib.h similarity index 88% rename from xla/service/cpu/runtime/thunk_testlib.h rename to xla/backends/cpu/runtime/thunk_testlib.h index 154c2b2897270..4da0650efee7c 100644 --- a/xla/service/cpu/runtime/thunk_testlib.h +++ b/xla/backends/cpu/runtime/thunk_testlib.h @@ -13,13 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_THUNK_TESTLIB_H_ -#define XLA_SERVICE_CPU_RUNTIME_THUNK_TESTLIB_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_THUNK_TESTLIB_H_ +#define XLA_BACKENDS_CPU_RUNTIME_THUNK_TESTLIB_H_ #include "absl/status/status.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/tsl/concurrency/async_value_ref.h" namespace xla::cpu { @@ -59,4 +59,4 @@ class ResourceUseThunk : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_THUNK_TESTLIB_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_THUNK_TESTLIB_H_ diff --git a/xla/service/cpu/runtime/topk_thunk.cc b/xla/backends/cpu/runtime/topk_thunk.cc similarity index 96% rename from xla/service/cpu/runtime/topk_thunk.cc rename to xla/backends/cpu/runtime/topk_thunk.cc index 6c238224166a5..0c72933dc1a3a 100644 --- a/xla/service/cpu/runtime/topk_thunk.cc +++ b/xla/backends/cpu/runtime/topk_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/topk_thunk.h" +#include "xla/backends/cpu/runtime/topk_thunk.h" #include #include @@ -21,8 +21,8 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/cpu/runtime_topk.h" #include "xla/stream_executor/device_memory.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/xla/service/cpu/runtime/topk_thunk.h b/xla/backends/cpu/runtime/topk_thunk.h similarity index 90% rename from xla/service/cpu/runtime/topk_thunk.h rename to xla/backends/cpu/runtime/topk_thunk.h index 7b2bfb63502bf..7e7fadb03852e 100644 --- a/xla/service/cpu/runtime/topk_thunk.h +++ b/xla/backends/cpu/runtime/topk_thunk.h @@ -13,16 +13,16 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_TOPK_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_TOPK_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_TOPK_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_TOPK_THUNK_H_ #include #include #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/tsl/concurrency/async_value_ref.h" namespace xla::cpu { @@ -56,4 +56,4 @@ class TopKThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_TOPK_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_TOPK_THUNK_H_ diff --git a/xla/service/cpu/runtime/while_thunk.cc b/xla/backends/cpu/runtime/while_thunk.cc similarity index 98% rename from xla/service/cpu/runtime/while_thunk.cc rename to xla/backends/cpu/runtime/while_thunk.cc index 486a0b93e72f5..6c1e81f5dee0d 100644 --- a/xla/service/cpu/runtime/while_thunk.cc +++ b/xla/backends/cpu/runtime/while_thunk.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/while_thunk.h" +#include "xla/backends/cpu/runtime/while_thunk.h" #include #include @@ -26,11 +26,11 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_executor.h" #include "xla/stream_executor/device_memory.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "tsl/platform/logging.h" diff --git a/xla/service/cpu/runtime/while_thunk.h b/xla/backends/cpu/runtime/while_thunk.h similarity index 92% rename from xla/service/cpu/runtime/while_thunk.h rename to xla/backends/cpu/runtime/while_thunk.h index e631e54842a52..c1de07de86ad5 100644 --- a/xla/service/cpu/runtime/while_thunk.h +++ b/xla/backends/cpu/runtime/while_thunk.h @@ -13,17 +13,17 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_CPU_RUNTIME_WHILE_THUNK_H_ -#define XLA_SERVICE_CPU_RUNTIME_WHILE_THUNK_H_ +#ifndef XLA_BACKENDS_CPU_RUNTIME_WHILE_THUNK_H_ +#define XLA_BACKENDS_CPU_RUNTIME_WHILE_THUNK_H_ #include #include #include #include "absl/status/statusor.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_executor.h" #include "xla/tsl/concurrency/async_value_ref.h" namespace xla::cpu { @@ -83,4 +83,4 @@ class WhileThunk final : public Thunk { } // namespace xla::cpu -#endif // XLA_SERVICE_CPU_RUNTIME_WHILE_THUNK_H_ +#endif // XLA_BACKENDS_CPU_RUNTIME_WHILE_THUNK_H_ diff --git a/xla/service/cpu/runtime/while_thunk_test.cc b/xla/backends/cpu/runtime/while_thunk_test.cc similarity index 97% rename from xla/service/cpu/runtime/while_thunk_test.cc rename to xla/backends/cpu/runtime/while_thunk_test.cc index fc6a32c8bd715..d4b874a72b380 100644 --- a/xla/service/cpu/runtime/while_thunk_test.cc +++ b/xla/backends/cpu/runtime/while_thunk_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/cpu/runtime/while_thunk.h" +#include "xla/backends/cpu/runtime/while_thunk.h" #include #include @@ -22,12 +22,12 @@ limitations under the License. #include #include +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_testlib.h" #include "xla/runtime/buffer_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_testlib.h" #include "xla/service/maybe_owning_device_memory.h" #include "xla/stream_executor/device_memory.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/xla/pjrt/cpu/BUILD b/xla/pjrt/cpu/BUILD index b2a79210039cc..48ad7b71ebf95 100644 --- a/xla/pjrt/cpu/BUILD +++ b/xla/pjrt/cpu/BUILD @@ -149,6 +149,9 @@ cc_library( "//xla:util", "//xla:xla_data_proto_cc", "//xla:xla_proto_cc", + "//xla/backends/cpu/runtime:buffer_allocations", + "//xla/backends/cpu/runtime:thunk", + "//xla/backends/cpu/runtime:thunk_executor", "//xla/client:executable_build_options", "//xla/client:xla_computation", "//xla/hlo/ir:hlo", @@ -185,9 +188,6 @@ cc_library( "//xla/service/cpu:cpu_runtime", "//xla/service/cpu:cpu_xfeed", "//xla/service/cpu:simple_orc_jit", - "//xla/service/cpu/runtime:buffer_allocations", - "//xla/service/cpu/runtime:thunk", - "//xla/service/cpu/runtime:thunk_executor", "//xla/stream_executor", "//xla/tsl/concurrency:async_value", "//xla/tsl/concurrency:ref_count", diff --git a/xla/pjrt/cpu/cpu_client.cc b/xla/pjrt/cpu/cpu_client.cc index 78779c83028af..65ed8589f4f2b 100644 --- a/xla/pjrt/cpu/cpu_client.cc +++ b/xla/pjrt/cpu/cpu_client.cc @@ -47,6 +47,9 @@ limitations under the License. #include "unsupported/Eigen/CXX11/Tensor" #include "mlir/IR/BuiltinOps.h" #include "xla/array.h" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include "xla/client/executable_build_options.h" #include "xla/client/xla_computation.h" #include "xla/debug_options_flags.h" @@ -83,9 +86,6 @@ limitations under the License. #include "xla/service/cpu/cpu_executable_run_options.h" #include "xla/service/cpu/cpu_runtime.h" #include "xla/service/cpu/cpu_xfeed.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_executor.h" #include "xla/service/cpu/simple_orc_jit.h" #include "xla/service/custom_call_status.h" #include "xla/service/custom_call_status_internal.h" diff --git a/xla/service/BUILD b/xla/service/BUILD index 6a7c391b8b221..005e67235c29f 100644 --- a/xla/service/BUILD +++ b/xla/service/BUILD @@ -7293,6 +7293,7 @@ cc_library( ":__subpackages__", "//tensorflow/compiler/tf2xla:__pkg__", "//xla/pjrt:__subpackages__", + "//xla/backends/cpu/runtime:__subpackages__", ]), deps = [ ":custom_call_status", diff --git a/xla/service/cpu/BUILD b/xla/service/cpu/BUILD index dcb7930be637c..aa5362b63427d 100644 --- a/xla/service/cpu/BUILD +++ b/xla/service/cpu/BUILD @@ -99,7 +99,7 @@ filegroup( "runtime_matmul_f64.cc", "runtime_matmul_s32.cc", "runtime_fork_join.cc", - "//xla/service/cpu/runtime:runtime_srcs", + "//xla/backends/cpu/runtime:runtime_srcs", #"runtime_handle_ffi_call.cc", # TODO(b/338344732): Add "runtime_handle_ffi_call.cc". ], visibility = internal_visibility([":friends"]), @@ -127,7 +127,7 @@ filegroup( "runtime_fork_join.h", "runtime_lightweight_check.h", "runtime_matmul.h", - "//xla/service/cpu/runtime:runtime_hdrs", + "//xla/backends/cpu/runtime:runtime_hdrs", #"runtime_handle_ffi_call.h", # TODO(b/338344732): Add "runtime_handle_ffi_call.h" ], visibility = internal_visibility([":friends"]), @@ -239,6 +239,7 @@ cc_library( "//xla:util", "//xla:xla_data_proto_cc", "//xla:xla_proto_cc", + "//xla/backends/cpu/runtime:thunk", "//xla/hlo/ir:hlo", "//xla/hlo/ir:hlo_module_group", "//xla/mlir_hlo", @@ -328,7 +329,6 @@ cc_library( "//xla/service:while_loop_simplifier", "//xla/service:while_loop_trip_count_annotator", "//xla/service:zero_sized_hlo_elimination", - "//xla/service/cpu/runtime:thunk", "//xla/service/llvm_ir:llvm_command_line_options", "//xla/service/llvm_ir:llvm_util", "//xla/service/spmd:stateful_rng_spmd_partitioner", @@ -557,6 +557,9 @@ cc_library( "//xla:types", "//xla:util", "//xla:xla_data_proto_cc", + "//xla/backends/cpu/runtime:buffer_allocations", + "//xla/backends/cpu/runtime:thunk", + "//xla/backends/cpu/runtime:thunk_executor", "//xla/hlo/ir:hlo", "//xla/service:buffer_assignment", "//xla/service:computation_layout", @@ -570,9 +573,6 @@ cc_library( "//xla/service:maybe_owning_device_memory", "//xla/service:shaped_buffer", "//xla/service:xla_debug_info_manager", - "//xla/service/cpu/runtime:buffer_allocations", - "//xla/service/cpu/runtime:thunk", - "//xla/service/cpu/runtime:thunk_executor", "//xla/stream_executor", "//xla/stream_executor:device_memory_allocator", "//xla/stream_executor/host:host_kernel_c_api", @@ -851,33 +851,33 @@ cc_library( "//xla:shape_util", "//xla:status_macros", "//xla:util", + "//xla/backends/cpu/runtime:all_gather_thunk", + "//xla/backends/cpu/runtime:all_reduce_thunk", + "//xla/backends/cpu/runtime:all_to_all_thunk", + "//xla/backends/cpu/runtime:call_thunk", + "//xla/backends/cpu/runtime:collective_permute_thunk", + "//xla/backends/cpu/runtime:collective_thunk", + "//xla/backends/cpu/runtime:conditional_thunk", + "//xla/backends/cpu/runtime:convolution_thunk", + "//xla/backends/cpu/runtime:copy_thunk", + "//xla/backends/cpu/runtime:custom_call_thunk", + "//xla/backends/cpu/runtime:dot_thunk", + "//xla/backends/cpu/runtime:fft_thunk", + "//xla/backends/cpu/runtime:infeed_thunk", + "//xla/backends/cpu/runtime:kernel_thunk", + "//xla/backends/cpu/runtime:logical_id_thunk", + "//xla/backends/cpu/runtime:outfeed_thunk", + "//xla/backends/cpu/runtime:reduce_scatter_thunk", + "//xla/backends/cpu/runtime:resource_use", + "//xla/backends/cpu/runtime:rng_state_thunk", + "//xla/backends/cpu/runtime:sort_thunk", + "//xla/backends/cpu/runtime:thunk", + "//xla/backends/cpu/runtime:topk_thunk", + "//xla/backends/cpu/runtime:while_thunk", "//xla/hlo/ir:hlo", "//xla/service:buffer_assignment", "//xla/service:collective_ops_utils", "//xla/service:hlo_module_config", - "//xla/service/cpu/runtime:all_gather_thunk", - "//xla/service/cpu/runtime:all_reduce_thunk", - "//xla/service/cpu/runtime:all_to_all_thunk", - "//xla/service/cpu/runtime:call_thunk", - "//xla/service/cpu/runtime:collective_permute_thunk", - "//xla/service/cpu/runtime:collective_thunk", - "//xla/service/cpu/runtime:conditional_thunk", - "//xla/service/cpu/runtime:convolution_thunk", - "//xla/service/cpu/runtime:copy_thunk", - "//xla/service/cpu/runtime:custom_call_thunk", - "//xla/service/cpu/runtime:dot_thunk", - "//xla/service/cpu/runtime:fft_thunk", - "//xla/service/cpu/runtime:infeed_thunk", - "//xla/service/cpu/runtime:kernel_thunk", - "//xla/service/cpu/runtime:logical_id_thunk", - "//xla/service/cpu/runtime:outfeed_thunk", - "//xla/service/cpu/runtime:reduce_scatter_thunk", - "//xla/service/cpu/runtime:resource_use", - "//xla/service/cpu/runtime:rng_state_thunk", - "//xla/service/cpu/runtime:sort_thunk", - "//xla/service/cpu/runtime:thunk", - "//xla/service/cpu/runtime:topk_thunk", - "//xla/service/cpu/runtime:while_thunk", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status", @@ -1066,7 +1066,7 @@ cc_library( deps = [ ":runtime_lightweight_check", "//xla:executable_run_options", - "//xla/service/cpu/runtime:conv_impl", + "//xla/backends/cpu/runtime:conv_impl", "//xla/tsl/framework/contraction:eigen_contraction_kernel", "//xla/tsl/framework/convolution:eigen_helpers", "@com_google_absl//absl/base:dynamic_annotations", @@ -1084,7 +1084,7 @@ cc_library( deps = [ ":runtime_lightweight_check", "//xla:executable_run_options", - "//xla/service/cpu/runtime:conv_impl", + "//xla/backends/cpu/runtime:conv_impl", "//xla/tsl/framework/contraction:eigen_contraction_kernel", "//xla/tsl/framework/convolution:eigen_helpers", "@com_google_absl//absl/base:dynamic_annotations", @@ -1212,7 +1212,7 @@ cc_library( copts = runtime_copts(), visibility = ["//visibility:public"], deps = [ - "//xla/service/cpu/runtime:conv_impl", + "//xla/backends/cpu/runtime:conv_impl", "//xla/tsl/framework/contraction:eigen_contraction_kernel", "//xla/tsl/framework/convolution:eigen_helpers", "@com_google_absl//absl/base:dynamic_annotations", @@ -1228,7 +1228,7 @@ cc_library( copts = runtime_copts(), visibility = ["//visibility:public"], deps = [ - "//xla/service/cpu/runtime:conv_impl", + "//xla/backends/cpu/runtime:conv_impl", "//xla/tsl/framework/contraction:eigen_contraction_kernel", "//xla/tsl/framework/convolution:eigen_helpers", "@com_google_absl//absl/base:dynamic_annotations", diff --git a/xla/service/cpu/cpu_compiler.cc b/xla/service/cpu/cpu_compiler.cc index fc292c9378976..13a47eb2b36b6 100644 --- a/xla/service/cpu/cpu_compiler.cc +++ b/xla/service/cpu/cpu_compiler.cc @@ -68,6 +68,7 @@ limitations under the License. #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" #include "mlir/Target/LLVMIR/Export.h" #include "mlir/Transforms/DialectConversion.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/cpu_function_runtime.h" #include "xla/hlo/ir/dfs_hlo_visitor_with_default.h" #include "xla/hlo/ir/hlo_casting_utils.h" @@ -113,7 +114,6 @@ limitations under the License. #include "xla/service/cpu/ir_emitter.h" #include "xla/service/cpu/ir_emitter2.h" #include "xla/service/cpu/parallel_task_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/cpu/simple_orc_jit.h" #include "xla/service/cpu/target_machine_features.h" #include "xla/service/cpu/thunk_emitter.h" diff --git a/xla/service/cpu/cpu_executable.cc b/xla/service/cpu/cpu_executable.cc index 0bec73523edd9..e1f4b21317065 100644 --- a/xla/service/cpu/cpu_executable.cc +++ b/xla/service/cpu/cpu_executable.cc @@ -41,6 +41,9 @@ limitations under the License. #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h" #include "llvm/IR/Mangler.h" #include "llvm/Support/Error.h" +#include "xla/backends/cpu/runtime/buffer_allocations.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include "xla/executable_run_options.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_input_output_alias_config.h" @@ -48,9 +51,6 @@ limitations under the License. #include "xla/literal.h" #include "xla/service/buffer_assignment.h" #include "xla/service/cpu/cpu_runtime.h" -#include "xla/service/cpu/runtime/buffer_allocations.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_executor.h" #include "xla/service/cpu/simple_orc_jit.h" #include "xla/service/custom_call_status.h" #include "xla/service/custom_call_status_internal.h" diff --git a/xla/service/cpu/cpu_executable.h b/xla/service/cpu/cpu_executable.h index 8c8883d867368..2c2aa248bcbe5 100644 --- a/xla/service/cpu/cpu_executable.h +++ b/xla/service/cpu/cpu_executable.h @@ -28,13 +28,13 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/thunk_executor.h" #include "xla/executable_run_options.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/literal.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/thunk_executor.h" #include "xla/service/cpu/simple_orc_jit.h" #include "xla/service/custom_call_status.h" #include "xla/service/custom_call_status_internal.h" diff --git a/xla/service/cpu/runtime_conv2d.cc b/xla/service/cpu/runtime_conv2d.cc index 4bc0d03fe8099..b7153feba5269 100644 --- a/xla/service/cpu/runtime_conv2d.cc +++ b/xla/service/cpu/runtime_conv2d.cc @@ -20,8 +20,8 @@ limitations under the License. #define EIGEN_USE_THREADS #include "absl/base/dynamic_annotations.h" +#include "xla/backends/cpu/runtime/conv_impl.h" #include "xla/executable_run_options.h" -#include "xla/service/cpu/runtime/conv_impl.h" #include "xla/service/cpu/runtime_lightweight_check.h" ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenConv2DF32( diff --git a/xla/service/cpu/runtime_conv3d.cc b/xla/service/cpu/runtime_conv3d.cc index 7e83269e289fd..d5690956e082a 100644 --- a/xla/service/cpu/runtime_conv3d.cc +++ b/xla/service/cpu/runtime_conv3d.cc @@ -20,8 +20,8 @@ limitations under the License. #define EIGEN_USE_THREADS #include "absl/base/dynamic_annotations.h" +#include "xla/backends/cpu/runtime/conv_impl.h" #include "xla/executable_run_options.h" -#include "xla/service/cpu/runtime/conv_impl.h" #include "xla/service/cpu/runtime_lightweight_check.h" ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenConv3DF32( diff --git a/xla/service/cpu/runtime_single_threaded_conv2d.cc b/xla/service/cpu/runtime_single_threaded_conv2d.cc index a770681987400..54de2423ec07e 100644 --- a/xla/service/cpu/runtime_single_threaded_conv2d.cc +++ b/xla/service/cpu/runtime_single_threaded_conv2d.cc @@ -18,7 +18,7 @@ limitations under the License. #include #include "absl/base/dynamic_annotations.h" -#include "xla/service/cpu/runtime/conv_impl.h" +#include "xla/backends/cpu/runtime/conv_impl.h" ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenSingleThreadedConv2DF16( diff --git a/xla/service/cpu/runtime_single_threaded_conv3d.cc b/xla/service/cpu/runtime_single_threaded_conv3d.cc index 08ff94d06e7e7..6902061dc2fe4 100644 --- a/xla/service/cpu/runtime_single_threaded_conv3d.cc +++ b/xla/service/cpu/runtime_single_threaded_conv3d.cc @@ -18,7 +18,7 @@ limitations under the License. #include #include "absl/base/dynamic_annotations.h" -#include "xla/service/cpu/runtime/conv_impl.h" +#include "xla/backends/cpu/runtime/conv_impl.h" ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenSingleThreadedConv3DF32( diff --git a/xla/service/cpu/thunk_emitter.cc b/xla/service/cpu/thunk_emitter.cc index 49ed9d3d0b343..aadfabf9360ce 100644 --- a/xla/service/cpu/thunk_emitter.cc +++ b/xla/service/cpu/thunk_emitter.cc @@ -28,6 +28,29 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/all_gather_thunk.h" +#include "xla/backends/cpu/runtime/all_reduce_thunk.h" +#include "xla/backends/cpu/runtime/all_to_all_thunk.h" +#include "xla/backends/cpu/runtime/call_thunk.h" +#include "xla/backends/cpu/runtime/collective_permute_thunk.h" +#include "xla/backends/cpu/runtime/collective_thunk.h" +#include "xla/backends/cpu/runtime/conditional_thunk.h" +#include "xla/backends/cpu/runtime/convolution_thunk.h" +#include "xla/backends/cpu/runtime/copy_thunk.h" +#include "xla/backends/cpu/runtime/custom_call_thunk.h" +#include "xla/backends/cpu/runtime/dot_thunk.h" +#include "xla/backends/cpu/runtime/fft_thunk.h" +#include "xla/backends/cpu/runtime/infeed_thunk.h" +#include "xla/backends/cpu/runtime/kernel_thunk.h" +#include "xla/backends/cpu/runtime/logical_id_thunk.h" +#include "xla/backends/cpu/runtime/outfeed_thunk.h" +#include "xla/backends/cpu/runtime/reduce_scatter_thunk.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/rng_state_thunk.h" +#include "xla/backends/cpu/runtime/sort_thunk.h" +#include "xla/backends/cpu/runtime/thunk.h" +#include "xla/backends/cpu/runtime/topk_thunk.h" +#include "xla/backends/cpu/runtime/while_thunk.h" #include "xla/cpu_function_runtime.h" #include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" @@ -42,29 +65,6 @@ limitations under the License. #include "xla/service/cpu/dot_op_emitter.h" #include "xla/service/cpu/ir_emission_utils.h" #include "xla/service/cpu/ir_emitter2.h" -#include "xla/service/cpu/runtime/all_gather_thunk.h" -#include "xla/service/cpu/runtime/all_reduce_thunk.h" -#include "xla/service/cpu/runtime/all_to_all_thunk.h" -#include "xla/service/cpu/runtime/call_thunk.h" -#include "xla/service/cpu/runtime/collective_permute_thunk.h" -#include "xla/service/cpu/runtime/collective_thunk.h" -#include "xla/service/cpu/runtime/conditional_thunk.h" -#include "xla/service/cpu/runtime/convolution_thunk.h" -#include "xla/service/cpu/runtime/copy_thunk.h" -#include "xla/service/cpu/runtime/custom_call_thunk.h" -#include "xla/service/cpu/runtime/dot_thunk.h" -#include "xla/service/cpu/runtime/fft_thunk.h" -#include "xla/service/cpu/runtime/infeed_thunk.h" -#include "xla/service/cpu/runtime/kernel_thunk.h" -#include "xla/service/cpu/runtime/logical_id_thunk.h" -#include "xla/service/cpu/runtime/outfeed_thunk.h" -#include "xla/service/cpu/runtime/reduce_scatter_thunk.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/rng_state_thunk.h" -#include "xla/service/cpu/runtime/sort_thunk.h" -#include "xla/service/cpu/runtime/thunk.h" -#include "xla/service/cpu/runtime/topk_thunk.h" -#include "xla/service/cpu/runtime/while_thunk.h" #include "xla/service/cpu/target_machine_features.h" #include "xla/service/hlo_module_config.h" #include "xla/shape.h" diff --git a/xla/service/cpu/thunk_emitter.h b/xla/service/cpu/thunk_emitter.h index 6921f76e75179..ee907adb675c5 100644 --- a/xla/service/cpu/thunk_emitter.h +++ b/xla/service/cpu/thunk_emitter.h @@ -23,14 +23,14 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/types/span.h" +#include "xla/backends/cpu/runtime/resource_use.h" +#include "xla/backends/cpu/runtime/thunk.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/service/buffer_assignment.h" #include "xla/service/cpu/ir_emitter2.h" -#include "xla/service/cpu/runtime/resource_use.h" -#include "xla/service/cpu/runtime/thunk.h" #include "xla/service/cpu/target_machine_features.h" #include "xla/service/hlo_module_config.h" #include "xla/shape_util.h"