Skip to content
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.

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

Add support for Adapt which enables converting tensor networks to GPU #187

Merged
merged 4 commits into from
May 30, 2024

Conversation

mtfishman
Copy link
Member

@mtfishman mtfishman commented May 30, 2024

With this we can now do:

using Metal: mtl
using NamedGraphs.NamedGraphGenerators: named_grid
using ITensorNetworks: random_tensornetwork, siteinds

g = named_grid((2, 2))
s = siteinds("S=1/2", g)
tn = random_tensornetwork(s)
tn_mtl = mtl(tn)

and we can see the tensors have been moved to GPU:

julia> tn[1, 1]
ITensor ord=3 (dim=2|id=976|"S=1/2,Site,n=1×1") (dim=1|id=434|"1×1,2×1") (dim=1|id=601|"1×1,1×2")
NDTensors.Dense{Float64, Vector{Float64}}

julia> tn_mtl[1, 1]
ITensor ord=3 (dim=2|id=976|"S=1/2,Site,n=1×1") (dim=1|id=434|"1×1,2×1") (dim=1|id=601|"1×1,1×2")
NDTensors.Dense{Float32, Metal.MtlVector{Float32, Metal.MTL.Private}}

@JoeyT1994 with this you should then be able to follow the same instructions that are shown here: https://itensor.github.io/ITensors.jl/dev/RunningOnGPUs.html to run calculation on GPU, for example for gate application before calling apply load the relevant GPU package and transfer the gates and tensor network to GPU using the corresponding conversion function, i.e. cu, mtl, roc, etc.

There may be parts of the library code that aren't generic enough for GPU, for example making implicit assumptions about the element type, constructing intermediate tensors on CPU instead of on the GPU device of the tensor network, etc. We had to go through a process of stamping out those kinds of issues in ITensors.jl and ITensorMPS.jl but they aren't hard to fix using Adapt.adapt.

@mtfishman mtfishman merged commit b4659b9 into main May 30, 2024
5 of 6 checks passed
@mtfishman mtfishman deleted the adapt branch May 30, 2024 17:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant