LinearSolve.jl/ext/LinearSolveAMDGPUExt.jl at eaecc6052e3f2f7b4e33b09ce868f762134ce78b · jClugstor/LinearSolve.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
module LinearSolveAMDGPUExt

using AMDGPU
using LinearSolve: LinearSolve, LinearCache, AMDGPUOffloadLUFactorization,
                   AMDGPUOffloadQRFactorization, init_cacheval, OperatorAssumptions
using LinearSolve.LinearAlgebra, LinearSolve.SciMLBase

# LU Factorization
function SciMLBase.solve!(cache::LinearSolve.LinearCache, alg::AMDGPUOffloadLUFactorization;
        kwargs...)
    if cache.isfresh
        fact = AMDGPU.rocSOLVER.getrf!(AMDGPU.ROCArray(cache.A))
        cache.cacheval = fact
        cache.isfresh = false
    end

    A_gpu, ipiv = cache.cacheval
    b_gpu = AMDGPU.ROCArray(cache.b)

    AMDGPU.rocSOLVER.getrs!('N', A_gpu, ipiv, b_gpu)

    y = Array(b_gpu)
    cache.u .= y
    SciMLBase.build_linear_solution(alg, y, nothing, cache)
end

function LinearSolve.init_cacheval(alg::AMDGPUOffloadLUFactorization, A, b, u, Pl, Pr,
        maxiters::Int, abstol, reltol, verbose::LinearVerbosity,
        assumptions::OperatorAssumptions)
    AMDGPU.rocSOLVER.getrf!(AMDGPU.ROCArray(A))
end

# QR Factorization
function SciMLBase.solve!(cache::LinearSolve.LinearCache, alg::AMDGPUOffloadQRFactorization;
        kwargs...)
    if cache.isfresh
        A_gpu = AMDGPU.ROCArray(cache.A)
        tau = AMDGPU.ROCVector{eltype(A_gpu)}(undef, min(size(A_gpu)...))
        AMDGPU.rocSOLVER.geqrf!(A_gpu, tau)
        cache.cacheval = (A_gpu, tau)
        cache.isfresh = false
    end

    A_gpu, tau = cache.cacheval
    b_gpu = AMDGPU.ROCArray(cache.b)

    # Apply Q^T to b
    AMDGPU.rocSOLVER.ormqr!('L', 'T', A_gpu, tau, b_gpu)

    # Solve the upper triangular system
    m, n = size(A_gpu)
    AMDGPU.rocBLAS.trsv!('U', 'N', 'N', n, A_gpu, b_gpu)

    y = Array(b_gpu[1:n])
    cache.u .= y
    SciMLBase.build_linear_solution(alg, y, nothing, cache)
end

function LinearSolve.init_cacheval(alg::AMDGPUOffloadQRFactorization, A, b, u, Pl, Pr,
        maxiters::Int, abstol, reltol, verbose::LinearVerbosity,
        assumptions::OperatorAssumptions)
    A_gpu = AMDGPU.ROCArray(A)
    tau = AMDGPU.ROCVector{eltype(A_gpu)}(undef, min(size(A_gpu)...))
    AMDGPU.rocSOLVER.geqrf!(A_gpu, tau)
    (A_gpu, tau)
end

end