Skip to content

Commit

Permalink
[NDTensors] HDF5 package extension (#1390)
Browse files Browse the repository at this point in the history
  • Loading branch information
@emstoudenmire
emstoudenmire authored Apr 17, 2024
1 parent b56184c commit fb77a8f
Show file tree
Hide file tree
Showing 11 changed files with 180 additions and 160 deletions.
8 changes: 6 additions & 2 deletions NDTensors/Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "NDTensors"
uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
authors = ["Matthew Fishman <[email protected]>"]
version = "0.3.0"
version = "0.3.1"

[deps]
Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
Expand All @@ -15,7 +15,6 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721"
InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
Expand All @@ -34,13 +33,15 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"

[weakdeps]
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"
AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"

[extensions]
NDTensorsCUDAExt = "CUDA"
NDTensorsHDF5Ext = "HDF5"
NDTensorsMetalExt = "Metal"
NDTensorsOctavianExt = "Octavian"
NDTensorsTBLISExt = "TBLIS"
Expand Down Expand Up @@ -75,3 +76,6 @@ TimerOutputs = "0.5.5"
TupleTools = "1.2.0"
VectorInterface = "0.4.2"
julia = "1.6"

[extras]
HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
8 changes: 8 additions & 0 deletions NDTensors/ext/NDTensorsHDF5Ext/NDTensorsHDF5Ext.jl
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
module NDTensorsHDF5Ext

include("blocksparse.jl")
include("dense.jl")
include("diag.jl")
include("empty.jl")

end # module NDTensorsHDF5Ext
68 changes: 68 additions & 0 deletions NDTensors/ext/NDTensorsHDF5Ext/blocksparse.jl
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
using HDF5: HDF5, attributes, create_group, open_group, read, write
using NDTensors: data, Block, blockoffsets, BlockOffsets, BlockSparse

# Helper function for HDF5 write/read of BlockSparse
function offsets_to_array(boff::BlockOffsets{N}) where {N}
nblocks = length(boff)
asize = (N + 1) * nblocks
n = 1
a = Vector{Int}(undef, asize)
for bo in pairs(boff)
for j in 1:N
a[n] = bo[1][j]
n += 1
end
a[n] = bo[2]
n += 1
end
return a
end

# Helper function for HDF5 write/read of BlockSparse
function array_to_offsets(a, N::Int)
asize = length(a)
nblocks = div(asize, N + 1)
boff = BlockOffsets{N}()
j = 0
for b in 1:nblocks
insert!(boff, Block(ntuple(i -> (a[j + i]), N)), a[j + N + 1])
j += (N + 1)
end
return boff
end

function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::String, B::BlockSparse)
g = create_group(parent, name)
attributes(g)["type"] = "BlockSparse{$(eltype(B))}"
attributes(g)["version"] = 1
if eltype(B) != Nothing
write(g, "ndims", ndims(B))
write(g, "data", data(B))
off_array = offsets_to_array(blockoffsets(B))
write(g, "offsets", off_array)
end
end

function HDF5.read(
parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Store}
) where {Store<:BlockSparse}
g = open_group(parent, name)
ElT = eltype(Store)
typestr = "BlockSparse{$ElT}"
if read(attributes(g)["type"]) != typestr
error("HDF5 group or file does not contain $typestr data")
end
N = read(g, "ndims")
off_array = read(g, "offsets")
boff = array_to_offsets(off_array, N)
# Attribute __complex__ is attached to the "data" dataset
# by the h5 library used by C++ version of ITensor:
if haskey(attributes(g["data"]), "__complex__")
M = read(g, "data")
nelt = size(M, 1) * size(M, 2)
data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
else
data = read(g, "data")
end
return BlockSparse(data, boff)
end
37 changes: 37 additions & 0 deletions NDTensors/ext/NDTensorsHDF5Ext/dense.jl
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
using HDF5: HDF5, attributes, create_group, open_group, read, write
using NDTensors: Dense

function HDF5.write(
parent::Union{HDF5.File,HDF5.Group}, name::String, D::Store
) where {Store<:Dense}
g = create_group(parent, name)
attributes(g)["type"] = "Dense{$(eltype(Store))}"
attributes(g)["version"] = 1
if eltype(D) != Nothing
write(g, "data", D.data)
end
end

