/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include "utils.h"

#pragma once

namespace kineto_stress_test {

// The unique identifier used for NCCL communication
extern ncclUniqueId nccl_id;

// The communicator object used for communication
extern ncclComm_t nccl_communicator;

// NCCL buffers on the device
extern float* pBuffNCCLSend;
extern float* pBuffNCCLRecv;

struct stress_test_args {
    // Number of threads that run the stress test
    uint32_t num_workers {1};

    // Number of operations per stress test
    uint32_t num_operations {10000};

    // Can improve event density. Cuda streams per worker
    uint32_t num_cuda_streams {1};

    // Simulates cuda mallocs happening in the PT Cuda Cache Allocator
    double prob_cuda_malloc {0.001};

    // The min number of compute iterations in the cuda kernel. If high
    // this reduces event density.
    uint32_t min_iters_kernel {1};

    // The max number of compute iterations in the cuda kernel. If high
    // this reduces event density.
    uint32_t max_iters_kernel {5};

    // The probability that instead of a kernel call we do a memset on the
    // input buffers, using a magic value
    double memset_prob {0.05};

    // The min idle time between kernel launches in microseconds
    uint32_t min_idle_us {1};

    // The max idle time between kernel launches in microseconds
    uint32_t max_idle_us {2};

    // If true, we randomly sleep a number of microseconds between kernel
    // launches.
    bool simulate_host_time {false};

    // If non-zero, we allocate UVM memory and use it
    bool use_uvm_buffers {false};

    // Size of a single buffer in FP32 elements in UVM
    uint64_t uvm_len {0};

    // If set true, the UVM allocation and initialization will be done in parallel
    // with cache allocation (e.g. cudaHostAlloc)
    bool parallel_uvm_alloc {false};

    // The probability of running a kernel that uses UVM
    double uvm_kernel_prob {0.001};

    // If true we need to run the binary using MPI on multiple ranks
    bool is_multi_rank {false};

    // Number of parallel processes to be spawned via MPI
    int32_t num_ranks {1};

    // Use this variable to pin a process to a specific GPU index.
    // Do not modify!
    int32_t rank {0};

    // Size of the NCCL buffers which needs to be at least the size
    // of the largest tensor
    uint32_t sz_nccl_buff_KB {1024};

    // Number of iterations between NCCL sync calls
    uint32_t num_iters_nccl_sync {100};

    // If true, we pre-allocate CUDA streams and reuse them throughout
    // the experiment
    bool pre_alloc_streams {false};

    // If true, h2d and d2h transfers would be scheduled on their own
    // CUDA stream
    bool use_memcpy_stream {false};

    // If true, kernels using UVM would be scheduled on their own
    // CUDA stream
    bool use_uvm_stream {false};

    // If true, we use cudaGetMemInfo throughout the stress test to
    // measure peak memory usage
    bool monitor_mem_usage {false};

    // Number of microseconds for trace collection. If 0 the trace is
    // not collected
    uint32_t trace_length_us {1000000};

    // Size of the CUPTI activity buffer in MB. If it's 0, we don't
    // explicitly set a value
    uint32_t cupti_buffer_mb {0};

    /* VARIABLES */

    // The CUDA streams vector
    cudaStream_t *compute_streams {nullptr};

    // The explicit memcpy stream
    cudaStream_t *memcpy_streams {nullptr};

    // The explicit UVM stream
    cudaStream_t *uvm_streams {nullptr};

    // UVM buffers
    float* uvm_a {nullptr};
    float* uvm_b {nullptr};
};

// We are using this to reduce the number of code lines
struct lcg_kernel_input {
  float const* __restrict__ d_a;
  float const* __restrict__ d_b;
  float* __restrict__ d_c;
  int len;
  float const* __restrict__ uvm_a;
  float const* __restrict__ uvm_b;
  uint64_t uvm_len;
  int iters;
};

// Use this function to vary the kernel name at runtime
void call_compute_kernel(
  uint32_t thread_blocks,
  uint32_t threads_per_block,
  uint32_t shmem_sz,
  cudaStream_t stream,
  lcg_kernel_input kernel_args,
  uint32_t op_id
);

void run_stress_test(
    uint32_t thread_id,
    uint32_t num_workers,
    stress_test_args test_args);

} // namespace kineto_stress_test
