PR #36190: [xla:gpu] Extract Command into a separate target

Imported from GitHub PR #36190

Keep splitting `command_buffer_cmd` into more manageable targets. This is an NFC change, simply moving code into separate build target and updating dependencies.

Add `AsyncStartCommand` and `AsyncDoneCommand` to be able to break circular dependency between Command and CommandExecutor (coming next PR).
Copybara import of the project:

--
65c6cb5 by Eugene Zhulenev <[email protected]>:

[xla:gpu] Extract CommandState into a separate target + improve documentation

--
d914208 by Eugene Zhulenev <[email protected]>:

[xla:gpu] CommandExecutor: Correctly split RecordCreate from RecordUpdate and cleanup public facing Record API

--
7cc56e5 by Eugene Zhulenev <[email protected]>:

fix BUILD

--
c76bbff by Eugene Zhulenev <[email protected]>:

[xla:gpu] Extract Command into a separate target

Merging this change closes #36190

COPYBARA_INTEGRATE_REVIEW=#36190 from ezhulenev:command-split-3 c76bbff
PiperOrigin-RevId: 855161008

Author: ezhulenev

xla/backends/gpu/profiler/BUILD

@@ -73,6 +73,7 @@ xla_test(
     ],
     deps = [
         ":kernel_name_tracer",
+        "//xla/backends/gpu/runtime:command",
         "//xla/backends/gpu/runtime:command_buffer_cmd",
         "//xla/backends/gpu/runtime:command_buffer_thunk",
         "//xla/backends/gpu/runtime:thunk",

xla/backends/gpu/profiler/kernel_name_tracer_test.cc

@@ -29,6 +29,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/ascii.h"
 #include "absl/strings/string_view.h"
+#include "xla/backends/gpu/runtime/command.h"
 #include "xla/backends/gpu/runtime/command_buffer_cmd.h"
 #include "xla/backends/gpu/runtime/command_buffer_thunk.h"
 #include "xla/backends/gpu/runtime/thunk.h"
@@ -155,7 +156,7 @@ void LaunchCommandBufferThunk(stream_executor::StreamExecutor* executor,
                       BufferUse::MemoryAccess::kWrite};
 
   // Prepare commands sequence for constructing command buffer.
-  CommandBufferCmdSequence commands;
+  CommandSequence commands;
   commands.Emplace<LaunchCmd>("AddI32", args, args_access,
                               LaunchDimensions(1, kLength),
                               /*shmem_bytes=*/0);

xla/backends/gpu/runtime/BUILD

@@ -80,6 +80,19 @@ xla_cc_test(
     ],
 )
 
+cc_library(
+    name = "command",
+    srcs = ["command.cc"],
+    hdrs = ["command.h"],
+    deps = [
+        ":command_state",
+        ":thunk",
+        "//xla/runtime:buffer_use",
+        "//xla/runtime:resource_use",
+        "//xla/stream_executor:command_buffer",
+    ],
+)
+
 cc_library(
     name = "command_buffer_cmd",
     srcs = ["command_buffer_cmd.cc"],
@@ -93,6 +106,7 @@ cc_library(
         ":collective_execution",
         ":collective_permute_thunk",
         ":collective_thunk",
+        ":command",
         ":command_state",
         ":copy_thunk",
         ":custom_call_thunk",
@@ -164,7 +178,6 @@ cc_library(
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
-        "@llvm-project//llvm:Support",
         "@tsl//tsl/profiler/lib:scoped_annotation",
     ],
 )
@@ -174,6 +187,7 @@ xla_test(
     srcs = ["command_buffer_cmd_test.cc"],
     backends = ["gpu"],
     deps = [
+        ":command",
         ":command_buffer_cmd",
         ":command_state",
         ":copy_thunk",
@@ -221,6 +235,7 @@ cc_library(
         ":collective_broadcast_thunk",
         ":collective_permute_thunk",
         ":collective_thunk",
+        ":command",
         ":command_buffer_cmd",
         ":command_state",
         ":conditional_thunk",
@@ -354,6 +369,7 @@ cc_library(
     srcs = ["command_buffer_thunk.cc"],
     hdrs = ["command_buffer_thunk.h"],
     deps = [
+        ":command",
         ":command_buffer_cmd",
         ":command_state",
         ":sequential_thunk",  # build_cleaner: keep
@@ -385,6 +401,7 @@ xla_test(
     srcs = ["command_buffer_thunk_test.cc"],
     backends = ["gpu"],
     deps = [
+        ":command",
         ":command_buffer_cmd",
         ":command_buffer_thunk",
         ":dynamic_slice_thunk",
@@ -436,6 +453,7 @@ xla_test(
     backends = ["gpu"],
     tags = ["cuda-only"],
     deps = [
+        ":command",
         ":command_buffer_cmd",
         ":command_buffer_thunk",
         ":thunk",

xla/backends/gpu/runtime/command.cc

@@ -0,0 +1,46 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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/backends/gpu/runtime/command.h"
+
+#include <string>
+
+namespace xla::gpu {
+
+std::string CommandTypeString(CommandType type) {
+  switch (type) {
+#define CASE_CMD_STRING(enum_name, cmd_name, ...) \
+  case CommandType::enum_name:                    \
+    return cmd_name;
+    XLA_GPU_COMMAND_LIST(CASE_CMD_STRING)
+#undef CASE_CMD_STRING
+  }
+}
+
+bool IsCollectiveCommand(CommandType type) {
+  switch (type) {
+    case CommandType::kAllGatherCmd:
+    case CommandType::kAllReduceCmd:
+    case CommandType::kAllToAllCmd:
+    case CommandType::kReduceScatterCmd:
+    case CommandType::kCollectiveBroadcastCmd:
+    case CommandType::kCollectivePermuteCmd:
+      return true;
+    default:
+      return false;
+  }
+}
+
+}  // namespace xla::gpu