asm_ops.h

#pragma once

#include <sys/mman.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include "check.h"

namespace sccl {
namespace asm_ops {

/*
标志名称描述与用途适用场景
__ATOMIC_RELAXED：最弱的内存顺序，无同步约束，仅保证原子性，适用于无数据依赖的场景
__ATOMIC_ACQUIRE：确保后续操作读取的共享数据可见，用于同步读取操作
__ATOMIC_RELEASE：确保当前操作对共享数据的修改对后续操作可见，用于同步
__ATOMIC_ACQ_REL：同时具备ACQUIRE和RELEASE语义，用于读写同步。确保在该原子操作之前的所有操作对其他线程可见，同时确保在该原子操作之后的所有操作对其他线程可见
__ATOMIC_SEQ_CST：顺序一致性约束，确保所有线程的操作按全局顺序执行
*/

/**
 * 以宽松内存序对引用计数进行原子加1操作
 * @param refs 指向引用计数的指针
 */
template <typename Int>
__host__ __device__ __forceinline__ void add_ref_count_increment_relaxed(Int* refs) {
    __atomic_fetch_add(refs, 1, __ATOMIC_RELAXED);
}

/**
 * 以顺序一致性内存顺序对引用计数进行原子加1操作
 * @param refs 指向引用计数的指针
 */
template <typename Int>
__host__ __device__ __forceinline__ void add_ref_count_increment_seq_cst(Int* refs) {
    __atomic_fetch_add(refs, 1, __ATOMIC_SEQ_CST);
}

/**
 * 以宽松内存序原子地增加引用计数
 * @param refs 指向引用计数变量的指针
 * @param nbytes 要增加的字节数
 */
template <typename Int>
__host__ __device__ __forceinline__ void add_ref_count_relaxed(Int* refs, int nbytes) {
    __atomic_fetch_add(refs, nbytes, __ATOMIC_RELAXED);
}

/**
 * 原子地减少引用计数并获取修改后的值（使用获取-释放内存序）
 * @param refs 指向引用计数的指针
 * @return 减少后的引用计数值
 */
template <typename Int>
__host__ __device__ __forceinline__ Int sub_ref_count_decrement_acq_rel(Int* refs) {
    return __atomic_sub_fetch(refs, 1, __ATOMIC_ACQ_REL);
}

////////////////////////////////////////////////////////////////////////////////////////////////
/*出发VMFault异常*/
__device__ __forceinline__ void trap() {
    // asm("trap;");
    __builtin_trap();
}

/**
 * @brief 执行全系统内存屏障（memory fence），确保所有线程都能看到最新的内存状态
 * @device 该函数仅在设备端（GPU）执行
 * @note 使用__threadfence_system()实现跨设备的全局内存一致性
 */
__device__ __forceinline__ void memory_fence() {
    // __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "");
    __threadfence_system();
}

__device__ __forceinline__ void memory_fence_gpu() {
    // __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");
    __threadfence();
}

__device__ __forceinline__ void memory_fence_cta() {
    // __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
    __threadfence_block();
}

template <typename Int>
__host__ __device__ __forceinline__ void st_relaxed_sys_global(Int* ptr, Int val) {
    __atomic_store_n(ptr, val, __ATOMIC_RELAXED);
    // asm volatile("st.relaxed.sys.global.s32 [%0], %1;"::"l"(ptr), "r"(val) : "memory");
}

template <typename Int>
__host__ __device__ __forceinline__ void st_release_sys_global(Int* ptr, Int val) {
    __atomic_store_n(ptr, val, __ATOMIC_RELEASE);
    // asm volatile("st.release.sys.global.s32 [%0], %1;"::"l"(ptr), "r"(val) : "memory");
}

template <typename Int>
__host__ __device__ __forceinline__ Int ld_acquire_sys_global(const Int* ptr) {
    Int ret;
    ret = __atomic_load_n(ptr, __ATOMIC_ACQUIRE);
    // asm volatile("ld.acquire.sys.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
    return ret;
}

} // namespace asm_ops
} // namespace sccl