Blocksparse CuTensor contraction backend

[kmp5VT]

No description provided.

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please run the ITensorFormatter to apply these changes.

Click here to view the suggested changes.

diff --git a/NDTensors/ext/NDTensorscuTENSORExt/contract.jl b/NDTensors/ext/NDTensorscuTENSORExt/contract.jl
index b50c220c..2f19b1cd 100644
--- a/NDTensors/ext/NDTensorscuTENSORExt/contract.jl
+++ b/NDTensors/ext/NDTensorscuTENSORExt/contract.jl
@@ -1,7 +1,6 @@
 using Base: ReshapedArray
 using NDTensors.Expose: Exposed, expose, unexpose
-using NDTensors: NDTensors, BlockSparseTensor, DenseTensor, array,
-blockdims, data, eachnzblock, inds, nblocks, nzblocks
+using NDTensors: NDTensors, BlockSparseTensor, DenseTensor, array, blockdims, data, eachnzblock, inds, nblocks, nzblocks
 using cuTENSOR: cuTENSOR, CuArray, CuTensor
 
 # Handle cases that can't be handled by `cuTENSOR.jl`
@@ -32,10 +31,11 @@ function to_cuTensorBS(T::BlockSparseTensor)
     nzblock_coords_t1 = [Int64.(x.data) for x in nzblocks(T)]
     block_per_mode_t1 = length.(block_extents_t1)
     is = [i for i in 1:ndims(T)]
-    return cuTENSOR.CuTensorBS(blocks_t1, block_per_mode_t1, block_extents_t1, nzblock_coords_t1, is);
+    return cuTENSOR.CuTensorBS(blocks_t1, block_per_mode_t1, block_extents_t1, nzblock_coords_t1, is)
 end
 
-function NDTensors._contract!(R::Exposed{<:CuArray, <:BlockSparseTensor},
+function NDTensors._contract!(
+        R::Exposed{<:CuArray, <:BlockSparseTensor},
         labelsR,
         tensor1::Exposed{<:CuArray, <:BlockSparseTensor},
         labelstensor1,
@@ -44,9 +44,9 @@ function NDTensors._contract!(R::Exposed{<:CuArray, <:BlockSparseTensor},
         grouped_contraction_plan,
         executor,
     )
-    N1 = ndims(unexpose(tensor1)) 
-    N2 = ndims(unexpose(tensor2)) 
-    NR = ndims(unexpose(R)) 
+    N1 = ndims(unexpose(tensor1))
+    N2 = ndims(unexpose(tensor2))
+    NR = ndims(unexpose(R))
     if NDTensors.using_CuTensorBS() && (N1 > 0) && (N2 > 0) && (NR > 0)
         # println("Using new function")
         cuR = ITensor_to_cuTensorBS(unexpose(R))
@@ -61,14 +61,14 @@ function NDTensors._contract!(R::Exposed{<:CuArray, <:BlockSparseTensor},
         return R
     else
         return NDTensors._contract!(
-        unexpose(R),
-        labelsR,
-        unexpose(tensor1),
-        labelstensor1,
-        unexpose(tensor2),
-        labelstensor2,
-        grouped_contraction_plan,
-        executor,
+            unexpose(R),
+            labelsR,
+            unexpose(tensor1),
+            labelstensor1,
+            unexpose(tensor2),
+            labelstensor2,
+            grouped_contraction_plan,
+            executor,
         )
     end
 end
diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl
index 0323f2c7..4936a30f 100644
--- a/NDTensors/src/NDTensors.jl
+++ b/NDTensors/src/NDTensors.jl
@@ -256,7 +256,6 @@ end
 
 function backend_octavian end
 
-
 _using_CuTensorBS = false
 
 using_CuTensorBS() = _using_CuTensorBS
diff --git a/NDTensors/src/blocksparse/contract_generic.jl b/NDTensors/src/blocksparse/contract_generic.jl
index 39b67fac..97afe393 100644
--- a/NDTensors/src/blocksparse/contract_generic.jl
+++ b/NDTensors/src/blocksparse/contract_generic.jl
@@ -71,19 +71,21 @@ function contract!(
     )
     return R
 end
-function _contract!(R::Exposed,
+function _contract!(
+        R::Exposed,
         labelsR,
         tensor1::Exposed,
         labelstensor1,
         tensor2::Exposed,
         labelstensor2,
         grouped_contraction_plan,
-        executor,
+        executor
     )
-    _contract!(unexpose(R), labelsR, 
-    unexpose(tensor1), labelstensor1,
-    unexpose(tensor2), labelstensor2,
-    grouped_contraction_plan,executor
+    return _contract!(
+        unexpose(R), labelsR,
+        unexpose(tensor1), labelstensor1,
+        unexpose(tensor2), labelstensor2,
+        grouped_contraction_plan, executor
     )
 end
 # Function barrier to improve type stability,

[mtfishman]

Great to see this, thanks @kmp5VT. I guess this relies on JuliaGPU/CUDA.jl#3057? Once that is merged, would we just need to install the latest version of cuTENSOR/cuTENSOR.jl and the new backend in this PR "just works"?

[emstoudenmire]

Looks nice and very minimal. Thanks Karl!

[kmp5VT]

@mtfishman In NDTensors I also added a internal variable _cutensor_blocksparse which you can enable/disable during runtime. So if the variable is enabled plus you have the changes in my cuda branch then it will automatically work. And from my tests, the new backend produces noticeable speedups over the previous blocksparse GPU backends

[kmp5VT]

Hi @mtfishman @emstoudenmire , the blocksparse cutensor wrapper is now merged into CUDA.jl main so this feature is no longer experimental. The biggest question I have is, is the current implementation (i.e. introducing a blocksparse cutensor global variable in NDTensors) sufficient or do you have an alternative implementation in mind? I would guess that we would want to remove the legacy canonical cutensor code (for blocksparse arrays) after benchmark tests have been collected.

[mtfishman]

Hi Karl, great to hear it, thanks for getting that implemented! We're very excited to try this out.

It's a good question. I think we should introduce this as "opt in" at first so we can have time to test it out, then if we're convinced it is either on par or better than the current implementation (which I'm guessing it is) we can make it the default when cuTENSOR is loaded. I'm thinking of trying out a new design based on ScopedValues for selecting algorithms/backends, would you mind if I made a PR based on this one testing out that design?

[kmp5VT]

@mtfishman No I don't mind at all. I can also give you access to this branch if you want to add it here. Let me know!

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 4.89%. Comparing base (ee2ebda) to head (ff5033c).

❗ There is a different number of reports uploaded between BASE (ee2ebda) and HEAD (ff5033c). Click for more details.

HEAD has 7 uploads less than BASE

Flag	BASE (ee2ebda)	HEAD (ff5033c)
docs	2	1
	12	6

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #1721       +/-   ##
==========================================
- Coverage   81.02%   4.89%   -76.13%     
==========================================
  Files          80      69       -11     
  Lines        5074    5021       -53     
==========================================
- Hits         4111     246     -3865     
- Misses        963    4775     +3812     

Flag	Coverage Δ
docs	4.92% <ø> (-0.07%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.