yapanuwan
/
DVR-jl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Jacobi runs on GPU

This commit is contained in:
ysyapa 2023-08-27 23:44:22 -04:00
parent ad5bac2bf2
commit 2263c26215
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
Hamiltonian.jl
View File

@ -19,7 +19,7 @@ struct Hamiltonian{T}
∂ = ∂_1DOF.(Ref(s), k, k') ∂ = ∂_1DOF.(Ref(s), k, k')
K = exp(-2im * convert(T, ϕ)) .* (∂ * ∂) # TODO: Calculate K matrix elements directly K = exp(-2im * convert(T, ϕ)) .* (∂ * ∂) # TODO: Calculate K matrix elements directly
if mode == gpu_cutensor if mode == gpu_cutensor
K = CuTensor(K, ['a', 'A']) K = CuTensor(CuArray(K), ['a', 'A'])
Vs = CuArray(Vs) Vs = CuArray(Vs)
end end
return new{T}(s, K, Vs, hermitian, mode) return new{T}(s, K, Vs, hermitian, mode)
@ -55,7 +55,7 @@ function LinearAlgebra.mul!(out::Array{Complex{T}}, H::Hamiltonian{T}, v::Array{
end end
"cuTENSOR contraction and accumulation (C = A * B + C)" "cuTENSOR contraction and accumulation (C = A * B + C)"
function contract_accumulate!(alpha::Numer, C::CuTensor, A::CuTensor, B::CuTensor)::CuTensor function contract_accumulate!(alpha::Number, C::CuTensor, A::CuTensor, B::CuTensor)::CuTensor
cuTENSOR.contraction!(alpha, A.data, A.inds, cuTENSOR.CUTENSOR_OP_IDENTITY, B.data, B.inds, cuTENSOR.CUTENSOR_OP_IDENTITY, cuTENSOR.contraction!(alpha, A.data, A.inds, cuTENSOR.CUTENSOR_OP_IDENTITY, B.data, B.inds, cuTENSOR.CUTENSOR_OP_IDENTITY,
one(eltype(C)), C.data, C.inds, cuTENSOR.CUTENSOR_OP_IDENTITY, cuTENSOR.CUTENSOR_OP_IDENTITY) one(eltype(C)), C.data, C.inds, cuTENSOR.CUTENSOR_OP_IDENTITY, cuTENSOR.CUTENSOR_OP_IDENTITY)
return C return C
@ -77,7 +77,7 @@ function LinearAlgebra.mul!(out::CuArray{Complex{T}}, H::Hamiltonian{T}, v::CuAr
for coord = 1:coords for coord = 1:coords
i = which_index(H.s, dim, coord) i = which_index(H.s, dim, coord)
@assert v_t.inds == inds_template "v indices permuted" @assert v_t.inds == inds_template "v indices permuted"
@assert H.K_diag.inds[2] == 'A' "K_diag indices permuted" @assert H.K.inds[2] == 'A' "K_diag indices permuted"
H.K.inds[1] = 'a' - 1 + i H.K.inds[1] = 'a' - 1 + i
v_t.inds[i] = 'A' v_t.inds[i] = 'A'
#synchronize(ctx) #synchronize(ctx)