与Julia生态系统集成
AMDGPU.jl将ROCm库与Julia生态系统集成在一起,提供统一的体验,使用AMDGPU.j、CPU或其他加速器支持的阵列几乎没有区别。
例如,rocBLAS用于常见的BLAS操作,Julia的算子为了提高效率而分派给他们。
julia> a = AMDGPU.rand(Float32, 1024, 1024);
julia> b = AMDGPU.rand(Float32, 1024, 1024);
julia> c = a * b; # 分派给rocBLAS进行矩阵乘法
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;
2.7.5性能
如果使用高效的构造,Julia GPU代码的性能与C++相当,有时甚至超过它,具体取决于工作负载。
使用AMDGPU.jl在Julia中实现,并在MI250x GPU上执行的memcopy,以及2D扩散内核的性能比较。
对于性能检查,可以使用分析来获得整个程序的时间线视图,如图2-11所示。
图2-11 HIP和AMDGPU网格点数目、内存时间线视图
还有@device_code_... (llvm, gcn, lowered)宏,以每个内核为基础转储不同的中间表示(未优化的llvm IR、优化的LLVMIR、汇编)。
下面是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 {
转换:
%.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.extract9
br i1 %.not, label %L92, label %L45
L45: ; 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 %15
store float %22, float addrspace(1)* %24, align 4
br label %L92
L92: ; preds = %L45, %conversion
; @ REPL[4]:6 within `vadd!`
ret void
}
用户可以很好地控制他们的代码,并可以针对高性能应用程序。
2.7.6应用程序和库
通过丰富的生态系统集成,实现应用程序非常容易,以下是其中一些支持AMD GPU的应用程序:
Nerf.jl:在原生Julia中即时实现NGP。
Whisper.jl:流行的语音转文本模型。
Diffusers.jl:稳定扩散。
GPU4GEO:使用LUMI超级计算机,针对LUMI-G的AMD MI250x GPU,对冰运动进行建模。