# Accelerating Tensor Computations in Julia with the GPU

How my code went from one week to one hour runtimes

# Introduction/Roadmap

Last year at JuliaCon, Matt Fishman and I gave a talk about our ongoing effort to port the ITensor code from C++ to Julia. At the time, I mentioned that we had begun trying to integrate a GPU-based tensor contraction backend and were looking forward to some significant speedups. We ended up completing this integration, and saw runtimes for representative parameters go from one week to one hour. In this post I'm going to go over:

- The physics problem we were writing code to solve
- Why the GPU is a good candidate to accelerate our simulations to solve this problem
- Why Julia is a natural choice if you want to take advantage of the GPU easily
- How, after we wrote the initial implementation, we made it faster by removing roadblocks
- Some final thoughts about why this worked out well for us

# What is ITensor

To understand why we thought this would be fruitful, and some performance traps we could already anticipate before writing a single line of GPU code, you have to understand a bit about the problem we're trying to solve. Tensor network algorithms are a very active area of research at the intersection of condensed matter physics, high energy physics, quantum information, and computer science. DMRG, which is the most successful numerical method for condensed matter systems in 1D, wasn't originally formulated as a tensor network algorithm, but “tensor network speak” turns out to be a natural language with which to discuss DMRG and its descendants. (If you're not sure what it's meant to solve, imagine a long chain of quantum objects, each with a fretful relationship with their neighbours. Will they overcome their differences and work together to form a magnet, or continue arguing with each other and fail to come to a consensus? DMRG is an efficient way to answer this question.) All this is to say that there's this class of algorithms physicists (and increasingly computer scientists) are interested in and they work well.

The driving idea of a tensor network algorithm is to take some high-dimensional optimization problem (solving for groundstates (eigenvectors corresponding to the minimal eigenvalue) of quantum many body systems is an example of this, since the full dimension of the system grows exponentially in the number of constituent particles) and compress it down to a much lower dimensional problem while retaining most of the important features. We do this by taking a `d^N`

length vector, where `d`

is the number of degrees of freedom of each constituent, and `N`

is the number of constituents, and breaking it down into a set of multidimensionsal tensors, the number of which hopefully scales like `N`

or at least much less than exponentially.

If you've studied linear algebra before, you've seen some simple examples of tensors: scalars, vectors, and matrices. We say a scalar is a 0-rank tensor, a vector a 1-rank tensor (since it has one index), a matrix a 2-rank tensor (since it has two indices), and then there are higher rank tensors, with three or more indices. When we multiply two matrices `A_ij`

and `B_jk`

together to get `C_ik`

, we're performing a “tensor contratction”, and we compute `C_ik`

by ssumming over indices `i`

and `j`

. Similarly, if we had high-rank tensors `A_ijkl`

and `B_lmin`

, we could contract them to get `C_jkmn`

- again, by summing over indices `i`

and `l`

.

Most tensor network algorithms are based on performing this decomposition and then iteratively improving it towards a target vector. Usually in physics that target is a physical state, but tensor networks have also been used for machine learning tasks and can represent quantum error correcting codes as well.

ITensor is a C++ package dedicated to providing both high level algorithms using tensor networks, like DMRG, and the low-level building blocks to create your own. ITensor makes it easy to create tensors out of indices, perform linear algebra factorizations (such as QR or SVD) on them, without forcing the user to worry about index ordering. You can read the ITensor tutorials for more information or watch our talk.

# Why use the GPU

It seems pretty reasonable that you could expect a speedup for many tensor network algorithms by using a GPU. By permuting indices, it's possible to reduce all contraction operations to matrix-matrix or matrix-vector multiplications, at which the GPU excels. Most tensor network algorithms have runtimes dominated by such contractions or by SVD. However, there are some performance gotchas we always need to consider when using the GPU:

- The device has comparatiely low memory. The most expensive cards have 32GB of onboard RAM, which is a lot, but many state-of-the-art DMRG calculations require over a terabyte of RAM or checkpointing by writing intermediate information to disk.
- There's high latency and low bandwidth for memory transfers. If we absolutely
*have to*copy memory from the host CPU to the device, we should try to do it all in one big blob, and not in many small chunks. Although the GPU can overlap computations and memory transfers writing code to handle this can be a bit complex. - The performance for single precision floats is much better than for double precision. Although we'll probably see a perfomance boost for doubles, it won't be as dramatic as for single precision
`Float32`

.

In addition, there is a danger in the most naive approach to handling tensor contractions, which is to just permute all involved tensors into the index layout necessary to write the operation as a matrix-matrix multiplication, and then sit back and call `GEMM`

. Although in many cases this will work quite well, especially if the permutation doesn't involve many indices, there are plenty of bad cases where a great deal of time could be spent alllocating destination arrays and permuting source arrays into them. The risk of this increases with the average number of indices on each tensor (since there are more “bad” permutations available).

For these reasons, we weren't sure if the GPU would be a good choice for my current reseach project. You can read the physics details here. We have a C++ implementation of this code which is CPU only and runs about 5000 lines. One of my goals with the project was to eventually open source the code in the hope that others might find it useful or improve it. However, C++, despite being a great langauge, can be intimidating to many people. I have some experience writing C code that uses CUDA which, despite the powerful API and really granular control over the device the programmer is provided, can also be intimidating and require you to keep a lot of balls in the air while you're writing the code. But the C++ solution, stuck on a single node as it was, with all parallelism coming from CPU BLAS spread over 28 cores, was taking up to a week to run to get a decent picture of the converged result. This was pretty frustrating from a development perspective because it meant the debug cycle of “something's wrong” - “OK, think I found it” - “is this a fix?” - “nope, something's still wrong” had to take place over multiple days.

Since Miles (the original author of ITensor) and Matt were already thinking of rewriting ITensor in Julia (see our talk for the motivations for this decision), I decided I would try to help and maybe try to add some GPU support to the new package. Many tensor network algorithms, not only this one, are dominated by tensor-tensor contractions as mentioned above. And since I had already had some experience working with Julia's GPU programming/wrapping infrastructre in `CuArrays.jl`

, I thought it wouldn't be so hard to integrate a GPU based tensor operations backend to `ITensors.jl`

. (In fact, we sometimes want to add or subtract tensors, not just contract them.)

Our first approach, and one I don't have benchmarks for, was the naive method described above - just permute everything and call `CUBLAS`

's general matrix-matrix multiplication routine. In general, handling GPU memory with `CuArrays.jl`

was very easy. An `ITensor`

is essentially an opaque `Vector`

with some indices along for the ride, which tell you in what order to index elements of the `Vector`

. It's analogous to `CartesianArray`

for those who have used Julia's multidimensional array support. Since our algorithms usually require us to somehow achieve a contraction, QR decomposition, and addition, we thought treating the `ITensor`

storage as essentially a blob you can permute and give to multiplication API calls would be enough. Usually in these algorithms you're not often accessing or manipulating single elements or slices of the `ITensor`

(although this is possible to do and easy in both the C++ and Julia versions), just the tensors themselves.

Here's the sum total of the code I needed to get a barebones `cuITensor`

that you could move on and off the GPU:

```
function cuITensor(::Type{T},inds::IndexSet) where {T<:Number}
return ITensor(Dense{float(T)}(CuArrays.zeros(float(T),dim(inds))), inds)
end
cuITensor(::Type{T},inds::Index...) where {T<:Number} = ITensor(T,IndexSet(inds...))
cuITensor(is::IndexSet) = cuITensor(Float64,is)
cuITensor(inds::Index...) = cuITensor(IndexSet(inds...))
cuITensor() = ITensor()
function cuITensor(x::S, inds::IndexSet{N}) where {S<:Number, N}
dat = CuVector{float(S)}(undef, dim(inds))
fill!(dat, float(x))
ITensor(Dense{S}(dat), inds)
end
cuITensor(x::S, inds::Index...) where {S<:Number} = cuITensor(x,IndexSet(inds...))
function cuITensor(A::Array{S},inds::IndexSet) where {S<:Number}
return ITensor(Dense(CuArray{S}(A)), inds)
end
function cuITensor(A::CuArray{S},inds::IndexSet) where {S<:Number}
return ITensor(Dense(A), inds)
end
cuITensor(A::Array{S}, inds::Index...) where {S<:Number} = cuITensor(A,IndexSet(inds...))
cuITensor(A::CuArray{S}, inds::Index...) where {S<:Number} = cuITensor(A,IndexSet(inds...))
cuITensor(A::ITensor) = cuITensor(A.store.data,A.inds)
function Base.collect(A::ITensor)
typeof(A.store.data) <: CuArray && return ITensor(collect(A.store), A.inds)
return A
end
```

