[xla:cpu] NFC: Move XLA:CPU runtime to xla/backends/cpu component

PiperOrigin-RevId: 663069847
openxla · Aug 14, 2024 · 5b5f580 · 5b5f580
1 parent ac359c4
commit 5b5f580
Show file tree

Hide file tree

Showing 85 changed files with 294 additions and 293 deletions.
diff --git a/xla/service/cpu/runtime/BUILD → xla/backends/cpu/runtime/BUILD b/xla/service/cpu/runtime/BUILD → xla/backends/cpu/runtime/BUILD
diff --git a/xla/service/cpu/runtime/all_gather_thunk.cc → xla/backends/cpu/runtime/all_gather_thunk.cc b/xla/service/cpu/runtime/all_gather_thunk.cc → 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 <cstdint>
 #include <memory>
@@ -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 → xla/backends/cpu/runtime/all_gather_thunk.h b/xla/service/cpu/runtime/all_gather_thunk.h → 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 <memory>
 
 #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 → xla/backends/cpu/runtime/all_reduce_thunk.cc b/xla/service/cpu/runtime/all_reduce_thunk.cc → 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 <cstdint>
 #include <cstring>
@@ -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 → xla/backends/cpu/runtime/all_reduce_thunk.h b/xla/service/cpu/runtime/all_reduce_thunk.h → 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 <memory>
 
 #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 → xla/backends/cpu/runtime/all_to_all_thunk.cc b/xla/service/cpu/runtime/all_to_all_thunk.cc → 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 <memory>
 #include <utility>
@@ -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 → xla/backends/cpu/runtime/all_to_all_thunk.h b/xla/service/cpu/runtime/all_to_all_thunk.h → 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 <memory>
 
 #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 → ...backends/cpu/runtime/buffer_allocations.h b/xla/service/cpu/runtime/buffer_allocations.h → ...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 <cstddef>
 #include <cstdint>
@@ -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/...ce/cpu/runtime/buffer_allocations_test.cc → ...ds/cpu/runtime/buffer_allocations_test.cc b/...ce/cpu/runtime/buffer_allocations_test.cc → ...ds/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 <cstddef>
 #include <vector>

diff --git a/xla/service/cpu/runtime/call_thunk.cc → xla/backends/cpu/runtime/call_thunk.cc b/xla/service/cpu/runtime/call_thunk.cc → 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 <memory>
 #include <utility>
 
 #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 → xla/backends/cpu/runtime/call_thunk.h b/xla/service/cpu/runtime/call_thunk.h → 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 <memory>
 
 #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/...e/cpu/runtime/collective_permute_thunk.cc → ...s/cpu/runtime/collective_permute_thunk.cc b/...e/cpu/runtime/collective_permute_thunk.cc → ...s/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 <cstdint>
 #include <memory>
@@ -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/...ce/cpu/runtime/collective_permute_thunk.h → ...ds/cpu/runtime/collective_permute_thunk.h b/...ce/cpu/runtime/collective_permute_thunk.h → ...ds/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 <cstdint>
 #include <memory>
@@ -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 → xla/backends/cpu/runtime/collective_thunk.cc b/xla/service/cpu/runtime/collective_thunk.cc → 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 <cstddef>
 #include <cstdint>
@@ -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 → xla/backends/cpu/runtime/collective_thunk.h b/xla/service/cpu/runtime/collective_thunk.h → 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 <cstdint>
 #include <memory>
@@ -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 → ...backends/cpu/runtime/conditional_thunk.cc b/xla/service/cpu/runtime/conditional_thunk.cc → ...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 <cstdint>
 #include <memory>
@@ -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 → xla/backends/cpu/runtime/conditional_thunk.h b/xla/service/cpu/runtime/conditional_thunk.h → 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 <memory>
 #include <vector>
 
 #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/...ice/cpu/runtime/conditional_thunk_test.cc → ...nds/cpu/runtime/conditional_thunk_test.cc b/...ice/cpu/runtime/conditional_thunk_test.cc → ...nds/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 <cstdint>
 #include <memory>
 #include <utility>
 #include <vector>
 
+#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 → xla/backends/cpu/runtime/conv_impl.cc b/xla/service/cpu/runtime/conv_impl.cc → 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 {