[NDTensors] Fix contracting dense with diag on GPU

[kmp5VT]

Description

In reference to this bug report

This bug is two parted. First I am working to make it possible to call tr on GPU based Tensors. The issue here is a scalar indexing problem where the code tries to getdiagindex of the tensor. For now I have tried to use @allowscalar and expose to solve the issue. Timing wise I have found CUDA based tr with @allowscalar is roughly 2x faster than CPU based trace.

The second bug is that delta functions are not being properly constructed on GPU. This is, impart, due to delta being a UniformDiag which does not carry information about where the diag will be allocated and assumes CPU. I am not sure yet how to fix this problem. Potentially we could replace the datatype of UniformDiag from <:Number to <:UnallocatedFill{<:Number}

Checklist:

• Trace works on CPU and GPU.
• Performance of trace on GPU is equal to or better than CPU
• Delta functions can be constructed on GPU
• All current unittests pass
• Unittests are made for the failing cases which prompted this PR.
• Ensure that all solutions also work with quantum numbers and block sparsity

[mtfishman]

I am not sure yet how to fix this problem. Potentially we could replace the datatype of UniformDiag from <:Number to <:UnallocatedFill{<:Number}

I would prefer not going that route since that is a much more involved change that would likely require rewriting a lot of the UniformDiag code (which is best left for the new DiagonalArrays design). It would be better to try to solve it in a more narrow way.

[kmp5VT]

@mtfishman So I am pushing an idea on how to fix the issue with \delta it involves reworking the dense function to look like this

dense(T::Tensor) = dense(unwrap_array_type(T), T)
dense(datat::Type{<:AbstractArray}, T::Tensor) = setstorage(T, adapt(datat, dense(storage(T))))

Since the dense for Diag already looks like this

function dense(T::DiagTensor)
  return dense(unwrap_array_type(T), T)
end

Update* I just checked and the code I pushed fixes both of the errors in the bug report on metal. Testing the other backends now.

[mtfishman]

What about using expose for that dispatch?

[kmp5VT]

@mtfishman For the dispatch of the dense function?

[mtfishman]

@mtfishman For the dispatch of the dense function?

Yes.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 44.14%. Comparing base (e08e131) to head (0dec0e5).
Report is 7 commits behind head on main.

❗ Current head 0dec0e5 differs from pull request most recent head 422159f

Please upload reports for the commit 422159f to get more accurate results.

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1453      +/-   ##
==========================================
+ Coverage   43.65%   44.14%   +0.49%     
==========================================
  Files         136      144       +8     
  Lines        8806     9374     +568     
==========================================
+ Hits         3844     4138     +294     
- Misses       4962     5236     +274     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mtfishman]

@kmp5VT this looks good to me.

The only remaining thing I can see to do is fix tensor contractions involving block sparse tensors with diagonal blocks. What's the status of that? I think we can fix that in a follow-up PR, I guess that can be fixed in a similar way (say by calling denseblocks and then moving the tensor to GPU).

[kmp5VT]

@kmp5VT this looks good to me.

The only remaining thing I can see to do is fix tensor contractions involving block sparse tensors with diagonal blocks. What's the status of that? I think we can fix that in a follow-up PR, I guess that can be fixed in a similar way (say by calling denseblocks and then moving the tensor to GPU).

@mtfishman Thats a good point. I hadn't checked the blocksparse implementation. I ran this code

using ITensorMPS: MPO, siteinds
using LinearAlgebra: tr
using Metal: mtl
s = siteinds("S=1/2", 2; conserve_qns=true)
o = MPO(s, "Id")
tr(mtl(o))

and it currently fails because of scalar indexing (because its fed into the cpu dense * diag code). I found that in diagblocksparse I can expose the block calls to contract which fixes the behavior

[mtfishman]

Great, glad to see it was a simple fix, simpler than I was picturing.

Seems like the only thing left is to add a test for the block sparse case, after that is this good to go?

[mtfishman]

Besides the final comment, this looks good, thanks.