/*
 * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 *
 * 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.
 */
#pragma once

#include <cuda_runtime.h>

#include <cstddef>
#include <cstdint>

#include "tensorrt_llm/common/config.h"

TRTLLM_NAMESPACE_BEGIN

namespace kernels {

struct HelixFieldInfo {
  uint8_t* dataPtr;
  int elementCount;  // Number of elements (e.g., kv_lora_rank for field 0, 1 for
                     // field 1)
  int elementSize;   // Size of each element in bytes (2 for half, 8 for float2)
  int stride;        // Stride between rows in bytes
};

struct HelixAllToAllParams {
  HelixFieldInfo sendFields[2];
  HelixFieldInfo recvFields[2];
  int entryCount;  // Number of entries per peer rank to process
  uint64_t* workspace;
  size_t workspaceStrideInU64;
  int cpRank;
  int cpSize;
  int channelCount;  // use 0 to auto-compute
  int maxChannelCount;
};

// ============================================================================
// Workspace Management Functions
// ============================================================================

/**
 * Compute number of channels for communication based on cpSize.
 *
 * @param cpSize Number of context parallel ranks
 * @param smCount Number of SMs available (0 = auto-detect)
 * @return Number of channels to use
 */
int computeHelixMaxChannelCount(int cpSize, int smCount = 0);

/**
 * Compute the workspace size required per rank for the all-to-all operation.
 *
 * @param cpSize Number of context parallel ranks
 * @return Size in bytes
 */
size_t computeHelixWorkspaceSizePerRank(int cpSize);

/**
 * Initialize workspace memory for a given rank.
 * Should be called once during setup.
 *
 * @param workspace Pointer to workspace memory (per-rank view)
 * @param cpSize Number of context parallel ranks
 * @param stream CUDA stream for asynchronous operations
 */
void initializeHelixWorkspace(uint64_t* workspace, int cpSize, cudaStream_t stream);

/**
 * Launch the helix all-to-all kernel.
 *
 * @param params Kernel parameters including field info and workspace
 * @param allowVariableField1 Whether to allow variable field 1
 * @param stream CUDA stream for kernel launch
 */
void launchHelixAllToAll(HelixAllToAllParams const& params, bool allowVariableField1,
                         bool enablePdl, cudaStream_t stream);

}  // namespace kernels

TRTLLM_NAMESPACE_END
