GPU functionality in NamedDimsArray

NDTensors/Project.toml

@@ -10,6 +10,7 @@ ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
 Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"

NDTensors/src/NDTensors.jl

@@ -37,6 +37,8 @@ for lib in [
   :SmallVectors,
   :SortedSets,
   :TagSets,
+  :UnspecifiedTypes,
+  :UnallocatedArrays,
 ]
   include("lib/$(lib)/src/$(lib).jl")
   @eval using .$lib: $lib

NDTensors/src/lib/NamedDimsArrays/ext/NamedDimsArraysTensorAlgebraExt/test/runtests.jl

@@ -3,17 +3,23 @@ using Test: @test, @testset, @test_broken
 using NDTensors.NamedDimsArrays: named, unname
 using NDTensors.TensorAlgebra: TensorAlgebra
 using LinearAlgebra: qr
-@testset "NamedDimsArraysTensorAlgebraExt contract (eltype=$(elt))" for elt in (
-  Float32, ComplexF32, Float64, ComplexF64
-)
+
+using NDTensors: NDTensors
+using GPUArraysCore: @allowscalar
+include(joinpath(pkgdir(NDTensors), "test", "NDTensorsTestUtils", "NDTensorsTestUtils.jl"))
+using .NDTensorsTestUtils: devices_list
+@testset "NamedDimsArraysTensorAlgebraExt contract (eltype=$(elt))" for dev in
+                                                                        devices_list(ARGS),
+  elt in (Float32, ComplexF32, Float64, ComplexF64)
+
   i = named(2, "i")
   j = named(2, "j")
   k = named(2, "k")
-  na1 = randn(elt, i, j)
-  na2 = randn(elt, j, k)
+  na1 = dev(randn(elt, i, j))
+  na2 = dev(randn(elt, j, k))
   na_dest = TensorAlgebra.contract(na1, na2)
   @test eltype(na_dest) === elt
