# 安装 CUDA.jl
using Pkg
Pkg.add("CUDA")

# 验证安装
using CUDA
CUDA.functional()         # true 表示可用
CUDA.version()            # CUDA 驱动版本
CUDA.runtime_version()    # CUDA 运行时版本

# GPU 信息
dev = CUDA.device()
CUDA.name(dev)            # GPU 名称
CUDA.totalmem(dev)        # 显存总量
CUDA.capability(dev)      # 计算能力

# 多 GPU 环境
CUDA.ndevices()           # GPU 数量
CUDA.devices()            # 所有设备列表

using CUDA

# 从普通数组创建
A_cpu = rand(Float32, 1000, 1000)
A_gpu = CuArray(A_cpu)              # CPU → GPU
A_back = Array(A_gpu)               # GPU → CPU

# 直接在 GPU 上创建
B_gpu = CUDA.zeros(Float32, 1000, 1000)    # 全零
C_gpu = CUDA.ones(Float32, 1000, 1000)     # 全一
D_gpu = CUDA.rand(Float32, 1000, 1000)     # 随机
E_gpu = CUDA.randn(Float32, 1000, 1000)    # 正态随机

# GPU 数组操作（与 CPU 数组语法完全一致）
F_gpu = A_gpu * B_gpu              # 矩阵乘法
G_gpu = A_gpu .+ C_gpu            # 逐元素加法
H_gpu = sin.(D_gpu)               # 逐元素 sin
sum_gpu = sum(D_gpu)              # 求和（结果在 GPU 上）
max_gpu = maximum(D_gpu)          # 最大值

# 类型转换
A_f64 = CuArray{Float64}(A_gpu)   # 转换精度

using CUDA

# 自定义 GPU 核函数
# 向量加法
function vector_add!(C, A, B)
    i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    if i <= length(C)
        C[i] = A[i] + B[i]
    end
    return nothing
end

n = 1024 * 1024
A = CUDA.rand(Float32, n)
B = CUDA.rand(Float32, n)
C = CUDA.zeros(Float32, n)

# 启动核函数
threads = 256
blocks = ceil(Int, n / threads)
@cuda threads=threads blocks=blocks vector_add!(C, A, B)

# 等待 GPU 完成
CUDA.synchronize()

# 验证
@assert Array(C) ≈ Array(A) + Array(B)

# 简单的矩阵乘法（非优化版本）
function matmul_kernel!(C, A, B)
    row = (blockIdx().y - 1) * blockDim().y + threadIdx().y
    col = (blockIdx().x - 1) * blockDim().x + threadIdx().x

    if row <= size(C, 1) && col <= size(C, 2)
        sum_val = 0.0f0
        for k in 1:size(A, 2)
            sum_val += A[row, k] * B[k, col]
        end
        C[row, col] = sum_val
    end
    return nothing
end

M, K, N = 256, 128, 256
A = CUDA.rand(Float32, M, K)
B = CUDA.rand(Float32, K, N)
C = CUDA.zeros(Float32, M, N)

# 2D 线程块
threads = (16, 16)
blocks = (ceil(Int, N/16), ceil(Int, M/16))
@cuda threads=threads blocks=blocks matmul_kernel!(C, A, B)
CUDA.synchronize()

using CUDA

# CUDA 线程层次结构：
# Grid → Block → Thread
# blockIdx().x/y/z   : 块索引
# blockDim().x/y/z   : 块维度
# threadIdx().x/y/z  : 线程索引
# gridDim().x/y/z    : 网格维度

# 全局线程索引
function global_index_1d()
    return (blockIdx().x - 1) * blockDim().x + threadIdx().x
end

function global_index_2d()
    row = (blockIdx().y - 1) * blockDim().y + threadIdx().y
    col = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    return row, col
end

# 通用核函数模板
function generic_kernel!(output, input)
    idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    if idx <= length(input)
        output[idx] = some_operation(input[idx])
    end
    return nothing
end

using CUDA

# 使用共享内存优化计算
# 共享内存是块内线程共享的高速缓存，比全局内存快 100 倍
function reduce_shared!(output, input)
    TILE = 256
    shared = @cuStaticSharedMem(Float32, 256)

    tid = threadIdx().x
    i = (blockIdx().x - 1) * blockDim().x + threadIdx().x

    # 加载到共享内存
    shared[tid] = i <= length(input) ? input[i] : 0.0f0
    sync_threads()

    # 归约（树形求和）
    stride = blockDim().x ÷ 2
    while stride > 0
        if tid <= stride
            shared[tid] += shared[tid + stride]
        end
        sync_threads()
        stride ÷= 2
    end

    if tid == 1
        output[blockIdx().x] = shared[1]
    end
    return nothing
end

using CUDA

# GPU 广播（自动并行化）
A = CUDA.rand(Float32, 1000, 1000)
B = CUDA.rand(Float32, 1000, 1000)

# Julia 广播语法自动在 GPU 上并行执行
C = sin.(A) .+ cos.(B) .+ 1.0f0
D = A .^ 2 .- B .* 2.0f0

# reduce 操作
total = sum(A)
max_val = maximum(A)
min_val = minimum(A)

# mapreduce
result = mapreduce(x -> x^2, +, A)    # 所有元素平方和

# GPU 上的逻辑索引
mask = A .> 0.5f0
count_true = sum(mask)

# 条件赋值
C = CUDA.zeros(Float32, 1000, 1000)
C[A .> 0.5f0] .= 1.0f0

