Merge branch 'main' into kmp5/debug/issue_1450

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.3.8"
+version = "0.3.12"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -51,11 +51,14 @@ NDTensorsTBLISExt = "TBLIS"
 NDTensorscuTENSORExt = "cuTENSOR"
 
 [compat]
+AMDGPU = "0.9"
 Accessors = "0.1.33"
 Adapt = "3.7, 4"
 ArrayLayouts = "1.4"
 BlockArrays = "0.16"
+CUDA = "5"
 Compat = "4.9"
+cuTENSOR = "2"
 Dictionaries = "0.4"
 EllipsisNotation = "1.8"
 FillArrays = "1"
@@ -68,6 +71,8 @@ InlineStrings = "1"
 LinearAlgebra = "1.6"
 MacroTools = "0.5"
 MappedArrays = "0.4"
+Metal = "1"
+Octavian = "0.3"
 PackageExtensionCompat = "1"
 Random = "1.6"
 SimpleTraits = "0.9.4"
@@ -76,10 +81,17 @@ SplitApplyCombine = "1.2.2"
 StaticArrays = "0.12, 1.0"
 Strided = "2"
 StridedViews = "0.2.2"
+TBLIS = "0.2"
 TimerOutputs = "0.5.5"
 TupleTools = "1.2.0"
 VectorInterface = "0.4.2"
 julia = "1.6"
 
 [extras]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
+TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9"
+cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl

@@ -1,5 +1,6 @@
 using ArrayLayouts: LayoutArray
 using BlockArrays: blockisequal
+using LinearAlgebra: Adjoint, Transpose
 using ..SparseArrayInterface:
   SparseArrayInterface,
   SparseArrayStyle,
@@ -73,6 +74,20 @@ function Base.copyto!(a_dest::LayoutArray, a_src::BlockSparseArrayLike)
   return a_dest
 end
 
+function Base.copyto!(
+  a_dest::AbstractMatrix, a_src::Transpose{T,<:AbstractBlockSparseMatrix{T}}
+) where {T}
+  sparse_copyto!(a_dest, a_src)
+  return a_dest
+end
+
+function Base.copyto!(
+  a_dest::AbstractMatrix, a_src::Adjoint{T,<:AbstractBlockSparseMatrix{T}}
+) where {T}
+  sparse_copyto!(a_dest, a_src)
+  return a_dest
+end
+
 function Base.permutedims!(a_dest, a_src::BlockSparseArrayLike, perm)
   sparse_permutedims!(a_dest, a_src, perm)
   return a_dest

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -3,8 +3,8 @@ using SplitApplyCombine: groupcount
 
 using Adapt: Adapt, WrappedArray
 