function HDF5.read(
parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Store}
) where {Store<:Dense}
g = open_group(parent, name)
ElT = eltype(Store)
typestr = "Dense{$ElT}"
if read(attributes(g)["type"]) != typestr
error("HDF5 group or file does not contain $typestr data")
end
if ElT == Nothing
return Dense{Nothing}()
end
# Attribute __complex__ is attached to the "data" dataset
# by the h5 library used by C++ version of ITensor:
if haskey(attributes(g["data"]), "__complex__")
M = read(g, "data")
nelt = size(M, 1) * size(M, 2)
data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
else
data = read(g, "data")
end
return Dense{ElT}(data)
end
38 changes: 38 additions & 0 deletions NDTensors/ext/NDTensorsHDF5Ext/diag.jl
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
using HDF5: HDF5, attributes, create_group, open_group, read, write
using NDTensors: datatype, Dense, Diag

function HDF5.write(
parent::Union{HDF5.File,HDF5.Group}, name::String, D::Store
) where {Store<:Diag}
g = create_group(parent, name)
attributes(g)["type"] = "Diag{$(eltype(Store)),$(datatype(Store))}"
attributes(g)["version"] = 1
if eltype(D) != Nothing
write(g, "data", D.data)
end
end

function HDF5.read(
parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Store}
) where {Store<:Diag}
g = open_group(parent, name)
ElT = eltype(Store)
DataT = datatype(Store)
typestr = "Diag{$ElT,$DataT}"
if read(attributes(g)["type"]) != typestr
error("HDF5 group or file does not contain $typestr data")
end
if ElT == Nothing
return Dense{Nothing}()
end
# Attribute __complex__ is attached to the "data" dataset
# by the h5 library used by C++ version of ITensor:
if haskey(attributes(g["data"]), "__complex__")
M = read(g, "data")
nelt = size(M, 1) * size(M, 2)
data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
else
data = read(g, "data")
end
return Diag{ElT,DataT}(data)
end
23 changes: 23 additions & 0 deletions NDTensors/ext/NDTensorsHDF5Ext/empty.jl
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
using HDF5: HDF5, attributes, create_group, open_group, read, write
using NDTensors: EmptyStorage

# XXX: this seems a bit strange and fragile?
# Takes the type very literally.
function HDF5.read(
parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{StoreT}
) where {StoreT<:EmptyStorage}
g = open_group(parent, name)
typestr = string(StoreT)
if read(attributes(g)["type"]) != typestr
error("HDF5 group or file does not contain $typestr data")
end
return StoreT()
end

function HDF5.write(
parent::Union{HDF5.File,HDF5.Group}, name::String, ::StoreT
) where {StoreT<:EmptyStorage}
g = create_group(parent, name)
attributes(g)["type"] = string(StoreT)
return attributes(g)["version"] = 1
end
66 changes: 0 additions & 66 deletions NDTensors/src/blocksparse/blocksparse.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -205,69 +205,3 @@ end
# end
# return R
#end

# Helper function for HDF5 write/read of BlockSparse
function offsets_to_array(boff::BlockOffsets{N}) where {N}
nblocks = length(boff)
asize = (N + 1) * nblocks
n = 1
a = Vector{Int}(undef, asize)
for bo in pairs(boff)
for j in 1:N
a[n] = bo[1][j]
n += 1
end
a[n] = bo[2]
n += 1
end
return a
end

# Helper function for HDF5 write/read of BlockSparse
function array_to_offsets(a, N::Int)
asize = length(a)
nblocks = div(asize, N + 1)
boff = BlockOffsets{N}()
j = 0
for b in 1:nblocks
insert!(boff, Block(ntuple(i -> (a[j + i]), N)), a[j + N + 1])
j += (N + 1)
end
return boff
end

function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::String, B::BlockSparse)
g = create_group(parent, name)
attributes(g)["type"] = "BlockSparse{$(eltype(B))}"
attributes(g)["version"] = 1
if eltype(B) != Nothing
write(g, "ndims", ndims(B))
write(g, "data", data(B))
off_array = offsets_to_array(blockoffsets(B))
write(g, "offsets", off_array)
end
end

function HDF5.read(
parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Store}
) where {Store<:BlockSparse}
g = open_group(parent, name)
ElT = eltype(Store)
typestr = "BlockSparse{$ElT}"
if read(attributes(g)["type"]) != typestr
error("HDF5 group or file does not contain $typestr data")
end
N = read(g, "ndims")
off_array = read(g, "offsets")
boff = array_to_offsets(off_array, N)
# Attribute __complex__ is attached to the "data" dataset
# by the h5 library used by C++ version of ITensor:
if haskey(attributes(g["data"]), "__complex__")
M = read(g, "data")
nelt = size(M, 1) * size(M, 2)
data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
else
data = read(g, "data")
end
return BlockSparse(data, boff)
end
35 changes: 0 additions & 35 deletions NDTensors/src/dense/dense.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -141,38 +141,3 @@ end
function convert(::Type{<:Dense{ElR,DataT}}, D::Dense) where {ElR,DataT}
return Dense(convert(DataT, data(D)))
end