Mostly, this handles different ways of providing the indices, and a few options for the input data type. I assumed that if you called the `cuITensor`

constructor but gave it an input CPU array, you probably wanted that array transferred to the GPU.
That's the easy part. Adding support for some other operations, like QR decomposition or eigensolving, wasn't much harder:

```
function eigenHermitian(T::CuDenseTensor{ElT,2,IndsT};
kwargs...) where {ElT,IndsT}
maxdim::Int = get(kwargs,:maxdim,minimum(dims(T)))
mindim::Int = get(kwargs,:mindim,1)
cutoff::Float64 = get(kwargs,:cutoff,0.0)
absoluteCutoff::Bool = get(kwargs,:absoluteCutoff,false)
doRelCutoff::Bool = get(kwargs,:doRelCutoff,true)
local DM, UM
if ElT <: Complex
DM, UM = CUSOLVER.heevd!('V', 'U', matrix(T))
else
DM, UM = CUSOLVER.syevd!('V', 'U', matrix(T))
end
DM_ = reverse(DM)
truncerr, docut, DM = truncate!(DM_;maxdim=maxdim, cutoff=cutoff, absoluteCutoff=absoluteCutoff, doRelCutoff=doRelCutoff)
dD = length(DM)
dV = reverse(UM, dims=2)
if dD < size(dV,2)
dV = CuMatrix(dV[:,1:dD])
end
# Make the new indices to go onto U and V
u = eltype(IndsT)(dD)
v = eltype(IndsT)(dD)
Uinds = IndsT((ind(T,1),u))
Dinds = IndsT((u,v))
dV_ = CuArrays.zeros(ElT, length(dV))
copyto!(dV_, vec(dV))
U = Tensor(Dense(dV_),Uinds)
D = Tensor(Diag(real.(DM)),Dinds)
return U,D
end
```

This weird looking method of getting `dV`

into `dV_`

is necessary because of the way `CuArrays.jl`

deals with reshapes. As we'll see later it doesn't seem to impact performance too much.
But of course, the big problem we wanted to solve was contractions. Because the CPU code also works by performing the permutation and calling `GEMM`

, it was relatively easy to pirate that over to the GPU:

```
function contract!(C::CuArray{T},
p::CProps,
A::CuArray{T},
B::CuArray{T},
α::Tα=1.0,
β::Tβ=0.0) where {T,Tα<:Number,Tβ<:Number}
# bunch of code to find permutations and permute α and β goes here!
CUBLAS.gemm_wrapper!(cref, tA,tB,aref,bref,promote_type(T,Tα)(α),promote_type(T,Tβ)(β))
if p.permuteC
permutedims!(C,reshape(cref,p.newCrange...),p.PC)
end
return
end
```

The design of `ITensors.jl`

specifies that the `ITensor`

type itself is not specialized on its storage type, so that from the user's point of view, they have a tensor-like object contracting with another tensor-like object, and the developers can worry about how to multiply a diagonal-like rank-6 tensor by a sparse rank-4 tensor. This makes it easier for users to implement the algorithms they need to do their research in, and it's one of the library's strengths. All that was needed to allow an `ITensor`

with GPU-backed storage to play nicely with an `ITensor`

with CPU-backed storage was few lines of edge case handling:

```
function contract!!(R::CuDenseTensor{<:Number,NR}, labelsR::NTuple{NR}, T1::DenseTensor{<:Number,N1}, labelsT1::NTuple{N1}, T2::CuDenseTensor{<:Number,N2}, labelsT2::NTuple{N2}) where {NR,N1,N2}
return contract!!(R, CuDenseTensor(cu(store(T1)), inds(T1)), labelsT1, T2, labelsT2)
end
function contract!!(R::CuDenseTensor{<:Number,NR}, labelsR::NTuple{NR}, T1::CuDenseTensor{<:Number,N1}, labelsT1::NTuple{N1}, T2::DenseTensor{<:Number,N2}, labelsT2::NTuple{N2}) where {NR,N1,N2}
return contract!!(R, T1, labelsT1, CuDenseTensor(cu(store(T2)), inds(T2)), labelsT2)
end
```

