/* Copyright 2024 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.
==============================================================================*/

#ifndef XLA_BACKENDS_GPU_COLLECTIVES_GPU_COLLECTIVES_H_
#define XLA_BACKENDS_GPU_COLLECTIVES_GPU_COLLECTIVES_H_

#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <optional>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "xla/backends/gpu/collectives/cancellation_token.h"
#include "xla/core/collectives/clique_id.h"
#include "xla/core/collectives/clique_key.h"
#include "xla/core/collectives/collectives.h"
#include "xla/core/collectives/communicator.h"
#include "xla/core/collectives/rank_id.h"
#include "xla/pjrt/distributed/key_value_store_interface.h"
#include "xla/runtime/device_id.h"
#include "xla/runtime/process_id.h"
#include "xla/stream_executor/device_address.h"
#include "xla/stream_executor/stream.h"
#include "xla/stream_executor/stream_executor.h"
#include "xla/xla_data.pb.h"

namespace xla::gpu {

// XLA:GPU extension of the Collectives interface with GPU-specific APIs.
class GpuCollectives : public Collectives {
 public:
  // Returns the default collectives implementation for the given platform.
  static GpuCollectives* Default(absl::string_view platform_name);

  // A callback to get a unique clique ids.
  using CliqueIdCallback =  // NOLINT
      std::function<absl::StatusOr<CliqueIds>(const CliqueKey&)>;

  // Topology describes how exactly the current process fits into the collection
  // of distributed processes, and based on this information the underlying
  // collective communication library can set up communication channels between
  // all devices.
  struct Topology {
    ProcessId process_id;
    size_t num_processes;
    size_t device_count_per_process;
    std::shared_ptr<KeyValueStoreInterface> kv_store;
    absl::flat_hash_map<GlobalDeviceId, ProcessId> device_to_process;
  };

  // Initializes the collectives backend with the provided topology information
  // and returns a callback that will generate unique ids for the cliques if
  // topology spans multiple processes and clique id generation requires
  // multi-process coordination. For local toplogies returns a nullptr callback.
  virtual absl::StatusOr<CliqueIdCallback> InitializeTopology(
      const Topology& topology) = 0;

  // GPU collectives device is just a wrapper around the StreamExecutor.
  class Device : public Collectives::Device {
   public:
    explicit Device(stream_executor::StreamExecutor* stream_executor);
    stream_executor::StreamExecutor* stream_executor() const;

   private:
    stream_executor::StreamExecutor* stream_executor_;
  };

  // GPU collectives executor is just a wrapper around the Stream.
  class Executor : public Communicator::Executor {
   public:
    explicit Executor(stream_executor::Stream* stream);
    stream_executor::Stream* stream() const;

   private:
    stream_executor::Stream* stream_;
  };

  static Executor On(se::Stream& stream) { return Executor(&stream); }

  // GPU communicator configuration.
  //
  // For NCCL backend see configuration options documentation at:
  // https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig
  struct Config : public Collectives::Config {
    bool split_share = false;
    int64_t max_nchannels = 0;

    // There are two types of NCCL communicators: blocking and non-blocking.
    // When a collective operation is called on a blocking communicator, the
    // communicator blocks until the operation has been scheduled on the GPU. A
    // non-blocking communicator, on the other hand, returns immediately.
    //
    // https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#c.blocking
    bool blocking_communicators = true;

    // If true, Communicator methods (e.g., AllReduce) return AsyncValueRefs
    // that are filled in asynchronously. If false, Communicator methods return
    // AsyncValueRefs that are already filled in.
    //
    // If blocking_communicators is false, then async_execution must be true.
    bool async_execution = false;
  };

  // A cancelable version of Collectives::CreateCommunicators.
  virtual absl::StatusOr<std::vector<std::unique_ptr<Communicator>>>
  CreateCommunicatorsWithCancel(const CliqueKey& clique_key,
                                const std::optional<CliqueIds>& clique_ids,
                                absl::Span<const DeviceRank> ranks,
                                const Collectives::Config& config,
                                std::shared_ptr<CancellationToken> cancel) {
    // By default, we ignore cancel.
    return CreateCommunicators(clique_key, clique_ids, ranks, config);
  }

  // A cancelable version of Collectives::SplitCommunicators.
  virtual absl::StatusOr<std::vector<std::unique_ptr<Communicator>>>
  SplitCommunicatorsWithCancel(absl::Span<const Communicator* const> comms,
                               int32_t color, absl::Span<const RankId> keys,
                               const Collectives::Config& config,
                               absl::Span<const DeviceRank> ranks,
                               std::shared_ptr<CancellationToken> cancel) {
    // By default, we ignore cancel.
    return SplitCommunicators(comms, color, keys, config, ranks);
  }

  // Returns true if GPU collectives are implemented.
  virtual bool IsImplemented() const = 0;

  // Returns true if GPU collectives support device-initiated communication.
  virtual bool SupportsDeviceComm() const { return false; }

  // Returns a slice of device memory `buff` containing `count` values of data
  // type `dtype` starting from `offset`.
  static stream_executor::DeviceAddressBase Slice(
      stream_executor::DeviceAddressBase buff, PrimitiveType dtype,
      size_t offset, size_t count);

  // TODO(b/410686553): Use smart wrapper instead of void*.
  virtual absl::StatusOr<void*> Allocate(uint64_t bytes) = 0;

  virtual absl::Status Deallocate(void* buffer) = 0;

  // Creates a single communicator.
  virtual absl::StatusOr<std::unique_ptr<Communicator>>
  CreateCommunicator() = 0;
};

}  // namespace xla::gpu

#endif  // XLA_BACKENDS_GPU_COLLECTIVES_GPU_COLLECTIVES_H_