function HDF5.write(
parent::Union{HDF5.File,HDF5.Group}, name::String, D::Store
) where {Store<:Dense}
g = create_group(parent, name)
attributes(g)["type"] = "Dense{$(eltype(Store))}"
attributes(g)["version"] = 1
if eltype(D) != Nothing
write(g, "data", D.data)
end
end

function HDF5.read(
parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Store}
) where {Store<:Dense}
g = open_group(parent, name)
ElT = eltype(Store)
typestr = "Dense{$ElT}"
if read(attributes(g)["type"]) != typestr
error("HDF5 group or file does not contain $typestr data")
end
if ElT == Nothing
return Dense{Nothing}()
end
# Attribute __complex__ is attached to the "data" dataset
# by the h5 library used by C++ version of ITensor:
if haskey(attributes(g["data"]), "__complex__")
M = read(g, "data")
nelt = size(M, 1) * size(M, 2)
data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
else
data = read(g, "data")
end
return Dense{ElT}(data)
end
36 changes: 0 additions & 36 deletions NDTensors/src/diag/diag.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,39 +132,3 @@ end
# Convert a Diag storage type to the closest Dense storage type
dense(::Type{<:NonuniformDiag{ElT,DataT}}) where {ElT,DataT} = Dense{ElT,DataT}
dense(::Type{<:UniformDiag{ElT}}) where {ElT} = Dense{ElT,default_datatype(ElT)}

function HDF5.write(
parent::Union{HDF5.File,HDF5.Group}, name::String, D::Store
) where {Store<:Diag}
g = create_group(parent, name)
attributes(g)["type"] = "Diag{$(eltype(Store)),$(datatype(Store))}"
attributes(g)["version"] = 1
if eltype(D) != Nothing
write(g, "data", D.data)
end
end

function HDF5.read(
parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Store}
) where {Store<:Diag}
g = open_group(parent, name)
ElT = eltype(Store)
DataT = datatype(Store)
typestr = "Diag{$ElT,$DataT}"
if read(attributes(g)["type"]) != typestr
error("HDF5 group or file does not contain $typestr data")
end
if ElT == Nothing
return Dense{Nothing}()
end
# Attribute __complex__ is attached to the "data" dataset
# by the h5 library used by C++ version of ITensor:
if haskey(attributes(g["data"]), "__complex__")
M = read(g, "data")
nelt = size(M, 1) * size(M, 2)
data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
else
data = read(g, "data")
end
return Diag{ElT,DataT}(data)
end
20 changes: 0 additions & 20 deletions NDTensors/src/empty/empty.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,23 +94,3 @@ end
using .TypeParameterAccessors: TypeParameterAccessors
TypeParameterAccessors.parenttype(empty::Type{<:EmptyStorage}) = storagetype(empty)
zero(empty::EmptyStorage) = empty
# XXX: this seems a bit strange and fragile?
# Takes the type very literally.
function HDF5.read(
parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{StoreT}
) where {StoreT<:EmptyStorage}
g = open_group(parent, name)
typestr = string(StoreT)
if read(attributes(g)["type"]) != typestr
error("HDF5 group or file does not contain $typestr data")
end
return StoreT()
end

function HDF5.write(
parent::Union{HDF5.File,HDF5.Group}, name::String, ::StoreT
) where {StoreT<:EmptyStorage}
g = create_group(parent, name)
attributes(g)["type"] = string(StoreT)
return attributes(g)["version"] = 1
end
1 change: 0 additions & 1 deletion NDTensors/src/imports.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@ using Random
using LinearAlgebra
using StaticArrays
using Functors
using HDF5
using SimpleTraits
using SplitApplyCombine
using Strided
Expand Down

2 comments on commit fb77a8f

@mtfishman
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register subdir=NDTensors

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/105128

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a NDTensors-v0.3.1 -m "<description of version>" fb77a8f275757a90afe373c2d5bd9ebea7755863
git push origin NDTensors-v0.3.1

Please sign in to comment.