I chose to copy the CPU storage to the device before the addition or contraction, hoping that this would occur rarely and that the performance gain in the main operation would offset the memory transfer time. Ideally this situation should never occur: we absolutely want to minimize memory transfers. However, if a user makes a mistake and forgets a `cuITensor(A)`

, their code won't error out. In fact, in the latest version of `ITensorsGPU.jl`

this dubious feature is disallowed, since in my code it was always the result of forgetting to initialize something on the GPU which should have been.

This was enough to get the barebones GPU support working. But I was still worried about the issue with the permutations, especially because the week-long simulations are those which are most memory intensive, and I was worried about running out of space on the device. Could there be a better solution?

`CUTENSOR`

and the story of how computers made my labour useless

Around this time we became aware of CUTENSOR, an NVIDIA library designed exactly for our use case: adding and contracting high rank tensors with indices in arbitrary order. However, this library was, of course, written in C. Luckily Julia makes it pretty easy to wrap C APIs, and we got started doing so in this epic PR to `CuArrays.jl`

. `CuArrays.jl`

already provides nice high- and low-level wrappers of CUDA C libraries in Julia, not only for dense or sparse linear algebra but also for random number generation and neural network primitives. So adding a multi-dimensional array library was a natural fit. During the process of getting these wrappers into a state fit for a public facing library, Tim Besard created some very nice scripts which automate the process of creating Julia wrappers for C functions, automating away many hours of labour I had performed years ago to get the sparse linear algebra and solver libraries working. Sic transit gloria mundi, I guess (generating these wrappers was not a glorious process). Now it will be easy for us to integrate changes to the `CUTENSOR`

API over time as more features are added.

`CUTENSOR`

's internals handle matching up elements for products and sums as part of the contraction process, so the permutations that `ITensors.jl`

performs for a CPU-based `ITensor`

are unnecessary. By overriding a few functions we're able to call the correct internal routines which feed through to `CUTENSOR`

:

```
function Base.:+(B::CuDenseTensor, A::CuDenseTensor)
opC = CUTENSOR.CUTENSOR_OP_IDENTITY
opA = CUTENSOR.CUTENSOR_OP_IDENTITY
opAC = CUTENSOR.CUTENSOR_OP_ADD
Ais = inds(A)
Bis = inds(B)
ind_dict = Vector{Index}()
for (idx, i) in enumerate(inds(A))
push!(ind_dict, i)
end
Adata = data(store(A))
Bdata = data(store(B))
reshapeBdata = reshape(Bdata,dims(Bis))
reshapeAdata = reshape(Adata,dims(Ais))
# probably a silly way to handle this, but it worked
ctainds = zeros(Int, length(Ais))
ctbinds = zeros(Int, length(Bis))
for (ii, ia) in enumerate(Ais)
ctainds[ii] = findfirst(x->x==ia, ind_dict)
end
for (ii, ib) in enumerate(Bis)
ctbinds[ii] = findfirst(x->x==ib, ind_dict)
end
ctcinds = copy(ctbinds)
C = CuArrays.zeros(eltype(Bdata), dims(Bis))
CUTENSOR.elementwiseBinary!(one(eltype(Adata)), reshapeAdata, ctainds, opA, one(eltype(Bdata)), reshapeBdata, ctbinds, opC, C, ctcinds, opAC)
copyto!(data(store(B)), vec(C))
return B
end
```

Once these wrappers and their tests were merged into `CuArrays.jl`

, I set about changing up how we were calling the contraction functions on the `ITensors.jl`

side. We decided to do this because within `CUTENSOR`

there were already highly optimized routines for various permutations, and we didn't want to try to reinvent the wheel with our permute-then-GEMM system. Switching to `CUTENSOR`

let us abstract away the permutation-handling, so the code interfacing with `CuArrays.jl`

was much simpler than under our previous approach. Dealing with optional dependencies, as `CuArrays.jl`

would have been for `ITensors.jl`

, is still kind of a pain in Julia, so I made a new package, `ITensorsGPU.jl`

, to hold all the CUDA-related logic. What's nice is that from the end-user's perspective, they just have copy the tensors to the GPU at the start of the simulation and afterwards everything works mostly seamlessly – they don't have to concern themselves with index orders or anything. It frees the user to focus more on high-level algorithm design.

# Extirpating memory copies

Copying memory back and forth from the device is extremely slow, and the code will perform best if we can eliminate as many as possible. One way to see how much time the device is spending on them is using NVIDIA's `nvprof`