# 多维规约
col_sums = sum(A, dims=1)     # 按列求和
row_means = mean(A, dims=2)   # 按行求均值

using CUDA

# CPU → GPU
A_cpu = rand(Float32, 1000, 1000)
A_gpu = CuArray(A_cpu)          # 方式1：显式转换
A_gpu = CUDA.cu(A_cpu)          # 方式2：cu 函数

# GPU → CPU
A_cpu = Array(A_gpu)            # 方式1
A_cpu = collect(A_gpu)          # 方式2（同 Array）

# 原地传输（避免分配）
A_cpu = rand(Float32, 100)
A_gpu = CUDA.zeros(Float32, 100)
copyto!(A_gpu, A_cpu)           # CPU → GPU 原地
copyto!(A_cpu, A_gpu)           # GPU → CPU 原地

# 异步传输（使用 CUDA.stream）
A_gpu = CUDA.zeros(Float32, 1000, 1000)
CUDA.@sync copyto!(A_gpu, A_cpu)    # 等待完成

# 检查数据所在设备
CUDA.device(A_gpu)              # 返回设备 ID

using CUDA

# 查看显存使用
CUDA.memory_status()            # 显示已用/总显存

# 手动垃圾回收
CUDA.reclaim()                  # 释放 GPU 内存

# 内存池管理
CUDA.pool_status()              # 内存池状态
CUDA.memory_pool()              # 当前内存池

# 预分配避免频繁分配
function good_practice()
    # 预分配输出
    output = CUDA.zeros(Float32, 1000, 1000)
    input = CUDA.rand(Float32, 1000, 1000)

    for _ in 1:100
        output .= sin.(input)    # 原地操作，不分配新内存
    end
end

# 避免不必要的 GPU→CPU 传输
function bad_practice(A_gpu)
    A_cpu = Array(A_gpu)         # 传输到 CPU
    result = sum(A_cpu)          # 在 CPU 上计算
    return CuArray(result)       # 传回 GPU（完全不必要！）
end

function good_practice(A_gpu)
    return sum(A_gpu)            # 直接在 GPU 上计算
end

using CUDA, BenchmarkTools

# GPU 计时（必须同步）
function gpu_time()
    A = CUDA.rand(Float32, 10000, 10000)
    B = CUDA.rand(Float32, 10000, 10000)

    # 预热（第一次运行会触发编译）
    C = A * B
    CUDA.synchronize()

    # 计时
    t = CUDA.@elapsed begin
        C = A * B
        CUDA.synchronize()
    end
    println("GPU 矩阵乘法: $(round(t*1000, digits=2)) ms")
    return C
end

# CPU 对比
function cpu_time()
    A = rand(Float32, 10000, 10000)
    B = rand(Float32, 10000, 10000)

    t = @elapsed C = A * B
    println("CPU 矩阵乘法: $(round(t*1000, digits=2)) ms")
    return C
end

# GPU 性能分析
CUDA.@profile A * B

using CUDA

# 列出所有 GPU
for dev in CUDA.devices()
    println("$(CUDA.name(dev)): $(CUDA.totalmem(dev) ÷ 1024^3) GB")
end

# 在指定 GPU 上工作
CUDA.device!(0) do
    A = CUDA.rand(Float32, 1000, 1000)
    result = sum(A * A)
    println("GPU 0 结果: $result")
end

# 多 GPU 数据并行
function multi_gpu_sum(data_chunks)
    n = CUDA.ndevices()
    results = Vector{Float32}(undef, n)
    Threads.@threads for i in 1:n
        CUDA.device!(i - 1) do
            results[i] = sum(CuArray(data_chunks[i]))
        end
    end
    return sum(results)
end

using CUDA, BenchmarkTools

n = 4096
A_cpu = rand(Float32, n, n)
B_cpu = rand(Float32, n, n)

# CPU
A_gpu = CuArray(A_cpu)
B_gpu = CuArray(B_cpu)
t_cpu = @elapsed A_cpu * B_cpu
t_gpu = CUDA.@elapsed begin
    A_gpu * B_gpu; CUDA.synchronize()
end
println("CPU: $(round(t_cpu*1000, digits=1))ms, GPU: $(round(t_gpu*1000, digits=1))ms")
println("加速比: $(round(t_cpu/t_gpu, digits=1))x")

using CUDA

function monte_carlo_pi_gpu(n::Int)
    x = CUDA.rand(Float32, n)
    y = CUDA.rand(Float32, n)
    inside = (x .^ 2 .+ y .^ 2) .<= 1.0f0
    return 4.0 * sum(inside) / n
end

n = 100_000_000
t_gpu = CUDA.@elapsed pi_gpu = monte_carlo_pi_gpu(n)
t_cpu = @elapsed pi_cpu = 4.0 * sum(rand(n).^2 + rand(n).^2 .<= 1.0) / n
println("GPU π ≈ $pi_gpu ($(round(t_gpu*1000, digits=1))ms)")
println("CPU π ≈ $pi_cpu ($(round(t_cpu*1000, digits=1))ms)")

# ❌ 避免：频繁传输
for i in 1:100
    A_cpu = Array(A_gpu)      # GPU → CPU
    result = process(A_cpu)
    A_gpu = CuArray(result)   # CPU → GPU
end

# ✅ 推荐：数据留在 GPU
for i in 1:100
    A_gpu .= process_gpu.(A_gpu)    # 全在 GPU 上完成
end