#pragma once

/**
 * HIP 到 CUDA 的兼容层
 * 
 * 使用 nvcc 编译时，自动将 HIP API 映射到 CUDA API
 * 使用 hipcc 编译时，使用原生 HIP 头文件
 */

#if defined(__NVCC__) || defined(__CUDACC__)

#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <stdio.h>

// Runtime API 映射
#define hipMalloc cudaMalloc
#define hipFree cudaFree
#define hipMemcpy cudaMemcpy
#define hipMemcpyHostToDevice cudaMemcpyHostToDevice
#define hipMemcpyDeviceToHost cudaMemcpyDeviceToHost
#define hipMemset cudaMemset
#define hipDeviceSynchronize cudaDeviceSynchronize
#define hipGetDeviceProperties cudaGetDeviceProperties
#define hipGetErrorString cudaGetErrorString

// Event API 映射
#define hipEvent_t cudaEvent_t
#define hipEventCreate cudaEventCreate
#define hipEventDestroy cudaEventDestroy
#define hipEventRecord cudaEventRecord
#define hipEventSynchronize cudaEventSynchronize
#define hipEventElapsedTime cudaEventElapsedTime

// 数据类型映射
#define hipDeviceProp_t cudaDeviceProp
#define hipError_t cudaError_t
#define hipSuccess cudaSuccess
// CUDA 使用 __nv_bfloat16，HIP 使用 hip_bfloat16
typedef __nv_bfloat16 hip_bfloat16;

// Shuffle 指令映射
// CUDA 9.0+ 需要使用带 _sync 后缀的版本，并传入 warp mask
// 0xffffffff 表示整个 warp 的所有线程都参与
#ifndef __shfl_down
#define __shfl_down(val, offset) __shfl_down_sync(0xffffffff, val, offset)
#endif

#else

#include <hip/hip_runtime.h>
#include <hip/hip_bfloat16.h>

#endif