tool. Working with the cluster means I usually do most of my development over SSH, so I used command line mode, which is really easy:

```
nvprof ~/software/julia/julia prof_run.jl
```

This generates some output about how much time the GPU spent doing various things, which is very long horizontally - scroll the box sideways if you can't see the function names:

```
==1386746== Profiling application: julia prof_run.jl
==1386746== Profiling result:
Type Time(%) Time Calls Avg Min Max Name
GPU activities: 14.02% 12.1800s 80792 150.76us 102.69us 408.00us void sytd2_upper_cta<double, double, int=5>(int, double*, int, double*, double*, double*)
7.44% 6.46842s 5077903 1.2730us 1.1190us 54.367us [CUDA memcpy HtoD]
6.02% 5.22885s 4430677 1.1800us 959ns 7.2640us ptxcall_anonymous19_1
5.06% 4.39301s 178968 24.546us 8.4480us 81.855us void cutensor_internal_namespace::tensor_contraction_kernel<cutensor_internal_namespace::tc_config_t<int=8, int=4, int=64, int=64, int=1, int=32, int=32, int=1, int=8, int=4, int=8, int=1, int=1, int=1, int=1, int=4, bool=1, bool=0, bool=0, bool=1, bool=0, cutensorOperator_t=1, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, bool=0>, double, double, double, double>(cutensor_internal_namespace::tc_params_t, int=1, int=4 const *, int=64 const *, cutensor_internal_namespace::tc_params_t, int=64*)
4.97% 4.31812s 346589 12.458us 5.4400us 29.088us void cutensor_internal_namespace::tensor_contraction_kernel<cutensor_internal_namespace::tc_config_t<int=8, int=4, int=64, int=64, int=1, int=16, int=16, int=1, int=8, int=4, int=4, int=1, int=1, int=1, int=1, int=4, bool=1, bool=1, bool=0, bool=0, bool=0, cutensorOperator_t=1, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, bool=0>, double, double, double, double>(cutensor_internal_namespace::tc_params_t, int=1, int=4 const *, int=64 const *, cutensor_internal_namespace::tc_params_t, int=64*)
4.96% 4.31109s 2397744 1.7970us 1.7270us 6.8160us ptxcall_setindex_kernel__26
4.18% 3.63043s 228988 15.854us 5.8560us 28.672us void cutensor_internal_namespace::tensor_contraction_kernel<cutensor_internal_namespace::tc_config_t<int=8, int=4, int=64, int=64, int=1, int=16, int=16, int=1, int=8, int=4, int=4, int=1, int=1, int=1, int=1, int=4, bool=1, bool=0, bool=0, bool=0, bool=0, cutensorOperator_t=1, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, bool=0>, double, double, double, double>(cutensor_internal_namespace::tc_params_t, int=1, int=4 const *, int=64 const *, cutensor_internal_namespace::tc_params_t, int=64*)
3.41% 2.96655s 2162240 1.3710us 1.2470us 3.4880us [CUDA memcpy DtoH]
2.77% 2.40280s 1827148 1.3150us 1.0870us 7.3600us ptxcall_anonymous19_14
2.70% 2.34219s 1327280 1.7640us 1.5360us 6.5920us ptxcall_setindex_kernel__15
5424,21 96%
...
```

You can see the `memcpyHtoD`

there, and it's taking up a lot of time! By carefully going through and creating `ITensors`

with the appropriate storage type in internal routines, like so:

```
# again probably a nicer way to do this
is_gpu = all([data(store(A[i,j])) isa CuArray for i in 1:Ny, j in 1:Nx)
N = spinI(findindex(A[row, col], "Site"); is_gpu=is_gpu)
function spinI(s::Index; is_gpu::Bool=false)::ITensor
I_data = is_gpu ? CuArrays.zeros(Float64, dim(s)*dim(s)) : zeros(Float64, dim(s), dim(s))
idi = diagind(reshape(I_data, dim(s), dim(s)), 0)
I_data[idi] = is_gpu ? CuArrays.ones(Float64, dim(s)) : ones(Float64, dim(s))
I = is_gpu ? cuITensor( I_data, IndexSet(s, s') ) : ITensor(vec(I_data), IndexSet(s, s'))
return I
end
```

