Skip to content

Commit b5cc3bb

Browse files
ezhulenevGoogle-ML-Automation
ezhulenev
authored and
Google-ML-Automation
committed
[xla:gpu] Move XLA:GPU runtime to xla/backends/gpu
PiperOrigin-RevId: 715158015
1 parent efa756b commit b5cc3bb

File tree

110 files changed

+450
-449
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

110 files changed

+450
-449
lines changed

‎xla/backends/gpu/codegen/emitters/BUILD‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ cc_library(
5050
"//xla:xla_data_proto_cc",
5151
"//xla/backends/gpu/codegen/ir:xla_gpu",
5252
"//xla/backends/gpu/codegen/transforms:passes",
53+
"//xla/backends/gpu/runtime:kernel_thunk",
5354
"//xla/codegen/emitters:computation_partitioner",
5455
"//xla/codegen/emitters:elemental_hlo_to_mlir",
5556
"//xla/codegen/emitters:type_util",
@@ -69,7 +70,6 @@ cc_library(
6970
"//xla/service/gpu:launch_dimensions",
7071
"//xla/service/gpu:target_util",
7172
"//xla/service/gpu/fusions:fusion_emitter",
72-
"//xla/service/gpu/runtime:kernel_thunk",
7373
"//xla/service/llvm_ir:llvm_util",
7474
"//xla/stream_executor:device_description",
7575
"//xla/tsl/framework/mlir:status_scoped_diagnostic_handler",

‎xla/backends/gpu/codegen/emitters/emitter_base.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ limitations under the License.
7878
#include "mlir/Transforms/Passes.h"
7979
#include "xla/backends/gpu/codegen/ir/xla_gpu_ops.h"
8080
#include "xla/backends/gpu/codegen/transforms/passes.h"
81+
#include "xla/backends/gpu/runtime/kernel_thunk.h"
8182
#include "xla/codegen/emitters/computation_partitioner.h"
8283
#include "xla/codegen/emitters/elemental_hlo_to_mlir.h"
8384
#include "xla/codegen/emitters/type_util.h"
@@ -97,7 +98,6 @@ limitations under the License.
9798
#include "xla/service/gpu/kernel_arguments.h"
9899
#include "xla/service/gpu/kernel_reuse_cache.h"
99100
#include "xla/service/gpu/launch_dimensions.h"
100-
#include "xla/service/gpu/runtime/kernel_thunk.h"
101101
#include "xla/service/gpu/target_util.h"
102102
#include "xla/service/llvm_ir/llvm_util.h"
103103
#include "xla/shape.h"

‎xla/service/gpu/runtime/BUILD‎ renamed to ‎xla/backends/gpu/runtime/BUILD‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -280,11 +280,11 @@ cc_library(
280280
"//xla:types",
281281
"//xla:util",
282282
"//xla:xla_data_proto_cc",
283+
"//xla/backends/gpu/runtime:thunk",
283284
"//xla/hlo/ir:hlo",
284285
"//xla/service:buffer_assignment",
285286
"//xla/service/gpu:buffer_allocations",
286287
"//xla/service/gpu:make_batch_pointers",
287-
"//xla/service/gpu/runtime:thunk",
288288
"//xla/stream_executor:blas",
289289
"//xla/stream_executor:device_memory",
290290
"//xla/stream_executor:gpu_solver_context",
@@ -476,7 +476,7 @@ cc_library(
476476
"@com_google_absl//absl/strings",
477477
"//xla/service:buffer_assignment",
478478
"//xla/service/gpu:buffer_allocations",
479-
"//xla/service/gpu/runtime:thunk",
479+
"//xla/backends/gpu/runtime:thunk",
480480
"//xla/stream_executor/gpu:gpu_stream",
481481
"//xla/stream_executor/gpu:gpu_types_header",
482482
"//xla:shape_util",
@@ -1168,11 +1168,11 @@ cc_library(
11681168
"//xla:types",
11691169
"//xla:util",
11701170
"//xla:xla_data_proto_cc",
1171+
"//xla/backends/gpu/runtime:thunk",
11711172
"//xla/hlo/ir:hlo",
11721173
"//xla/service:buffer_assignment",
11731174
"//xla/service/gpu:buffer_allocations",
11741175
"//xla/service/gpu:make_batch_pointers",
1175-
"//xla/service/gpu/runtime:thunk",
11761176
"//xla/stream_executor:blas",
11771177
"//xla/stream_executor:device_memory",
11781178
"//xla/stream_executor:stream",

‎xla/service/gpu/runtime/annotation.cc‎ renamed to ‎xla/backends/gpu/runtime/annotation.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#include "xla/service/gpu/runtime/annotation.h"
16+
#include "xla/backends/gpu/runtime/annotation.h"
1717

1818
#include <algorithm>
1919
#include <cstddef>

‎xla/service/gpu/runtime/annotation.h‎ renamed to ‎xla/backends/gpu/runtime/annotation.h‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#ifndef XLA_SERVICE_GPU_RUNTIME_ANNOTATION_H_
17-
#define XLA_SERVICE_GPU_RUNTIME_ANNOTATION_H_
16+
#ifndef XLA_BACKENDS_GPU_RUNTIME_ANNOTATION_H_
17+
#define XLA_BACKENDS_GPU_RUNTIME_ANNOTATION_H_
1818

1919
#include <cstdint>
2020
#include <optional>
@@ -107,4 +107,4 @@ std::optional<tsl::profiler::ScopedAnnotation> GetKernelAnnotation(
107107

108108
} // namespace xla::gpu
109109

110-
#endif // XLA_SERVICE_GPU_RUNTIME_ANNOTATION_H_
110+
#endif // XLA_BACKENDS_GPU_RUNTIME_ANNOTATION_H_

‎xla/service/gpu/runtime/cholesky_thunk.cc‎ renamed to ‎xla/backends/gpu/runtime/cholesky_thunk.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#include "xla/service/gpu/runtime/cholesky_thunk.h"
16+
#include "xla/backends/gpu/runtime/cholesky_thunk.h"
1717

1818
#include <complex>
1919
#include <cstdint>
@@ -23,9 +23,9 @@ limitations under the License.
2323
#include "absl/functional/any_invocable.h"
2424
#include "absl/status/status.h"
2525
#include "absl/strings/str_format.h"
26+
#include "xla/backends/gpu/runtime/thunk.h"
2627
#include "xla/service/buffer_assignment.h"
2728
#include "xla/service/gpu/make_batch_pointers.h"
28-
#include "xla/service/gpu/runtime/thunk.h"
2929
#include "xla/stream_executor/blas.h"
3030
#include "xla/stream_executor/device_memory.h"
3131
#include "xla/stream_executor/gpu_solver_context.h"

‎xla/service/gpu/runtime/cholesky_thunk.h‎ renamed to ‎xla/backends/gpu/runtime/cholesky_thunk.h‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,17 +13,17 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#ifndef XLA_SERVICE_GPU_RUNTIME_CHOLESKY_THUNK_H_
17-
#define XLA_SERVICE_GPU_RUNTIME_CHOLESKY_THUNK_H_
16+
#ifndef XLA_BACKENDS_GPU_RUNTIME_CHOLESKY_THUNK_H_
17+
#define XLA_BACKENDS_GPU_RUNTIME_CHOLESKY_THUNK_H_
1818

1919
#include <cstdint>
2020
#include <memory>
2121

2222
#include "absl/functional/any_invocable.h"
2323
#include "absl/status/status.h"
2424
#include "absl/status/statusor.h"
25+
#include "xla/backends/gpu/runtime/thunk.h"
2526
#include "xla/service/buffer_assignment.h"
26-
#include "xla/service/gpu/runtime/thunk.h"
2727
#include "xla/stream_executor/blas.h"
2828
#include "xla/stream_executor/device_memory.h"
2929
#include "xla/stream_executor/gpu_solver_context.h"
@@ -84,4 +84,4 @@ struct CholeskyParams {
8484
} // namespace gpu
8585
} // namespace xla
8686

87-
#endif // XLA_SERVICE_GPU_RUNTIME_CHOLESKY_THUNK_H_
87+
#endif // XLA_BACKENDS_GPU_RUNTIME_CHOLESKY_THUNK_H_

‎xla/service/gpu/runtime/command_buffer_cmd.cc‎ renamed to ‎xla/backends/gpu/runtime/command_buffer_cmd.cc‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#include "xla/service/gpu/runtime/command_buffer_cmd.h"
16+
#include "xla/backends/gpu/runtime/command_buffer_cmd.h"
1717

1818
#include <algorithm>
1919
#include <cassert>
@@ -42,6 +42,14 @@ limitations under the License.
4242
#include "llvm/ADT/STLExtras.h"
4343
#include "xla/backends/gpu/collectives/gpu_clique_key.h"
4444
#include "xla/backends/gpu/collectives/gpu_collectives.h"
45+
#include "xla/backends/gpu/runtime/annotation.h"
46+
#include "xla/backends/gpu/runtime/dynamic_slice_thunk.h"
47+
#include "xla/backends/gpu/runtime/nccl_all_gather_thunk.h"
48+
#include "xla/backends/gpu/runtime/nccl_all_reduce_thunk.h"
49+
#include "xla/backends/gpu/runtime/nccl_all_to_all_thunk.h"
50+
#include "xla/backends/gpu/runtime/nccl_collective_broadcast_thunk.h"
51+
#include "xla/backends/gpu/runtime/nccl_collective_thunk.h"
52+
#include "xla/backends/gpu/runtime/thunk.h"
4553
#include "xla/debug_options_flags.h"
4654
#include "xla/executable_run_options.h"
4755
#include "xla/ffi/call_frame.h"
@@ -56,14 +64,6 @@ limitations under the License.
5664
#include "xla/service/gpu/kernels/custom_kernel.h"
5765
#include "xla/service/gpu/launch_dimensions.h"
5866
#include "xla/service/gpu/matmul_utils.h"
59-
#include "xla/service/gpu/runtime/annotation.h"
60-
#include "xla/service/gpu/runtime/dynamic_slice_thunk.h"
61-
#include "xla/service/gpu/runtime/nccl_all_gather_thunk.h"
62-
#include "xla/service/gpu/runtime/nccl_all_reduce_thunk.h"
63-
#include "xla/service/gpu/runtime/nccl_all_to_all_thunk.h"
64-
#include "xla/service/gpu/runtime/nccl_collective_broadcast_thunk.h"
65-
#include "xla/service/gpu/runtime/nccl_collective_thunk.h"
66-
#include "xla/service/gpu/runtime/thunk.h"
6767
#include "xla/service/gpu/stream_executor_util.h"
6868
#include "xla/shape_util.h"
6969
#include "xla/status_macros.h"

‎xla/service/gpu/runtime/command_buffer_cmd.h‎ renamed to ‎xla/backends/gpu/runtime/command_buffer_cmd.h‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#ifndef XLA_SERVICE_GPU_RUNTIME_COMMAND_BUFFER_CMD_H_
17-
#define XLA_SERVICE_GPU_RUNTIME_COMMAND_BUFFER_CMD_H_
16+
#ifndef XLA_BACKENDS_GPU_RUNTIME_COMMAND_BUFFER_CMD_H_
17+
#define XLA_BACKENDS_GPU_RUNTIME_COMMAND_BUFFER_CMD_H_
1818

1919
#include <algorithm>
2020
#include <cstddef>
@@ -37,6 +37,10 @@ limitations under the License.
3737
#include "absl/synchronization/mutex.h"
3838
#include "absl/types/span.h"
3939
#include "xla/backends/gpu/collectives/gpu_clique_key.h"
40+
#include "xla/backends/gpu/runtime/custom_call_thunk.h"
41+
#include "xla/backends/gpu/runtime/dynamic_slice_thunk.h"
42+
#include "xla/backends/gpu/runtime/nccl_collective_thunk.h"
43+
#include "xla/backends/gpu/runtime/thunk.h"
4044
#include "xla/ffi/api/c_api.h"
4145
#include "xla/hlo/ir/hlo_computation.h"
4246
#include "xla/runtime/buffer_use.h"
@@ -46,10 +50,6 @@ limitations under the License.
4650
#include "xla/service/gpu/kernels/custom_kernel.h"
4751
#include "xla/service/gpu/launch_dimensions.h"
4852
#include "xla/service/gpu/matmul_utils.h"
49-
#include "xla/service/gpu/runtime/custom_call_thunk.h"
50-
#include "xla/service/gpu/runtime/dynamic_slice_thunk.h"
51-
#include "xla/service/gpu/runtime/nccl_collective_thunk.h"
52-
#include "xla/service/gpu/runtime/thunk.h"
5353
#include "xla/shape.h"
5454
#include "xla/stream_executor/command_buffer.h"
5555
#include "xla/stream_executor/device_memory.h"
@@ -1189,4 +1189,4 @@ class DynamicSliceFusionCmd : public CommandBufferCmd {
11891189

11901190
} // namespace xla::gpu
11911191

1192-
#endif // XLA_SERVICE_GPU_RUNTIME_COMMAND_BUFFER_CMD_H_
1192+
#endif // XLA_BACKENDS_GPU_RUNTIME_COMMAND_BUFFER_CMD_H_

‎xla/service/gpu/runtime/command_buffer_cmd_emitter.cc‎ renamed to ‎xla/backends/gpu/runtime/command_buffer_cmd_emitter.cc‎

Lines changed: 19 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#include "xla/service/gpu/runtime/command_buffer_cmd_emitter.h"
16+
#include "xla/backends/gpu/runtime/command_buffer_cmd_emitter.h"
1717

1818
#include <memory>
1919
#include <optional>
@@ -23,25 +23,25 @@ limitations under the License.
2323
#include "absl/container/inlined_vector.h"
2424
#include "absl/status/status.h"
2525
#include "absl/status/statusor.h"
26+
#include "xla/backends/gpu/runtime/command_buffer_cmd.h"
27+
#include "xla/backends/gpu/runtime/conditional_thunk.h"
28+
#include "xla/backends/gpu/runtime/copy_thunk.h"
29+
#include "xla/backends/gpu/runtime/cudnn_thunk.h"
30+
#include "xla/backends/gpu/runtime/custom_call_thunk.h"
31+
#include "xla/backends/gpu/runtime/gemm_thunk.h"
32+
#include "xla/backends/gpu/runtime/gpublas_lt_matmul_thunk.h"
33+
#include "xla/backends/gpu/runtime/kernel_thunk.h"
34+
#include "xla/backends/gpu/runtime/memset_thunk.h"
35+
#include "xla/backends/gpu/runtime/nccl_all_gather_thunk.h"
36+
#include "xla/backends/gpu/runtime/nccl_all_reduce_thunk.h"
37+
#include "xla/backends/gpu/runtime/nccl_all_to_all_thunk.h"
38+
#include "xla/backends/gpu/runtime/nccl_collective_thunk.h"
39+
#include "xla/backends/gpu/runtime/replica_id_thunk.h"
40+
#include "xla/backends/gpu/runtime/sequential_thunk.h"
41+
#include "xla/backends/gpu/runtime/thunk.h"
42+
#include "xla/backends/gpu/runtime/wait_for_streams_thunk.h"
43+
#include "xla/backends/gpu/runtime/while_thunk.h"
2644
#include "xla/runtime/buffer_use.h"
27-
#include "xla/service/gpu/runtime/command_buffer_cmd.h"
28-
#include "xla/service/gpu/runtime/conditional_thunk.h"
29-
#include "xla/service/gpu/runtime/copy_thunk.h"
30-
#include "xla/service/gpu/runtime/cudnn_thunk.h"
31-
#include "xla/service/gpu/runtime/custom_call_thunk.h"
32-
#include "xla/service/gpu/runtime/gemm_thunk.h"
33-
#include "xla/service/gpu/runtime/gpublas_lt_matmul_thunk.h"
34-
#include "xla/service/gpu/runtime/kernel_thunk.h"
35-
#include "xla/service/gpu/runtime/memset_thunk.h"
36-
#include "xla/service/gpu/runtime/nccl_all_gather_thunk.h"
37-
#include "xla/service/gpu/runtime/nccl_all_reduce_thunk.h"
38-
#include "xla/service/gpu/runtime/nccl_all_to_all_thunk.h"
39-
#include "xla/service/gpu/runtime/nccl_collective_thunk.h"
40-
#include "xla/service/gpu/runtime/replica_id_thunk.h"
41-
#include "xla/service/gpu/runtime/sequential_thunk.h"
42-
#include "xla/service/gpu/runtime/thunk.h"
43-
#include "xla/service/gpu/runtime/wait_for_streams_thunk.h"
44-
#include "xla/service/gpu/runtime/while_thunk.h"
4545
#include "xla/util.h"
4646
#include "tsl/platform/errors.h"
4747
#include "tsl/platform/statusor.h"

0 commit comments

Comments
 (0)