nvshmem_collective_thunk.cc
imports tsl/platform/casts.himports xla/tsl/platform/errors.h
SHA-256bfcad9ce76966b0882792b800c4890541c70a268acc1595fb8bccf0acaf77d19
MaleculeMd(In₄)
Evidence
1/* Copyright 2025 The OpenXLA Authors.
2Licensed under the Apache License, Version 2.0 (the "License");
3you may not use this file except in compliance with the License.
4You may obtain a copy of the License at
5 http://www.apache.org/licenses/LICENSE-2.0
6Unless required by applicable law or agreed to in writing, software
7distributed under the License is d …
8:42… IND, either express or implied.
9See the License for the specific language governing permissions and
10limitations under the License.
11==============================================================================*/
12
13#include "xla/backends/gpu/runtime/nvshmem_collective_thunk.h"
14
15#include <cstdint>
16#include <memory>
17#include <optional>
18#include <utility>
19#include <vector>
20
21#include "absl/base/casts.h"
22#include "absl/container/flat_hash_map.h"
23#include "absl/container/flat_hash_set.h"
24#include "absl/log/log.h"
25#include "absl/status/status.h"
26#include "absl/status/statusor.h"
27#include "absl/strings/str_format.h"
28#include "absl/synchronization/mutex.h"
29#include "xla/backends/gpu/collectives/gpu_clique_key.h"
30#include "xla/backends/gpu/collectives/gpu_collectives.h"
31#include "xla/backends/gpu/runtime/collective_clique_requests.h"
32#include "xla/backends/gpu/runtime/collective_execution.h"
33#include "xla/backends/gpu/runtime/collective_thunk.h"
34#include "xla/backends/gpu/runtime/nvshmem_collective_thunk.pb.h"
35#include "xla/backends/gpu/runtime/thunk.h"
36#include "xla/backends/gpu/runtime/thunk.pb.h"
37#include "xla/backends/gpu/runtime/thunk_kind.pb.h"
38#include "xla/core/collectives/collectives.h"
39#include "xla/core/collectives/collectives_registry.h"
40#include "xla/primitive_util.h"
41#include "xla/runtime/device_id.h"
42#include "xla/service/collective_ops_utils.h"
43#include "xla/shape.h"
44#include "xla/status_macros.h"
45#include "xla/stream_executor/event.h"
46#include "xla/stream_executor/stream.h"
47#include "xla/stream_executor/stream_executor.h"
48#include "xla/tsl/platform/errors.h"
49#include "xla/tsl/platform/statusor.h"
50#include "xla/xla_data.pb.h"
51#include "tsl/platform/casts.h"
52#include "xla/tsl/platform/status_macros.h"
53
54namespace xla {
55namespace gpu {
242:7… turn std::nullopt;
243 }
244 // We rely on the fact that the pointer to async_events_ is unique.
245 return absl::bit_cast<AsyncEventsUniqueId>(async_even …