it's possible to dramatically cut down on this, getting a final profiling report of

```
==3303697== Profiling application: julia prof_run.jl
==3303697== Profiling result:
Type Time(%) Time Calls Avg Min Max Name
GPU activities: 13.78% 19.2803s 343307 56.160us 15.328us 7.2638ms cutensor_internal_namespace::contraction_kernel(cutensor_internal_namespace::KernelParam_double_iden_1_2_false_false_double_iden_1_2_false_false_double_1_double_double_tb_128_128_8_simt_sm50_256)
13.47% 18.8464s 440326 42.800us 15.071us 653.75us cutensor_internal_namespace::contraction_kernel(cutensor_internal_namespace::KernelParam_double_iden_1_2_false_false_double_iden_1_2_true_false_double_1_double_double_tb_128_128_8_simt_sm50_256)
9.70% 13.5765s 262876 51.646us 15.360us 197.02us cutensor_internal_namespace::contraction_kernel(cutensor_internal_namespace::KernelParam_double_iden_1_2_true_false_double_iden_1_2_true_false_double_1_double_double_tb_128_128_8_simt_sm50_256)
8.90% 12.4562s 114666 108.63us 15.648us 4.7480ms cutensor_internal_namespace::contraction_kernel(cutensor_internal_namespace::KernelParam_double_iden_1_2_true_false_double_iden_1_2_false_false_double_1_double_double_tb_128_128_8_simt_sm50_256)
7.96% 11.1327s 305932 36.389us 22.559us 76.831us void gesvdbj_batch_32x16<double, double>(int, int const *, int const *, int const *, int, double const *, int, double, double*, double*, int*, double, int, double)
5.15% 7.21024s 2439080 2.9560us 2.0480us 6.8470us void ormtr_gerc<double, int=5, int=3, int=1>(int, double const *, int, int, double*, unsigned long, double const *, int, double const *)
3.62% 5.06010s 1663200 3.0420us 1.6630us 5.8240us void sytd2_symv_upper<double, int=4>(int, double const *, double const *, unsigned long, double const *, double*)
3.43% 4.80411s 2439080 1.9690us 1.5680us 5.4390us void ormtr_gemv_c<double, int=4>(int, int, double const *, unsigned long, double const *, int, double*)
2.84% 3.97216s 1636800 2.4260us 2.1120us 5.6000us void larfg_kernel_fast<double, double, int=6>(int, double*, double*, int, double*)
2.57% 3.59984s 2123225 1.6950us 864ns 90.111us [CUDA memcpy DtoD]
2.37% 3.31885s 1663200 1.9950us 1.1520us 4.5760us void sytd2_her2k_kernel<double, int=8, int=4>(int, double*, unsigned long, double const *, int, double const *)
2.25% 3.14692s 611864 5.1430us 4.6390us 12.704us void svd_column_rotate_batch_32x16<double, int=5, int=3>(int, int const *, int const *, int, double*, int, double*, int, double const *, int*)
2.17% 3.03782s 54611 55.626us 3.1990us 315.74us void cutensor_internal_namespace::reduction_kernel<bool=1, int=2, int=6, int=256, int=1, int=256, bool=1, cutensorOperator_t=1, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, double, double, double, double, double>(double, double const *, double const *, cutensor_internal_namespace::reduction_kernel<bool=1, int=2, int=6, int=256, int=1, int=256, bool=1, cutensorOperator_t=1, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, double, double, double, double, double>, double const *, double const **, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, cutensor_internal_namespace::reduction_params_t)
2.11% 2.95838s 29720 99.541us 30.367us 1.3328ms cutensor_internal_namespace::contraction_kernel(cutensor_internal_namespace::KernelParam_double_iden_1_2_true_false_double_iden_1_2_false_true_double_1_double_double_tb_128_128_8_simt_sm50_256)
2.08% 2.91405s 1663200 1.7520us 1.5040us 4.5440us void sytd2_compute_w_kernel<double, int=8, int=1>(double const *, int, double const *, double const *, int, double*)
1.83% 2.56571s 1384952 1.8520us 992ns 6.0160us [CUDA memcpy DtoH]
1.28% 1.79144s 305932 5.8550us 5.4070us 8.4800us void svd_row_rotate_batch_32x16<double>(int, int const *, int const *, int, double*, int, double const *, int*)
1.05% 1.47607s 1157165 1.2750us 831ns 8.8960us ptxcall_anonymous21_4
1.03% 1.44738s 62378 23.203us 10.560us 57.983us void geqr2_smem<double, double, int=8, int=6, int=4>(int, int, double*, unsigned long, double*, int)
0.72% 1.00992s 252101 4.0060us 1.5680us 805.62us void cutensor_internal_namespace::tensor_elementwise_kernel<cutensor_internal_namespace::ElementwiseConfig<unsigned int=1, int=128, unsigned int=64, unsigned int=2>, cutensor_internal_namespace::ElementwiseStaticOpPack<cutensorOperator_t=1, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t, cutensorOperator_t>, double, double, double, double>(cutensor_internal_namespace::ElementwiseParameters, int, int, cutensorOperator_t=1, unsigned int=64 const *, cutensor_internal_namespace::ElementwiseParameters, unsigned int=2 const *, cutensor_internal_namespace::ElementwiseParameters, cutensor_internal_namespace::ElementwiseConfig<unsigned int=1, int=128, unsigned int=64, unsigned int=2> const *, unsigned int=2 const **, bool, bool, bool, bool, cutensor_internal_namespace::ElementwiseOpPack)
0.71% 990.60ms 46039 21.516us 9.3440us 52.063us void geqr2_smem<double, double, int=8, int=5, int=5>(int, int, double*, unsigned long, double*, int)
0.54% 748.94ms 139738 5.3590us 895ns 11.776us void syevj_parallel_order_set_kernel<int=512>(int, int*)
0.51% 707.93ms 35644 19.861us 12.192us 34.368us void geqr2_smem<double, double, int=8, int=7, int=3>(int, int, double*, unsigned long, double*, int)
0.49% 684.11ms 542074 1.2620us 927ns 5.9200us copy_info_kernel(int, int*)
0.48% 669.92ms 298670 2.2420us 1.2470us 10.144us [CUDA memcpy HtoD]
```

