Skip to content

Very basic CUDA support #66

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kshyatt wants to merge 1 commit into from
Closed

Very basic CUDA support #66

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,16 @@ TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"

[weakdeps]
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"

[extensions]
MPSKitModelsCUDAExt = ["CUDA", "cuTENSOR"]

[compat]
CUDA = "5"
cuTENSOR = "2"
MPSKit = "0.13.1"
MacroTools = "0.5"
PrecompileTools = "1"
Expand All @@ -23,9 +32,11 @@ TupleTools = "1"
julia = "1.10"

[extras]
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"

[targets]
test = ["Test", "TestExtras", "QuadGK"]
test = ["Test", "TestExtras", "QuadGK", "CUDA", "cuTENSOR"]
183 changes: 183 additions & 0 deletions ext/MPSKitModelsCUDAExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,183 @@
module MPSKitModelsCUDAExt

using MPSKitModels, CUDA, cuTENSOR

import MPSKitModels: build_a_plus_left!, build_a_plus_right!, build_a_min_left!, build_a_min_right!
import MPSKitModels: e_plusmin_up, e_plusmin_down, e_number_up, e_number_down, spinmatrices
using MPSKitModels: SU2Irrep, U1Irrep, Trivial, fusiontrees, sectortype, block, dual

function e_number_down(::Type{TA}, ::Type{Trivial} = Trivial, ::Type{Trivial} = Trivial) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.single_site_operator(TA, Trivial, Trivial)
I = sectortype(t)
CUDA.@allowscalar begin
t[(I(1), I(1))][2, 2] = 1
t[(I(0), I(0))][2, 2] = 1
end
return t
end
function e_number_down(::Type{TA}, ::Type{Trivial}, ::Type{U1Irrep}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.single_site_operator(TA, Trivial, U1Irrep)
I = sectortype(t)
CUDA.@allowscalar begin
t[(I(1, -1 // 2), dual(I(1, -1 // 2)))][1, 1] = 1
t[(I(0, 0), I(0, 0))][2, 2] = 1
end
return t
end
function e_number_down(::Type{TA}, ::Type{U1Irrep}, ::Type{Trivial}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.single_site_operator(TA, U1Irrep, Trivial)
I = sectortype(t)
CUDA.@allowscalar begin
block(t, I(1, 1))[2, 2] = 1 # expected to be [1,2]
block(t, I(0, 2))[1, 1] = 1
end
return t
end
function e_number_down(::Type{TA}, ::Type{U1Irrep}, ::Type{U1Irrep}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.single_site_operator(TA, U1Irrep, U1Irrep)
I = sectortype(t)
CUDA.@allowscalar begin
block(t, I(1, 1, -1 // 2)) .= 1
block(t, I(0, 2, 0)) .= 1
end
return t
end
function e_number_up(::Type{TA}, ::Type{Trivial}, ::Type{Trivial}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.single_site_operator(TA, Trivial, Trivial)
I = sectortype(t)
CUDA.@allowscalar begin
t[(I(1), I(1))][1, 1] = 1
t[(I(0), I(0))][2, 2] = 1
end
return t
end
function e_number_up(::Type{TA}, ::Type{U1Irrep}, ::Type{Trivial}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.single_site_operator(TA, U1Irrep, Trivial)
I = sectortype(t)
CUDA.@allowscalar begin
block(t, I(1, 1))[1, 1] = 1
block(t, I(0, 2))[1, 1] = 1
end
return t
end
function e_number_up(::Type{TA}, ::Type{U1Irrep}, ::Type{U1Irrep}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.single_site_operator(TA, U1Irrep, U1Irrep)
I = sectortype(t)
CUDA.@allowscalar begin
block(t, I(1, 1, 1 // 2)) .= 1
block(t, I(0, 2, 0)) .= 1
end
return t
end
function e_plusmin_up(::Type{TA}, ::Type{Trivial}, ::Type{Trivial}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.two_site_operator(TA, Trivial, Trivial)
I = sectortype(t)
CUDA.@allowscalar begin
t[(I(1), I(0), dual(I(0)), dual(I(1)))][1, 1, 1, 1] = 1
t[(I(1), I(1), dual(I(0)), dual(I(0)))][1, 2, 1, 2] = 1
t[(I(0), I(0), dual(I(1)), dual(I(1)))][2, 1, 2, 1] = -1
t[(I(0), I(1), dual(I(1)), dual(I(0)))][2, 2, 2, 2] = -1
end
return t
end
function e_plusmin_up(::Type{TA}, ::Type{U1Irrep}, ::Type{Trivial}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.two_site_operator(TA, U1Irrep, Trivial)
I = sectortype(t)
CUDA.@allowscalar begin
t[(I(1, 1), I(0, 0), dual(I(0, 0)), dual(I(1, 1)))][1, 1, 1, 1] = 1
t[(I(1, 1), I(1, 1), dual(I(0, 0)), dual(I(0, 2)))][1, 2, 1, 1] = 1
t[(I(0, 2), I(0, 0), dual(I(1, 1)), dual(I(1, 1)))][1, 1, 2, 1] = -1
t[(I(0, 2), I(1, 1), dual(I(1, 1)), dual(I(0, 2)))][1, 2, 2, 1] = -1
end
return t
end

function e_plusmin_down(::Type{TA}, ::Type{Trivial}, ::Type{Trivial}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.two_site_operator(TA, Trivial, Trivial)
I = sectortype(t)
CUDA.@allowscalar begin
t[(I(1), I(0), dual(I(0)), dual(I(1)))][2, 1, 1, 2] = 1
t[(I(1), I(1), dual(I(0)), dual(I(0)))][2, 1, 1, 2] = -1
t[(I(0), I(0), dual(I(1)), dual(I(1)))][2, 1, 1, 2] = 1
t[(I(0), I(1), dual(I(1)), dual(I(0)))][2, 1, 1, 2] = -1
end
return t
end
function e_plusmin_down(::Type{TA}, ::Type{Trivial}, ::Type{U1Irrep}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.two_site_operator(T, Trivial, U1Irrep)
I = sectortype(t)
CUDA.@allowscalar begin
t[(I(1, -1 // 2), I(0, 0), dual(I(0, 0)), dual(I(1, -1 // 2)))][1, 1, 1, 1] = 1
t[(I(1, -1 // 2), I(1, 1 // 2), dual(I(0, 0)), dual(I(0, 0)))][1, 1, 1, 2] = -1
t[(I(0, 0), I(0, 0), dual(I(1, 1 // 2)), dual(I(1, -1 // 2)))][2, 1, 1, 1] = 1
t[(I(0, 0), I(1, 1 // 2), dual(I(1, 1 // 2)), dual(I(0, 0)))][2, 1, 1, 2] = -1
end
return t
end
function e_plusmin_down(::Type{TA}, ::Type{U1Irrep}, ::Type{Trivial}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.two_site_operator(TA, U1Irrep, Trivial)
I = sectortype(t)
CUDA.@allowscalar begin
t[(I(1, 1), I(0, 0), dual(I(0, 0)), dual(I(1, 1)))][2, 1, 1, 2] = 1
t[(I(1, 1), I(1, 1), dual(I(0, 0)), dual(I(0, 2)))][2, 1, 1, 1] = -1
t[(I(0, 2), I(0, 0), dual(I(1, 1)), dual(I(1, 1)))][1, 1, 1, 2] = 1
t[(I(0, 2), I(1, 1), dual(I(1, 1)), dual(I(0, 2)))][1, 1, 1, 1] = -1
end
return t
end
function e_plusmin_down(::Type{TA}, ::Type{U1Irrep}, ::Type{U1Irrep}) where {TA <: CuArray}
t = MPSKitModels.HubbardOperators.two_site_operator(T, U1Irrep, U1Irrep)
I = sectortype(t)
t[(I(1, 1, -1 // 2), I(0, 0, 0), dual(I(0, 0, 0)), dual(I(1, 1, -1 // 2)))] .= 1
t[(I(1, 1, -1 // 2), I(1, 1, 1 // 2), dual(I(0, 0, 0)), dual(I(0, 2, 0)))] .= -1
t[(I(0, 2, 0), I(0, 0, 0), dual(I(1, 1, 1 // 2)), dual(I(1, 1, -1 // 2)))] .= 1
t[(I(0, 2, 0), I(1, 1, 1 // 2), dual(I(1, 1, 1 // 2)), dual(I(0, 2, 0)))] .= -1
return t
end

function build_a_plus_left!(::Type{U1Irrep}, ::Type{TA}, a⁺) where {TA <: CuArray}
for (f1, f2) in fusiontrees(a⁺)
c₁, c₂ = f1.uncoupled[1], f2.uncoupled[1]
if c₁.charge == c₂.charge + 1
CUDA.@allowscalar a⁺[f1, f2] .= -sqrt(c₁.charge)
end
end
return
end

function build_a_plus_right!(::Type{U1Irrep}, ::Type{TA}, a⁺) where {TA <: CuArray}
for (f1, f2) in fusiontrees(a⁺)
c₁, c₂ = f1.uncoupled[2], f2.uncoupled[1]
if c₁.charge == c₂.charge + 1
CUDA.@allowscalar a⁺[f1, f2] .= -sqrt(c₁.charge)
end
end
return
end

function build_a_min_left!(::Type{U1Irrep}, ::Type{TA}, a⁻) where {TA <: CuArray}
for (f1, f2) in fusiontrees(a⁻)
c₁, c₂ = f1.uncoupled[1], f2.uncoupled[1]
if c₁.charge + 1 == c₂.charge
CUDA.@allowscalar a⁻[f1, f2] .= -sqrt(c₂.charge)
end
end
return
end

function build_a_min_right!(::Type{U1Irrep}, ::Type{TA}, a⁻) where {TA <: CuArray}
for (f1, f2) in fusiontrees(a⁻)
c₁, c₂ = f1.uncoupled[2], f2.uncoupled[1]
if c₁.charge + 1 == c₂.charge
CUDA.@allowscalar a⁻[f1, f2] .= -sqrt(c₂.charge)
end
end
return
end

function spinmatrices(s::Union{Rational{Int}, Int}, ::Type{TorA}) where {T, TorA <: CuArray{T}}
hSx, hSy, hSz, hI = spinmatrices(s, T)
return CuArray(hSx), CuArray(hSy), CuArray(hSz), CuArray(hI)
end

end
68 changes: 34 additions & 34 deletions src/models/hamiltonians.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,12 +205,12 @@ function heisenberg_XYZ(lattice::AbstractLattice; kwargs...)
return heisenberg_XYZ(ComplexF64, lattice; kwargs...)
end
function heisenberg_XYZ(
T::Type{<:Number} = ComplexF64, lattice::AbstractLattice = InfiniteChain(1);
::Type{TorA}, lattice::AbstractLattice = InfiniteChain(1);
Jx = 1.0, Jy = 1.0, Jz = 1.0, spin = 1
)
term = rmul!(S_xx(T, Trivial; spin = spin), Jx) +
rmul!(S_yy(T, Trivial; spin = spin), Jy) +
rmul!(S_zz(T, Trivial; spin = spin), Jz)
) where {TorA}
term = rmul!(S_xx(TorA, Trivial; spin = spin), Jx) +
rmul!(S_yy(TorA, Trivial; spin = spin), Jy) +
rmul!(S_zz(TorA, Trivial; spin = spin), Jz)
return @mpoham sum(nearest_neighbours(lattice)) do (i, j)
return term{i, j}
end
Expand Down Expand Up @@ -239,19 +239,19 @@ function bilinear_biquadratic_model(
return bilinear_biquadratic_model(ComplexF64, symmetry, lattice; kwargs...)
end
function bilinear_biquadratic_model(
elt::Type{<:Number}, lattice::AbstractLattice;
::Type{TorA}, lattice::AbstractLattice;
kwargs...
)
return bilinear_biquadratic_model(elt, Trivial, lattice; kwargs...)
) where {TorA}
return bilinear_biquadratic_model(TorA, Trivial, lattice; kwargs...)
end
function bilinear_biquadratic_model(
elt::Type{<:Number} = ComplexF64, symmetry::Type{<:Sector} = Trivial,
::Type{TorA}, symmetry::Type{<:Sector} = Trivial,
lattice::AbstractLattice = InfiniteChain(1);
spin = 1, J = 1.0, θ = 0.0
)
) where {TorA}
return @mpoham sum(nearest_neighbours(lattice)) do (i, j)
return J * cos(θ) * S_exchange(elt, symmetry; spin = spin){i, j} +
J * sin(θ) * (S_exchange(elt, symmetry; spin = spin)^2){i, j}
return J * cos(θ) * S_exchange(TorA, symmetry; spin = spin){i, j} +
J * sin(θ) * (S_exchange(TorA, symmetry; spin = spin)^2){i, j}
end
end

Expand Down Expand Up @@ -328,18 +328,18 @@ function hubbard_model(
)
return hubbard_model(ComplexF64, particle_symmetry, spin_symmetry; kwargs...)
end
function hubbard_model(elt::Type{<:Number}, lattice::AbstractLattice; kwargs...)
return hubbard_model(elt, Trivial, Trivial, lattice; kwargs...)
function hubbard_model(::Type{TorA}, lattice::AbstractLattice; kwargs...) where {TorA}
return hubbard_model(TorA, Trivial, Trivial, lattice; kwargs...)
end
function hubbard_model(
T::Type{<:Number} = ComplexF64, particle_symmetry::Type{<:Sector} = Trivial,
::Type{TorA}, particle_symmetry::Type{<:Sector} = Trivial,
spin_symmetry::Type{<:Sector} = Trivial, lattice::AbstractLattice = InfiniteChain(1);
t = 1.0, U = 1.0, mu = 0.0, n::Integer = 0
)
hopping = e⁺e⁻(T, particle_symmetry, spin_symmetry) +
e⁻e⁺(T, particle_symmetry, spin_symmetry)
interaction_term = nꜛnꜜ(T, particle_symmetry, spin_symmetry)
N = e_number(T, particle_symmetry, spin_symmetry)
) where {TorA}
hopping = e⁺e⁻(TorA, particle_symmetry, spin_symmetry) +
e⁻e⁺(TorA, particle_symmetry, spin_symmetry)
interaction_term = nꜛnꜜ(TorA, particle_symmetry, spin_symmetry)
N = e_number(TorA, particle_symmetry, spin_symmetry)
return @mpoham begin
sum(nearest_neighbours(lattice)) do (i, j)
return -t * hopping{i, j}
Expand Down Expand Up @@ -377,14 +377,14 @@ function bose_hubbard_model(
return bose_hubbard_model(ComplexF64, symmetry, lattice; kwargs...)
end
function bose_hubbard_model(
elt::Type{<:Number} = ComplexF64, symmetry::Type{<:Sector} = Trivial,
::Type{TorA}, symmetry::Type{<:Sector} = Trivial,
lattice::AbstractLattice = InfiniteChain(1);
cutoff::Integer = 5, t = 1.0, U = 1.0, mu = 0.0, n = 0
)
hopping_term = a_plusmin(elt, symmetry; cutoff = cutoff) +
a_minplus(elt, symmetry; cutoff = cutoff)
N = a_number(elt, symmetry; cutoff = cutoff)
interaction_term = contract_onesite(N, N - id(domain(N)))
) where {TorA}
hopping_term = a_plusmin(TorA, symmetry; cutoff = cutoff) +
a_minplus(TorA, symmetry; cutoff = cutoff)
N = a_number(TorA, symmetry; cutoff = cutoff)
interaction_term = contract_onesite(N, N - id(TorA, domain(N)))

H = @mpoham begin
sum(nearest_neighbours(lattice)) do (i, j)
Expand Down Expand Up @@ -436,19 +436,19 @@ function tj_model(
)
return tj_model(ComplexF64, particle_symmetry, spin_symmetry; kwargs...)
end
function tj_model(elt::Type{<:Number}, lattice::AbstractLattice; kwargs...)
return tj_model(elt, Trivial, Trivial, lattice; kwargs...)
function tj_model(::Type{TorA}, lattice::AbstractLattice; kwargs...) where {TorA}
return tj_model(TorA, Trivial, Trivial, lattice; kwargs...)
end
function tj_model(
T::Type{<:Number} = ComplexF64, particle_symmetry::Type{<:Sector} = Trivial,
::Type{TorA}, particle_symmetry::Type{<:Sector} = Trivial,
spin_symmetry::Type{<:Sector} = Trivial, lattice::AbstractLattice = InfiniteChain(1);
t = 2.5, J = 1.0, mu = 0.0, slave_fermion::Bool = false
)
hopping = TJOperators.e_plusmin(T, particle_symmetry, spin_symmetry; slave_fermion) +
TJOperators.e_minplus(T, particle_symmetry, spin_symmetry; slave_fermion)
num = TJOperators.e_number(T, particle_symmetry, spin_symmetry; slave_fermion)
) where {TorA}
hopping = TJOperators.e_plusmin(TorA, particle_symmetry, spin_symmetry; slave_fermion) +
TJOperators.e_minplus(TorA, particle_symmetry, spin_symmetry; slave_fermion)
num = TJOperators.e_number(TorA, particle_symmetry, spin_symmetry; slave_fermion)
heisenberg = TJOperators.S_exchange(
T, particle_symmetry, spin_symmetry;
TorA, particle_symmetry, spin_symmetry;
slave_fermion
) -
(1 / 4) * (num ⊗ num)
Expand Down
Loading
Loading