-  @test unname(na_dest, (i, k)) ≈ unname(na1) * unname(na2)
+  @allowscalar @test unname(na_dest, (i, k)) ≈ unname(na1) * unname(na2)
 end
 @testset "NamedDimsArraysTensorAlgebraExt QR (eltype=$(elt))" for elt in (
   Float32, ComplexF32, Float64, ComplexF64

NDTensors/src/lib/NamedDimsArrays/test/test_basic.jl

@@ -10,10 +10,14 @@ using NDTensors.NamedDimsArrays:
   namedaxes,
   namedsize,
   unname
+using NDTensors: NDTensors
+using GPUArraysCore: @allowscalar
+include(joinpath(pkgdir(NDTensors), "test", "NDTensorsTestUtils", "NDTensorsTestUtils.jl"))
+using .NDTensorsTestUtils: devices_list
 
 @testset "NamedDimsArrays $(@__FILE__)" begin
-  @testset "Basic functionality" begin
-    a = randn(3, 4)
+  @testset "Basic functionality on $dev" for dev in devices_list(ARGS)
+    a = dev(randn(3, 4))
     na = named(a, ("i", "j"))
     # TODO: Just call this `size`?
     i, j = namedsize(na)
@@ -22,86 +26,96 @@ using NDTensors.NamedDimsArrays:
     @test !isnamed(a)
     @test isnamed(na)
     @test dimnames(na) == ("i", "j")
-    @test na[1, 1] == a[1, 1]
-    na[1, 1] = 11
-    @test na[1, 1] == 11
+    @allowscalar begin
+      @test na[1, 1] == a[1, 1]
+      na[1, 1] = 11
+      @test na[1, 1] == 11
+    end
     # TODO: Should `size` output `namedsize`?
     @test size(na) == (3, 4)
     @test namedsize(na) == (named(3, "i"), named(4, "j"))
     @test length(na) == 12
     # TODO: Should `axes` output `namedaxes`?
     @test axes(na) == (1:3, 1:4)
     @test namedaxes(na) == (named(1:3, "i"), named(1:4, "j"))
-    @test randn(named(3, "i"), named(4, "j")) isa NamedDimsArray
-    @test na["i" => 1, "j" => 2] == a[1, 2]
-    @test na["j" => 2, "i" => 1] == a[1, 2]
-    na["j" => 2, "i" => 1] = 12
-    @test na[1, 2] == 12
-    @test na[j => 1, i => 2] == a[2, 1]
-    @test na[aj => 1, ai => 2] == a[2, 1]
-    na[j => 1, i => 2] = 21
-    @test na[2, 1] == 21
-    na[aj => 1, ai => 2] = 2211
-    @test na[2, 1] == 2211
+    @test dev(randn(named(3, "i"), named(4, "j"))) isa NamedDimsArray
+    @allowscalar begin
+      @test na["i" => 1, "j" => 2] == a[1, 2]
+      @test na["j" => 2, "i" => 1] == a[1, 2]
+      na["j" => 2, "i" => 1] = 12
+      @test na[1, 2] == 12
+      @test na[j => 1, i => 2] == a[2, 1]
+      @test na[aj => 1, ai => 2] == a[2, 1]
+      na[j => 1, i => 2] = 21
+      @test na[2, 1] == 21
+      na[aj => 1, ai => 2] = 2211
+      @test na[2, 1] == 2211
+    end
     na′ = align(na, ("j", "i"))
     @test a == permutedims(unname(na′), (2, 1))
     na′ = align(na, (j, i))
     @test a == permutedims(unname(na′), (2, 1))
     na′ = align(na, (aj, ai))
     @test a == permutedims(unname(na′), (2, 1))
   end
-  @testset "Shorthand constructors (eltype=$elt)" for elt in (
-    Float32, ComplexF32, Float64, ComplexF64
-  )
+
+  @testset "Shorthand constructors (eltype=$elt)" for dev in devices_list(ARGS),
+    elt in (Float32, ComplexF32, Float64, ComplexF64)
+
     i, j = named.((2, 2), ("i", "j"))
     value = rand(elt)
-    for na in (zeros(elt, i, j), zeros(elt, (i, j)))
+    for na in (dev(zeros(elt, i, j)), dev(zeros(elt, (i, j))))
       @test eltype(na) === elt
       @test na isa NamedDimsArray
       @test dimnames(na) == ("i", "j")
-      @test iszero(na)
+      @allowscalar @test iszero(na)
     end
-    for na in (fill(value, i, j), fill(value, (i, j)))
+    for na in (dev(fill(value, i, j)), dev(fill(value, (i, j))))
       @test eltype(na) === elt
       @test na isa NamedDimsArray
       @test dimnames(na) == ("i", "j")
-      @test all(isequal(value), na)
+      @allowscalar @test all(isequal(value), na)
     end
-    for na in (rand(elt, i, j), rand(elt, (i, j)))
+    for na in (dev(rand(elt, i, j)), dev(rand(elt, (i, j))))
       @test eltype(na) === elt
       @test na isa NamedDimsArray
       @test dimnames(na) == ("i", "j")
-      @test !iszero(na)
-      @test all(x -> real(x) > 0, na)
+      @allowscalar begin
+        @test !iszero(na)
+        @test all(x -> real(x) > 0, na)
+      end
     end
-    for na in (randn(elt, i, j), randn(elt, (i, j)))
+    for na in (dev(randn(elt, i, j)), dev(randn(elt, (i, j))))
       @test eltype(na) === elt
       @test na isa NamedDimsArray
       @test dimnames(na) == ("i", "j")
-      @test !iszero(na)
+      @allowscalar @test !iszero(na)
     end
   end
-  @testset "Shorthand constructors (eltype=unspecified)" begin
+
+  @testset "Shorthand constructors (eltype=unspecified)" for dev in devices_list(ARGS)
     i, j = named.((2, 2), ("i", "j"))
     default_elt = Float64
-    for na in (zeros(i, j), zeros((i, j)))
+    for na in (dev(zeros(i, j)), dev(zeros((i, j))))
       @test eltype(na) === default_elt
       @test na isa NamedDimsArray
       @test dimnames(na) == ("i", "j")
-      @test iszero(na)
+      @allowscalar @test iszero(na)
     end
-    for na in (rand(i, j), rand((i, j)))
+    for na in (dev(rand(i, j)), dev(rand((i, j))))
       @test eltype(na) === default_elt
       @test na isa NamedDimsArray
       @test dimnames(na) == ("i", "j")
-      @test !iszero(na)
-      @test all(x -> real(x) > 0, na)
+      @allowscalar begin
+        @test !iszero(na)
+        @test all(x -> real(x) > 0, na)
+      end
     end
-    for na in (randn(i, j), randn((i, j)))
+    for na in (dev(randn(i, j)), dev(randn((i, j))))
       @test eltype(na) === default_elt
       @test na isa NamedDimsArray
       @test dimnames(na) == ("i", "j")
-      @test !iszero(na)
+      @allowscalar @test !iszero(na)
     end
   end
 end

NDTensors/src/lib/TensorAlgebra/src/contract/allocate_output.jl

@@ -1,3 +1,4 @@
+using NDTensors.SetParameters: Position, set_parameter
 function allocate_output(
   ::typeof(contract),
   alg::Algorithm,
@@ -13,7 +14,10 @@ function allocate_output(
   # TODO: Define `output_type(contract, alg, labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2, α, β)`.
   # TODO: Define `output_structure(contract, alg, labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2, α, β)`.
   # TODO: Define `allocate(type, structure)`.
-  return Array{promote_type(eltype(a1), eltype(a2))}(undef, length.(axes_dest))
+  elt = promote_type(eltype(a1), eltype(a2))
+  @assert typeof(a1) == typeof(a2)
+  array = set_parameter(typeof(a1), Position{1}(), elt)
+  return array(undef, length.(axes_dest))
 end
 
 # TODO: Generalize to `output_structure`.

NDTensors/src/lib/TensorAlgebra/src/contract/contract_matricize/contract.jl

@@ -1,3 +1,4 @@
+using NDTensors.Unwrap: expose
 function contract!(
   alg::Algorithm"matricize",
   a_dest::AbstractArray,
@@ -12,7 +13,7 @@ function contract!(
   a_dest_matricized = matricize(a_dest, biperm_dest)
   a1_matricized = matricize(a1, biperm1)
   a2_matricized = matricize(a2, biperm2)
-  mul!(a_dest_matricized, a1_matricized, a2_matricized, α, β)
+  mul!(expose(a_dest_matricized), expose(a1_matricized), expose(a2_matricized), α, β)
   perm_dest = flatten(biperm_dest)
   # TODO: Create a function `unmatricize` or `unfusedims`.
   # unmatricize!(a_dest, a_dest_matricized, axes(a_dest), perm_dest)

NDTensors/src/lib/UnallocatedArrays/README.md

@@ -0,0 +1,4 @@
+# UnallocatedArrays
+
+A module defining a set of unallocated immutable lazy arrays which will be used to quickly construct 
+tensors and allocating as little data as possible.

NDTensors/src/lib/UnallocatedArrays/src/UnallocatedArrays.jl

@@ -0,0 +1,11 @@
+module UnallocatedArrays
+include("abstractfill/abstractfill.jl")
+include("abstractfill/set_types.jl")
+
+include("unallocatedfill.jl")
+include("unallocatedzeros.jl")
+include("abstractunallocatedarray.jl")
+include("set_types.jl")
+
+export UnallocatedFill, UnallocatedZeros, alloctype, set_alloctype, allocate
+end

NDTensors/src/lib/UnallocatedArrays/src/abstractfill/abstractfill.jl

@@ -0,0 +1,19 @@
+using FillArrays: AbstractFill
+using NDTensors.SetParameters: Position, get_parameter, set_parameters
+## Here are functions specifically defined for UnallocatedArrays
+## not implemented by FillArrays
+## TODO this might need a more generic name maybe like compute unit
+function alloctype(A::AbstractFill)
+  return A.alloc
+end
+
+## TODO this fails if the parameter is a type
+function alloctype(Atype::Type{<:AbstractFill})
+  return get_parameter(Atype, Position{4}())
+end
+
+allocate(A::AbstractFill) = alloctype(A)(parent(A))
+
+set_eltype(T::Type{<:AbstractFill}, elt::Type) = set_parameters(T, Position{1}(), elt)
+set_ndims(T::Type{<:AbstractFill}, n) = set_parameters(T, Position{2}(), n)
+set_axestype(T::Type{<:AbstractFill}, ax::Type) = set_parameters(T, Position{3}(), ax)

NDTensors/src/lib/UnallocatedArrays/src/abstractfill/set_types.jl

@@ -0,0 +1,75 @@
+using FillArrays: AbstractFill
+using NDTensors.SetParameters: SetParameters, Position
+using NDTensors.UnspecifiedTypes: UnspecifiedZero
+## TODO make unit tests for all of these functions
+## TODO remove P4
+# `SetParameters.jl` overloads.
+SetParameters.get_parameter(::Type{<:AbstractFill{P1}}, ::Position{1}) where {P1} = P1
+SetParameters.get_parameter(::Type{<:AbstractFill{<:Any,P2}}, ::Position{2}) where {P2} = P2
+function SetParameters.get_parameter(
+  ::Type{<:AbstractFill{<:Any,<:Any,P3}}, ::Position{3}
+) where {P3}
+  return P3
+end
+
+## Setting paramaters
+# right now I am just defining the necessary ones for my implementation still working on full implementation
+# Set parameter 1
+SetParameters.set_parameter(T::Type{<:AbstractFill}, ::Position{1}, P1) = T{P1}
+
+# Set parameter 2
+function SetParameters.set_parameter(
+  T::Type{<:AbstractFill{P1}}, ::Position{2}, P2
+) where {P1}
+  return T{P2}
+end
+
+# Set parameter 3
+function SetParameters.set_parameter(
+  T::Type{<:AbstractFill{P1,P2}}, ::Position{3}, P3
+) where {P1,P2}
+  return T{P3}
+end
+
+## TODO define a specify_parameters function
+## To quickly specify P1, P2, and P3 
+## default parameters
+function SetParameters.default_parameter(::Type{<:AbstractFill}, ::Position{1})
+  return UnspecifiedTypes.UnspecifiedZero
+end
+SetParameters.default_parameter(::Type{<:AbstractFill}, ::Position{2}) = 0
+SetParameters.default_parameter(::Type{<:AbstractFill}, ::Position{3}) = Tuple{}
+
+SetParameters.nparameters(::Type{<:AbstractFill}) = Val(3)
+
+# Set parameter 1
+## Right now using AbstractArray
+## TODO These are more difficult because T is technically defined so need some way to strip T of it {} types
+# function set_parameter(
+#   T::Type{<:AbstractFill{<:Any,<:Any,P3}}, ::Position{1}, P1
+# ) where {P3}
+#   return T{P1,<:Any,P3}
+# end
+# function set_parameter(
+#   T::Type{<:AbstractFill{<:Any,P2,P3}}, ::Position{1}, P1
+# ) where {P2,P3}
+#   return T{P1,P2,P3}
+# end
+# function set_parameter(
+#   T::Type{<:AbstractFill{<:Any,<:Any,P3}}, ::Position{2}, P2
+# ) where {P3}
+#   return T{<:Any,P2,P3}
+# end
+# function set_parameter(
+#   T::Type{<:AbstractFill{P1,<:Any,P3}}, ::Position{2}, P2
+# ) where {P1,P3}
+#   return T
+# end
+# set_parameter(T::Type{<:AbstractFill}, ::Position{3}, P3) = T{<:Any,<:Any,P3}
+# set_parameter(T::Type{<:AbstractFill{P1}}, ::Position{3}, P3) where {P1} = T{<:Any,P3}
+# function set_parameter(T::Type{<:AbstractFill{<:Any,P2}}, ::Position{3}, P3) where {P2}
+#   return T{<:Any,P2,P3}
+# end
+
+# Set parameter 2
+## using AbstractArray

NDTensors/src/lib/UnallocatedArrays/src/abstractunallocatedarray.jl

@@ -0,0 +1,52 @@
+using FillArrays: FillArrays, getindex_value
+using NDTensors.SetParameters: set_parameters
+using Adapt: adapt
+
+const UnallocatedArray{ElT,N,AxesT,AllocT} = Union{
+  <:UnallocatedFill{ElT,N,AxesT,AllocT},<:UnallocatedZeros{ElT,N,AxesT,AllocT}
+}
+
+@inline Base.axes(A::UnallocatedArray) = axes(parent(A))
+Base.size(A::UnallocatedArray) = size(parent(A))
+function FillArrays.getindex_value(A::UnallocatedArray)
+  return getindex_value(parent(A))
+end
+
+function Base.complex(A::UnallocatedArray)
+  return set_alloctype(
+    complex(parent(A)), set_parameters(alloctype(A), Position{1}(), complex(eltype(A)))
+  )
+end
+
+function Base.transpose(a::UnallocatedArray)
+  return set_alloctype(transpose(parent(a)), alloctype(a))
+end
+
+function Base.adjoint(a::UnallocatedArray)
+  return set_alloctype(adjoint(parent(a)), alloctype(a))
+end
+
+function Base.similar(a::UnallocatedArray)
+  return alloctype(a)(undef, size(a))
+end
+
+function set_alloctype(T::Type{<:UnallocatedArray}, alloc::Type{<:AbstractArray})
+  return set_parameters(T, Position{4}(), alloc)
+end
+
+## This overloads the definition defined in `FillArrays.jl`
+for STYPE in (:AbstractArray, :AbstractFill)
+  @eval begin
+    @inline $STYPE{T}(F::UnallocatedArray{T}) where {T} = F
+    @inline $STYPE{T,N}(F::UnallocatedArray{T,N}) where {T,N} = F
+  end
+end
+
+## TODO fix this because reshape loses alloctype
+#FillArrays.reshape(a::Union{<:UnallocatedFill, <:UnallocatedZeros}, dims) = set_alloctype(reshape(parent(a), dims), allocate(a))
+
+# function Adapt.adapt_storage(to::Type{<:AbstractArray}, x::Union{<:UnallocatedFill, <:UnallocatedZeros})
+#   return set_alloctype(parent(x), to)
+# end
+
+# function Adapt.adapt_storage(to::Type{<:Number}, x::)