This run was allowed to go on for longer to show that now the runtime on the GPU is dominated by useful work contracting tensors or computing factorizations.

# Head to head performance fight

With a real physics use case to test, the nice folks over at NVIDIA ran a performance comparison for us. These are some representative (actually, rather small) simulations using our code, and you can see that the GPU acceleration is helping a lot.

The `x`

-axis in this figure is a little opaque - it refers to varying problem sizes we might want to simulate, specifically an `N`

by `N`

lattice with a maximum tensor interior dimension of `m`

- thus, `N - m`

in the labels.
And now I'm able to run simulations that used to take a week in an hour thanks to the GPU acceleration, with no loss of accuracy so far. We've had a big success using the GPU, and so far haven't run up against the device memory limit. Having a Julia code makes it easier to maintain and easier to reason about new code as it's written. We're looking forward to using this backend to accelerate other tensor network algorithms and make it quicker to test out new ideas.

# Some Takeaways

- The
`CuArrays.jl`

package makes it pretty easy and pleasant to integrate the GPU into your codebase. However, there are still some pain points (in the sparse linear algebra code especially) if someone is looking for a project to contribute to. - It's important to pick a problem like this that is about 95% of the way to the perfect problem for the GPU if you want to brag without having to do much work.
- You should actually check to make sure that you aren't copying memory back and forth when you don't need to. If you are copying more than you think you should, you can try to figure out where it's coming from by inserting a
`stacktrace`

call into the`cudaMemcpy`

calls at the`CUDAdrv.jl`

package. That should tell you the location up the call stack in your code where the copy to/from the device is triggered.

There were several things that made this transition successful:

- The simplicity of writing Julia wrappers for C code (this let us quickly interface with
`CUTENSOR`

) - The existing ease of use of
`CuArrays.jl`

- now it's easy to extend this to a new library, and it made writing the constructors for our own types quite straightforward - The fact that our problem was well-suited for the GPU

It's certainly true that we could have achieved the same, or possibly better, had we modified the C++ ITensor code to use the GPU. But I think it's fair to say it would have taken much more time, and would have been less accessible to other people in condensed matter physics. We were willing to settle for a slightly less than optimal speedup if the code to achieve it got written at all.

If you want to look at some of the unpleasant internals of this, feel free to check out `ITensorsGPU.jl`

and try things out for yourself. If you're interested in learning more about tensor network algorithms, check out Miles’ site here.