/* Copyright 2023 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_RUNTIME_COMMAND_BUFFER_THUNK_H_
#define XLA_BACKENDS_GPU_RUNTIME_COMMAND_BUFFER_THUNK_H_

#include <cstdint>
#include <memory>
#include <string>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/functional/function_ref.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/synchronization/mutex.h"
#include "xla/backends/gpu/runtime/command_executor.h"
#include "xla/backends/gpu/runtime/command_state.h"
#include "xla/backends/gpu/runtime/sequential_thunk.h"
#include "xla/backends/gpu/runtime/thunk.h"
#include "xla/service/buffer_assignment.h"
#include "xla/stream_executor/command_buffer.h"
#include "xla/stream_executor/device_address.h"
#include "xla/stream_executor/stream_executor.h"

namespace xla::gpu {

class CommandBufferThunk : public Thunk {
 public:
  CommandBufferThunk(CommandBufferCmdExecutor commands, ThunkInfo thunk_info,
                     std::unique_ptr<SequentialThunk> thunks = nullptr,
                     bool enable_command_buffers_during_profiling = false);

  const std::unique_ptr<SequentialThunk>& thunks() const { return thunks_; }

  absl::Status Prepare(const PrepareParams& params) override;
  absl::Status Initialize(const InitializeParams& params) override;
  absl::Status ExecuteOnStream(const ExecuteParams& params) override;

  // Return the allocation address that was lazilly allocated inside command
  // buffer. This API is required when the buffers are allocated inside command
  // buffer but will be consumed by non-command buffer operations.
  absl::StatusOr<se::DeviceAddressBase> GetCommandBufferAllocationAddress(
      const ExecuteParams& params, int64_t index);

  void ForAllThunks(absl::FunctionRef<void(const Thunk*)> fn) const override;
  void ForAllThunksMutable(absl::FunctionRef<void(Thunk*)> fn) override;

  std::string ToString(int indent) const override;

 private:
  // Command buffer instantiated on a `se::StreamExecutor` instance, and
  // auxiliary state required for efficient command buffer updates.
  struct ExecutorCommandBuffer {
    explicit ExecutorCommandBuffer(
        std::unique_ptr<se::CommandBuffer> command_buffer);

    // Updates recorded buffer allocation for the given `commands` using the
    // buffer allocations passed in `params`. Returns buffer allocations that
    // changed since the last update. Returned buffer allocations are sorted by
    // the buffer allocation index.
    std::vector<BufferAllocation::Index> UpdateBufferAllocations(
        const CommandBufferCmdExecutor& commands,
        const Thunk::ExecuteParams& params)
        ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex);

    // se::CommandBuffer is not thread safe, and we guard it with a mutex to
    // guarantee that we do not mutate it concurrently.
    absl::Mutex mutex;
    std::unique_ptr<se::CommandBuffer> command_buffer ABSL_GUARDED_BY(mutex);

    // A manager for an external state attached by commands in a command
    // sequence to a command buffer.
    CommandStateManager state ABSL_GUARDED_BY(mutex);

    // Mapping from buffer allocation index to the device memory passed at
    // that index to the last call of `commands_.Record(...)` for
    // `command_buffer`. We can just use a vector instead of map because
    // `BufferAllocation::Index` is a unique identifier assigned
    // contiguously and thus can be used as array index.
    //
    // If no device memory addresses changed from a previous call to
    // `Record`, we can skip command buffer update and simply submit it for
    // execution on a stream. All other pieces of information (like thread
    // and block sizes) captured by commands at construction time and do not
    // change.
    std::vector<se::DeviceAddressBase> recorded_allocs ABSL_GUARDED_BY(mutex);

    // Number of command buffer executions since last update.
    int64_t num_executions ABSL_GUARDED_BY(mutex) = 0;

    // For GPU backend, NCCL may call cuda-graph un-supported host side API
    // during graph capturing (e.g. cuCtxEnablePeerAccess), this will break XLA
    // cuda graph run. To work around the issue, this PR introduces a warm up
    // iteration for command buffer thunk, during warm up iteration, command
    // buffer thunk are executed through normal thunks. The warm up iteration
    // will do the proper NCCL setup, so later iterations running through
    // command buffer does not need to call NCCL setup APIs.
    bool warmup_done ABSL_GUARDED_BY(mutex) = false;
  };

  // Command buffer thunk owns commands buffers instantiated on all executors.
  struct State {
    absl::Mutex mutex;
    absl::flat_hash_map<se::StreamExecutor*,
                        std::shared_ptr<ExecutorCommandBuffer>>
        command_buffers ABSL_GUARDED_BY(mutex);
  };

  // Returns a command buffer instantiated for `executor` or creates new one.
  absl::StatusOr<std::shared_ptr<ExecutorCommandBuffer>>
  GetOrCreateCommandBuffer(se::StreamExecutor* executor);

  // Each individual command buffer allocates state on device (CUDA graph) and
  // it adds up pretty quickly. To prevent OOM errors we proactively evict
  // command buffers from device by clearing command buffer thunk state. We use
  // global state to track all command buffer thunks in a process and coordinate
  // command buffer eviction.
  struct GlobalState;

  // Returns a global state of tracked command buffers thunks.
  static GlobalState* GetGlobalState();

  // Adds command buffer thunk state for tracking.
  static void TrackCommandBuffers(std::weak_ptr<State> state);

  // Evicts all previously instantiated command buffers.
  static void EvictCommandBuffers();

  // Commands executor that initializes command buffers on each stream executor.
  CommandBufferCmdExecutor commands_;

  // Thunk sequence that executes the same commands as in `commands_` but using
  // thunk mechanism. We use it as a fallback mechanism to work around CUPTI
  // bugs that lead to memory corruption when CUPTI traces CUDA graph execution.
  std::unique_ptr<SequentialThunk> thunks_;

  // When true, allows command buffers to be used while profiling active.
  // TODO(b/355487968): Remove this option when validation complete.
  bool enable_command_buffers_during_profiling_;

  // Command buffer thunk state allocated in heap to allow global (per-process)
  // management of instantiated command buffers.
  std::shared_ptr<State> state_;
};

}  // namespace xla::gpu

#endif  // XLA_BACKENDS_GPU_RUNTIME_COMMAND_BUFFER_THUNK_H_