-const WrappedAbstractBlockSparseArray{T,N,A} = WrappedArray{
-  T,N,<:AbstractBlockSparseArray,<:AbstractBlockSparseArray{T,N}
+const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{
+  T,N,AbstractBlockSparseArray,AbstractBlockSparseArray{T,N}
 }
 
 # TODO: Rename `AnyBlockSparseArray`.

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearray/blocksparsearray.jl

@@ -10,7 +10,7 @@ struct BlockSparseArray{
   N,
   A<:AbstractArray{T,N},
   Blocks<:AbstractArray{A,N},
-  Axes<:NTuple{N,AbstractUnitRange{Int}},
+  Axes<:Tuple{Vararg{AbstractUnitRange,N}},
 } <: AbstractBlockSparseArray{T,N}
   blocks::Blocks
   axes::Axes

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -8,6 +8,7 @@ using BlockArrays:
   blocks,
   blocklengths,
   findblockindex
+using LinearAlgebra: Adjoint, Transpose
 using ..SparseArrayInterface: perm, iperm, nstored
 ## using MappedArrays: mappedarray
 
@@ -86,35 +87,96 @@ end
 
 # BlockArrays
 
-using ..SparseArrayInterface: SparseArrayInterface, AbstractSparseArray
+using ..SparseArrayInterface:
+  SparseArrayInterface, AbstractSparseArray, AbstractSparseMatrix
 
-# Represents the array of arrays of a `SubArray`
-# wrapping a block spare array, i.e. `blocks(array)` where `a` is a `SubArray`.
+_perm(::PermutedDimsArray{<:Any,<:Any,P}) where {P} = P
+_getindices(t::Tuple, indices) = map(i -> t[i], indices)
+_getindices(i::CartesianIndex, indices) = CartesianIndex(_getindices(Tuple(i), indices))
+
+# Represents the array of arrays of a `PermutedDimsArray`
+# wrapping a block spare array, i.e. `blocks(array)` where `a` is a `PermutedDimsArray`.
 struct SparsePermutedDimsArrayBlocks{T,N,Array<:PermutedDimsArray{T,N}} <:
        AbstractSparseArray{T,N}
   array::Array
 end
 function blocksparse_blocks(a::PermutedDimsArray)
   return SparsePermutedDimsArrayBlocks(a)
 end
-_perm(::PermutedDimsArray{<:Any,<:Any,P}) where {P} = P
-_getindices(t::Tuple, indices) = map(i -> t[i], indices)
-_getindices(i::CartesianIndex, indices) = CartesianIndex(_getindices(Tuple(i), indices))
-function SparseArrayInterface.stored_indices(a::SparsePermutedDimsArrayBlocks)
-  return map(I -> _getindices(I, _perm(a.array)), stored_indices(blocks(parent(a.array))))
-end
 function Base.size(a::SparsePermutedDimsArrayBlocks)
   return _getindices(size(blocks(parent(a.array))), _perm(a.array))
 end
-function Base.getindex(a::SparsePermutedDimsArrayBlocks, index::Vararg{Int})
+function Base.getindex(
+  a::SparsePermutedDimsArrayBlocks{<:Any,N}, index::Vararg{Int,N}
+) where {N}
   return PermutedDimsArray(
     blocks(parent(a.array))[_getindices(index, _perm(a.array))...], _perm(a.array)
   )
 end
+function SparseArrayInterface.stored_indices(a::SparsePermutedDimsArrayBlocks)
+  return map(I -> _getindices(I, _perm(a.array)), stored_indices(blocks(parent(a.array))))
+end
+# TODO: Either make this the generic interface or define
+# `SparseArrayInterface.sparse_storage`, which is used
+# to defined this.
+SparseArrayInterface.nstored(a::SparsePermutedDimsArrayBlocks) = length(stored_indices(a))
 function SparseArrayInterface.sparse_storage(a::SparsePermutedDimsArrayBlocks)
   return error("Not implemented")
 end
 
+reverse_index(index) = reverse(index)
+reverse_index(index::CartesianIndex) = CartesianIndex(reverse(Tuple(index)))
+
+# Represents the array of arrays of a `Transpose`
+# wrapping a block spare array, i.e. `blocks(array)` where `a` is a `Transpose`.
+struct SparseTransposeBlocks{T,Array<:Transpose{T}} <: AbstractSparseMatrix{T}
+  array::Array
+end
+function blocksparse_blocks(a::Transpose)
+  return SparseTransposeBlocks(a)
+end
+function Base.size(a::SparseTransposeBlocks)
+  return reverse(size(blocks(parent(a.array))))
+end
+function Base.getindex(a::SparseTransposeBlocks, index::Vararg{Int,2})
+  return transpose(blocks(parent(a.array))[reverse(index)...])
+end
+function SparseArrayInterface.stored_indices(a::SparseTransposeBlocks)
+  return map(reverse_index, stored_indices(blocks(parent(a.array))))
+end
+# TODO: Either make this the generic interface or define
+# `SparseArrayInterface.sparse_storage`, which is used
+# to defined this.
+SparseArrayInterface.nstored(a::SparseTransposeBlocks) = length(stored_indices(a))
+function SparseArrayInterface.sparse_storage(a::SparseTransposeBlocks)
+  return error("Not implemented")
+end
+
+# Represents the array of arrays of a `Adjoint`
+# wrapping a block spare array, i.e. `blocks(array)` where `a` is a `Adjoint`.
+struct SparseAdjointBlocks{T,Array<:Adjoint{T}} <: AbstractSparseMatrix{T}
+  array::Array
+end
+function blocksparse_blocks(a::Adjoint)
+  return SparseAdjointBlocks(a)
+end
+function Base.size(a::SparseAdjointBlocks)
+  return reverse(size(blocks(parent(a.array))))
+end
+function Base.getindex(a::SparseAdjointBlocks, index::Vararg{Int,2})
+  return blocks(parent(a.array))[reverse(index)...]'
+end
+function SparseArrayInterface.stored_indices(a::SparseAdjointBlocks)
+  return map(reverse_index, stored_indices(blocks(parent(a.array))))
+end
+# TODO: Either make this the generic interface or define
+# `SparseArrayInterface.sparse_storage`, which is used
+# to defined this.
+SparseArrayInterface.nstored(a::SparseAdjointBlocks) = length(stored_indices(a))
+function SparseArrayInterface.sparse_storage(a::SparseAdjointBlocks)
+  return error("Not implemented")
+end
+
 # TODO: Move to `BlockArraysExtensions`.
 # This takes a range of indices `indices` of array `a`
 # and maps it to the range of indices within block `block`.
@@ -167,9 +229,6 @@ end
 function Base.size(a::SparseSubArrayBlocks)
   return length.(axes(a))
 end
-function SparseArrayInterface.stored_indices(a::SparseSubArrayBlocks)
-  return stored_indices(view(blocks(parent(a.array)), axes(a)...))
-end
 function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::CartesianIndex{N}) where {N}
   return a[Tuple(I)...]
 end
@@ -192,6 +251,13 @@ function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) whe
   # TODO: Implement this properly.
   return true
 end
+function SparseArrayInterface.stored_indices(a::SparseSubArrayBlocks)
+  return stored_indices(view(blocks(parent(a.array)), axes(a)...))
+end
+# TODO: Either make this the generic interface or define
+# `SparseArrayInterface.sparse_storage`, which is used
+# to defined this.
+SparseArrayInterface.nstored(a::SparseSubArrayBlocks) = length(stored_indices(a))
 function SparseArrayInterface.sparse_storage(a::SparseSubArrayBlocks)
   return error("Not implemented")
 end