// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::sync::atomic::{AtomicI32, AtomicUsize, Ordering};

use cudax::nccl;

static GLOBAL_RANK: AtomicI32 = AtomicI32::new(-1);
static RANK_COUNT: AtomicUsize = AtomicUsize::new(0);

static NCCL_COMM: AtomicUsize = AtomicUsize::new(0);

pub fn set_global_rank(rank: i32) {
    GLOBAL_RANK.store(rank, Ordering::Release);
}

pub fn get_global_rank() -> i32 {
    GLOBAL_RANK.load(Ordering::Acquire)
}

pub fn set_rank_count(count: usize) {
    RANK_COUNT.store(count, Ordering::Release);
}

pub fn get_rank_count() -> usize {
    RANK_COUNT.load(Ordering::Acquire)
}

pub fn set_comm(comm: nccl::ncclComm_t) {
    NCCL_COMM.store(comm, Ordering::Release);
}

pub fn get_comm() -> nccl::ncclComm_t {
    NCCL_COMM.load(Ordering::Acquire) as nccl::ncclComm_t
}