Programming AMD GPUs with Julia — ROCm Blogs
Julia 是一种高级的、通用的动态编程语言,通过 LLVM 自动编译为高效的本地代码,并支持多个平台。借助 LLVM,还可以支持编程 GPU,包括 AMD GPU。
AMDGPU.jl Julia 包是编程 AMD GPU 的主要入口点。它提供对高级数组编程和低级内核编程的支持,并与丰富的 Julia 生态系统集成,为用户提供统一的体验。
快速入门
安装 AMDGPU.jl 和安装常规的 Julia 包一样简单。在 Julia REPL 中执行以下命令(`]` 符号会进入 Pkg REPL 模式,`Backspace` 退出该模式):
] add AMDGPU
用户还需要有一个正常工作的 ROCm 安装,请参阅 requirements 部分。
示例:逐元素加法
一旦你成功安装了AMDGPU.jl包,你就可以导入该包并开始使用了:
julia> using AMDGPUjulia> a = AMDGPU.ones(Float32, 1024); # ';' 抑制 REPL 中的输出julia> b = AMDGPU.ones(Float32, 1024);julia> c = a .+ b; # '.' 为 '+' 操作符进行函数广播julia> sum(c) 2048
在上面的示例中,我们导入了该包,在GPU上分配了`a`和`b`数组并填充了`1`,计算了逐元素和(注意用 .+
进行操作符广播),最后计算了整个`c`数组的总和。
为了比较一下,我们还可以在CPU上进行相同的计算:
julia> ch = Array(a) .+ Array(b);julia> Array(c) ≈ ch # '≈' 可以用 'isapprox(Array(c), ch)' 替换以避免使用Unicode符号Unicode true
这里我们首先将 a
和 b
从GPU传输到CPU,计算逐元素和,将其存储在 ch
中,并使用 ≈
操作符(Julia支持Unicode输入)与 c
进行比较。
示例: 元素级加法的GPU内核
我们还可以通过编写自定义GPU内核来执行相同的计算。
julia> function vadd!(c, a, b)i = workitemIdx().x + (workgroupIdx().x - 1) * workgroupDim().xif i ≤ length(c)@inbounds c[i] = a[i] + b[i]endreturnend vadd! (generic function with 1 method)
类似于C++/HIP内核,Julia支持AMD GPU特定的内在函数,可在内核中使用。在这里,我们以与常规HIP内核相同的方式计算单个workitem的索引`i`。然后我们将`a`和`b`的元素级结果写入`c`中(`@inbounds`宏禁用了边界检查,从而提高了性能)。
注意: 由于内核是常规的GPU函数,所有内核应返回nothing。
使用方便的`@roc`宏可以启动内核:
julia> groupsize = 256;julia> gridsize = cld(length(c), groupsize);julia> @roc groupsize=groupsize gridsize=gridsize vadd!(c, a, b);julia> Array(c) ≈ ch true
注意: 所有内核启动都是异步的,因此在需要时用户必须显式使用`AMDGPU.synchronize()`进行同步。然而,在GPU到CPU的传输过程中,AMDGPU.jl在内部进行了同步。
我们可以看到,用Julia进行内核编程是非常容易的,并且内核在功能上并不受限,与HIP保持同步。
与Julia生态系统的集成
AMDGPU.jl将ROCm库与Julia生态系统集成,提供了一种统一的体验,在使用由AMDGPU.jl、CPU或其他加速器支持的数组方面几乎没有区别。
例如,使用rocBLAS进行常见的BLAS操作,并且Julia的运算符会派发到这些操作以提高效率。
julia> a = AMDGPU.rand(Float32, 1024, 1024);julia> b = AMDGPU.rand(Float32, 1024, 1024);julia> c = a * b; # dispatches to rocBLAS for matrix multiplication
可以使用Flux.jl 或Lux.jl进行机器学习,它们支持常见的构建块:
julia> using AMDGPU, Flux;julia> Flux.gpu_backend!("AMDGPU");julia> model = Conv((3, 3), 3 => 7, relu; bias=false);julia> x = AMDGPU.rand(Float32, (100, 100, 3, 50)); # 随机生成的图像,形状为 WxHxCxB。julia> y = model(x); # 派发到MIOpen进行卷积
可以使用Zygote.jl给定任何Julia函数计算梯度:
julia> θ = AMDGPU.rand(Float32, 16, 16);julia> x = AMDGPU.rand(Float32, 16, 16);julia> loss, grads = Zygote.withgradient(θ) do θsum(θ * x)end;
以及更多功能!
性能
在使用有效的构造的前提下,Julia GPU代码的性能与C++相当,有时甚至会超越C++,这取决于工作负载。
下面是一个在MI250x GPU上使用AMDGPU.jl实现的memcopy和2D扩散内核的性能对比。
对于性能检查,可以使用profiling来获取整个程序的时间线视图。
使用`@device_code_...(
llvm`,`gcn`,`lowered`)宏可以在每个内核的基础上转储不同的中间表示(未优化的LLVM IR,优化的LLVM IR,汇编)。
下面是为上面定义的`vadd!`内核生成的优化后的LLVM IR:
julia> @device_code_llvm @roc launch=false vadd!(c, a, b)
; @ REPL[4]:1 within `vadd!` define amdgpu_kernel void @_Z5vadd_14ROCDeviceArrayI7Float32Li1ELi1EES_IS0_Li1ELi1EES_IS0_Li1ELi1EE({ i64, i64, i64, i64, i64, i64, i32, i32, i64, i64, i64, i64 } %state,{ [1 x i64], i8 addrspace(1)*, i64 } %0,{ [1 x i64], i8 addrspace(1)*, i64 } %1,{ [1 x i64], i8 addrspace(1)*, i64 } %2 ) local_unnamed_addr #1 { conversion:%.fca.2.extract9 = extractvalue { [1 x i64], i8 addrspace(1)*, i64 } %0, 2 ; @ REPL[4]:2 within `vadd!`%3 = call i32 @llvm.amdgcn.workitem.id.x()%4 = add nuw nsw i32 %3, 1%5 = call i32 @llvm.amdgcn.workgroup.id.x()%6 = zext i32 %5 to i64%7 = call i8 addrspace(4)* @llvm.amdgcn.dispatch.ptr()%8 = getelementptr inbounds i8, i8 addrspace(4)* %7, i64 4%9 = bitcast i8 addrspace(4)* %8 to i16 addrspace(4)*%10 = load i16, i16 addrspace(4)* %9, align 4%11 = zext i16 %10 to i64%12 = mul nuw nsw i64 %11, %6%13 = zext i32 %4 to i64%14 = add nuw nsw i64 %12, %13 ; @ REPL[4]:3 within `vadd!`%.not = icmp sgt i64 %14, %.fca.2.extract9br i1 %.not, label %L92, label %L45L45: ; preds = %conversion%.fca.1.extract = extractvalue { [1 x i64], i8 addrspace(1)*, i64 } %2, 1%.fca.1.extract2 = extractvalue { [1 x i64], i8 addrspace(1)*, i64 } %1, 1%.fca.1.extract8 = extractvalue { [1 x i64], i8 addrspace(1)*, i64 } %0, 1 ; @ REPL[4]:4 within `vadd!`%15 = add nsw i64 %14, -1%16 = bitcast i8 addrspace(1)* %.fca.1.extract2 to float addrspace(1)*%17 = getelementptr inbounds float, float addrspace(1)* %16, i64 %15%18 = load float, float addrspace(1)* %17, align 4%19 = bitcast i8 addrspace(1)* %.fca.1.extract to float addrspace(1)*%20 = getelementptr inbounds float, float addrspace(1)* %19, i64 %15%21 = load float, float addrspace(1)* %20, align 4%22 = fadd float %18, %21%23 = bitcast i8 addrspace(1)* %.fca.1.extract8 to float addrspace(1)*%24 = getelementptr inbounds float, float addrspace(1)* %23, i64 %15store float %22, float addrspace(1)* %24, align 4br label %L92L92: ; preds = %L45, %conversion ; @ REPL[4]:6 within `vadd!`ret void }
通过这种方式,用户可以对自己的代码进行精细控制,并可以针对高性能应用程序来优化。
应用与库
由于与丰富生态系统的集成,实现应用变得极其容易,这里列举了一些支持AMD GPU的应用:
-
Nerf.jl:用原生Julia实现的Instant-NGP。
-
Whisper.jl: 流行的语音转文本模型。
-
Diffusers.jl: 稳定扩散1.5模型。
-
GPU4GEO: 使用LUMI超级计算机模拟冰的运动,针对LUMI-G的AMD MI250x GPU。
尝试一下
AMDGPU.jl 支持Linux和Windows操作系统以及各种设备!
致谢
特别感谢Anton Smirnov和Julia社区为本文做出的贡献。ROCm软件生态系统因社区项目(如Julia)而更加强大,让你以全新的方式使用AMD GPU。如果你有项目想在这里分享,请提出问题或PR。