diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 59174aa6..a31b7e61 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -20,6 +20,7 @@ jobs: uses: "QuantumKitHub/QuantumKitHubActions/.github/workflows/TestGroups.yml@main" with: fast: ${{ github.event.pull_request.draft == true }} + exclude: '["testsuite"]' nthreads: 4 timeout-minutes: 90 # Enzyme plays a large role here, so also test macOS and Windows on PRs, once ubuntu passes diff --git a/Project.toml b/Project.toml index 31c63a38..d242ad43 100644 --- a/Project.toml +++ b/Project.toml @@ -42,56 +42,26 @@ TensorOperationsJLArraysExt = "JLArrays" [compat] AMDGPU = "2" -Aqua = "0.6, 0.7, 0.8" -Adapt = "4" Bumper = "0.6, 0.7" CUDACore = "6" ChainRulesCore = "1" -ChainRulesTestUtils = "1" -DynamicPolynomials = "0.5, 0.6" Enzyme = "0.13.183" -EnzymeTestUtils = "0.2" GPUArrays = "11" JLArrays = "0.3" LRUCache = "1" LinearAlgebra = "1.6" -Logging = "1.6" Mooncake = "0.5" PackageExtensionCompat = "1" -ParallelTestRunner = "2" PrecompileTools = "1.1" Preferences = "1.4" PtrArrays = "1.2" -Random = "1" -Strided = "2.6" +Strided = "2.6.7" StridedViews = "0.5" TBLIS = "0.3" -Test = "1" TupleTools = "1.6" VectorInterface = "0.4.1, 0.5, 0.6, 0.7" -cuRAND = "6" cuTENSOR = "6" julia = "1.10" -[extras] -Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" -Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e" -CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" -ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" -ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" -cuRAND = "20fd9a0b-12d5-4c2f-a8af-7c34e9e60431" -DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" -EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a" -JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" -Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" -Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" -ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" -Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9" -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1" - -[targets] -test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "ChainRulesCore", "cuRAND", "CUDACore", "cuTENSOR", "Aqua", "Logging", "Bumper", "Mooncake", "Enzyme", "EnzymeTestUtils", "Adapt", "JLArrays", "AMDGPU", "TBLIS", "ParallelTestRunner"] +[workspace] +projects = ["test"] diff --git a/ext/TensorOperationsTBLISExt.jl b/ext/TensorOperationsTBLISExt.jl index 8e44c45c..2e4fe6c9 100644 --- a/ext/TensorOperationsTBLISExt.jl +++ b/ext/TensorOperationsTBLISExt.jl @@ -157,6 +157,8 @@ function TO.tensortrace!( throw(ArgumentError("output tensor must not be aliased with input tensor")) T = eltype(C) + # TBLIS does not handle traces over dimensions of size zero, which only scale `C` + any(i -> iszero(size(A, i)), q[1]) && return TO._scale!(C, β) einA, einC = trace_labels(p, q) Av, Cv = SV(A), SV(C) lenA, strideA = tblis_dims(Av) diff --git a/ext/TensorOperationscuTENSORExt.jl b/ext/TensorOperationscuTENSORExt.jl index ad52efd5..eb6c271f 100644 --- a/ext/TensorOperationscuTENSORExt.jl +++ b/ext/TensorOperationscuTENSORExt.jl @@ -59,16 +59,17 @@ end #------------------------------------------------------------------------------------------- # A Base wrapper over `CuArray` will first pass via the `select_backend` methods for # `AbstractArray` and be converted into a `StridedView` if it satisfies `isstrided`. Hence, -# we only need to capture `CuStridedView` here. -function TO.select_backend(::typeof(TO.tensoradd!), ::CuStridedView, ::CuStridedView) +# we only need to capture `CuStridedView` here. cuTENSOR requires all arguments to have the +# same element type; other combinations fall back to the `StridedView` methods. +function TO.select_backend(::typeof(TO.tensoradd!), ::CuStridedView{T}, ::CuStridedView{T}) where {T} return cuTENSORBackend() end -function TO.select_backend(::typeof(TO.tensortrace!), ::CuStridedView, ::CuStridedView) +function TO.select_backend(::typeof(TO.tensortrace!), ::CuStridedView{T}, ::CuStridedView{T}) where {T} return cuTENSORBackend() end function TO.select_backend( - ::typeof(TO.tensorcontract!), ::CuStridedView, ::CuStridedView, ::CuStridedView - ) + ::typeof(TO.tensorcontract!), ::CuStridedView{T}, ::CuStridedView{T}, ::CuStridedView{T} + ) where {T} return cuTENSORBackend() end @@ -83,6 +84,9 @@ function TO.tensoradd!( α::Number, β::Number, backend::cuTENSORBackend, allocator = CUDAAllocator() ) + TO.argcheck_tensoradd(C, A, pA) + TO.dimcheck_tensoradd(C, A, pA) + isempty(C) && return C C_cuda, isview = _custrided(C, allocator) A_cuda, = _custrided(A, allocator) tensoradd!(C_cuda, A_cuda, pA, conjA, α, β, backend, allocator) @@ -97,6 +101,11 @@ function TO.tensorcontract!( α::Number, β::Number, backend::cuTENSORBackend, allocator = CUDAAllocator() ) + TO.argcheck_tensorcontract(C, A, pA, B, pB, pAB) + TO.dimcheck_tensorcontract(C, A, pA, B, pB, pAB) + isempty(C) && return C + # cuTENSOR does not accept dimensions of size zero; contracting over them only scales `C` + any(i -> iszero(size(A, i)), pA[2]) && return TO._scale!(C, β) C_cuda, isview = _custrided(C, allocator) A_cuda, = _custrided(A, allocator) B_cuda, = _custrided(B, allocator) @@ -112,6 +121,13 @@ function TO.tensortrace!( α::Number, β::Number, backend::cuTENSORBackend, allocator = CUDAAllocator() ) + TO.argcheck_tensortrace(C, A, p, q) + TO.dimcheck_tensortrace(C, A, p, q) + isempty(C) && return C + # cuTENSOR does not accept dimensions of size zero; tracing over them only scales `C` + any(i -> iszero(size(A, i)), q[1]) && return TO._scale!(C, β) + # a trace without traced indices is a permutation, which cuTENSOR cannot plan as a reduction + isempty(q[1]) && return tensoradd!(C, A, p, conjA, α, β, backend, allocator) C_cuda, isview = _custrided(C, allocator) A_cuda, = _custrided(A, allocator) tensortrace!(C_cuda, A_cuda, p, q, conjA, α, β, backend, allocator) diff --git a/src/implementation/base.jl b/src/implementation/base.jl index 2410436a..9642eb57 100644 --- a/src/implementation/base.jl +++ b/src/implementation/base.jl @@ -104,6 +104,8 @@ function tensortrace!( szA = size(A) so = TupleTools.getindices(szA, linearize(p)) st = prod(TupleTools.getindices(szA, q[1])) + # traced dimensions of size zero: there is nothing to add to `β * C` + iszero(st) && return _scale!(C, β) Ã = reshape( PermutedDimsArray(A, (linearize(p)..., linearize(q)...)), (prod(so), st, st) ) @@ -141,6 +143,8 @@ function tensortrace!( szA = size(A) so = TupleTools.getindices(szA, linearize(p)) st = prod(TupleTools.getindices(szA, q[1])) + # traced dimensions of size zero: there is nothing to add to `β * C` + iszero(st) && return _scale!(C, β) perm = (linearize(p)..., linearize(q)...) Atemp = tensoralloc_add(eltype(A), A, (perm, ()), conjA, Val(true), allocator) Ã = reshape(permutedims!(Atemp, A, perm), (prod(so), st, st)) diff --git a/src/implementation/strided.jl b/src/implementation/strided.jl index 3dd4bf11..29a75e4e 100644 --- a/src/implementation/strided.jl +++ b/src/implementation/strided.jl @@ -101,6 +101,9 @@ function tensorcontract!( ) @nospecialize backend allocator + (Base.mightalias(C, A) || Base.mightalias(C, B)) && + throw(ArgumentError("output tensor must not be aliased with input tensor")) + # standardize input types for compilation time α′ = standardize_scalartype(C, α) β′ = standardize_scalartype(C, β) diff --git a/src/utils.jl b/src/utils.jl index 368c765e..8602e2c6 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -9,3 +9,7 @@ function _kron(Es::NTuple{N, Any}, ba) where {N} ) return tensorproduct(E1, ((1, 2), ()), false, E2, p2, false, p, One(), ba...) end + +# `C = β * C`, for operations without anything to add to `C`, such as traces or contractions +# over dimensions of size zero; `β = 0` overwrites `C`, which may contain `NaN`s +_scale!(C::AbstractArray, β::Number) = iszero(β) ? fill!(C, zero(eltype(C))) : (C .= β .* C) diff --git a/test/Project.toml b/test/Project.toml new file mode 100644 index 00000000..3dc75059 --- /dev/null +++ b/test/Project.toml @@ -0,0 +1,57 @@ +name = "TensorOperationsTests" + +[deps] +AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e" +CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" +DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a" +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67" +TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9" +TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6" +VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8" +cuBLAS = "182d3088-87b7-4494-8cad-fc6afaa545bc" +cuRAND = "20fd9a0b-12d5-4c2f-a8af-7c34e9e60431" +cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1" + +[sources] +TensorOperations = {path = ".."} + +[compat] +AMDGPU = "2" +Adapt = "4" +Aqua = "0.6, 0.7, 0.8" +Bumper = "0.6, 0.7" +CUDACore = "6" +ChainRulesCore = "1" +ChainRulesTestUtils = "1" +DynamicPolynomials = "0.5, 0.6" +Enzyme = "0.13.183" +EnzymeTestUtils = "0.2" +JLArrays = "0.3" +LinearAlgebra = "1.6" +Logging = "1.6" +Mooncake = "0.5" +ParallelTestRunner = "2" +Random = "1" +Strided = "2.6.7" +TBLIS = "0.3" +Test = "1" +TupleTools = "1.6" +VectorInterface = "0.4.1, 0.5, 0.6, 0.7" +cuBLAS = "6" +cuRAND = "6" +cuTENSOR = "6" diff --git a/test/core/backends/basecopy.jl b/test/core/backends/basecopy.jl new file mode 100644 index 00000000..a2dcf861 --- /dev/null +++ b/test/core/backends/basecopy.jl @@ -0,0 +1,15 @@ +using TensorOperations +using TensorOperations: BaseCopy +using Test + +@isdefined(TestSuite) || include("../../testsuite/TestSuite.jl") + +# `BaseCopy` is the reference for all other backends, and is itself checked against loops +@testset "BaseCopy" begin + TestSuite.test_backend(BaseCopy(); reference = TestSuite.LoopReference(), diagonal = true) +end +@testset "BaseCopy with generic element types" begin + TestSuite.test_backend( + BaseCopy(); reference = TestSuite.LoopReference(), eltypes = (BigFloat, Complex{BigFloat}) + ) +end diff --git a/test/core/backends/baseview.jl b/test/core/backends/baseview.jl new file mode 100644 index 00000000..ad821bc9 --- /dev/null +++ b/test/core/backends/baseview.jl @@ -0,0 +1,9 @@ +using TensorOperations +using TensorOperations: BaseView +using Test + +@isdefined(TestSuite) || include("../../testsuite/TestSuite.jl") + +@testset "BaseView" begin + TestSuite.test_backend(BaseView(); diagonal = true) +end diff --git a/test/core/backends/default.jl b/test/core/backends/default.jl new file mode 100644 index 00000000..9e42910d --- /dev/null +++ b/test/core/backends/default.jl @@ -0,0 +1,9 @@ +using TensorOperations +using TensorOperations: DefaultBackend +using Test + +@isdefined(TestSuite) || include("../../testsuite/TestSuite.jl") + +@testset "DefaultBackend" begin + TestSuite.test_backend(DefaultBackend(); diagonal = true) +end diff --git a/test/core/backends/strided.jl b/test/core/backends/strided.jl new file mode 100644 index 00000000..c22b70bc --- /dev/null +++ b/test/core/backends/strided.jl @@ -0,0 +1,20 @@ +using TensorOperations +using TensorOperations: StridedBLAS, StridedNative, ManualAllocator, BufferAllocator +using Test + +@isdefined(TestSuite) || include("../../testsuite/TestSuite.jl") + +@testset "StridedNative" begin + TestSuite.test_backend(StridedNative(); diagonal = true, aliasing = true) +end +@testset "StridedBLAS" begin + TestSuite.test_backend(StridedBLAS(); aliasing = true) +end +@testset "$backend with generic element types" for backend in (StridedNative(), StridedBLAS()) + TestSuite.test_backend(backend; eltypes = (BigFloat, Complex{BigFloat})) +end +@testset "$backend with allocator $allocator" for backend in (StridedNative(), StridedBLAS()), + allocator in (ManualAllocator(), BufferAllocator()) + + TestSuite.test_tensorcontract(backend; allocator, eltypes = (Float64, ComplexF64)) +end diff --git a/test/core/methods.jl b/test/core/methods.jl index 501e9ace..2497a509 100644 --- a/test/core/methods.jl +++ b/test/core/methods.jl @@ -131,175 +131,3 @@ backendlist = (BaseCopy(), BaseView(), StridedNative(), StridedBLAS()) @test C2 == C3 == C4 end end - -# test in-place methods -#----------------------- -# test different versions of in-place methods, -# with changing element type and with nontrivial strides -@testset "in-place methods with backend $b" for b in backendlist - @testset "tensorcopy!" begin - Abig = randn(Float64, (30, 30, 30, 30)) - A = view(Abig, 1 .+ 3 * (0:9), 2 .+ 2 * (0:6), 5 .+ 4 * (0:6), 4 .+ 3 * (0:8)) - p = (3, 1, 4, 2) - Cbig = zeros(ComplexF64, (50, 50, 50, 50)) - C = view(Cbig, 13 .+ (0:6), 11 .+ 4 * (0:9), 15 .+ 4 * (0:8), 4 .+ 3 * (0:6)) - Acopy = tensorcopy(A, 1:4) - Ccopy = tensorcopy(C, 1:4) - pA = (p, ()) - α = randn(Float64) - tensorcopy!(C, A, pA, false, α, b) - tensorcopy!(Ccopy, Acopy, pA, false, 1.0, b) - @test C ≈ α * Ccopy - @test_throws IndexError tensorcopy!(C, A, ((1, 2, 3), ()), false, 1.0, b) - @test_throws DimensionMismatch tensorcopy!(C, A, ((1, 2, 3, 4), ()), false, 1.0, b) - @test_throws IndexError tensorcopy!(C, A, ((1, 2, 2, 3), ()), false, 1.0, b) - end - - @testset "tensoradd!" begin - Abig = randn(Float64, (30, 30, 30, 30)) - A = view(Abig, 1 .+ 3 * (0:9), 2 .+ 2 * (0:6), 5 .+ 4 * (0:6), 4 .+ 3 * (0:8)) - p = (3, 1, 4, 2) - Cbig = zeros(ComplexF64, (50, 50, 50, 50)) - C = view(Cbig, 13 .+ (0:6), 11 .+ 4 * (0:9), 15 .+ 4 * (0:8), 4 .+ 3 * (0:6)) - Acopy = tensorcopy(p, A, 1:4) - Ccopy = tensorcopy(1:4, C, 1:4) - α = randn(Float64) - β = randn(Float64) - tensoradd!(C, A, (p, ()), false, α, β, b) - Ccopy = β * Ccopy + α * Acopy - @test C ≈ Ccopy - @test_throws IndexError tensoradd!(C, A, ((1, 2, 3), ()), false, 1.2, 0.5, b) - @test_throws DimensionMismatch tensoradd!( - C, A, ((1, 2, 3, 4), ()), false, 1.2, 0.5, b - ) - @test_throws IndexError tensoradd!(C, A, ((1, 1, 2, 3), ()), false, 1.2, 0.5, b) - end - - @testset "tensortrace!" begin - Abig = rand(Float64, (30, 30, 30, 30)) - A = view(Abig, 1 .+ 3 * (0:8), 2 .+ 2 * (0:14), 5 .+ 4 * (0:6), 7 .+ 2 * (0:8)) - Bbig = rand(ComplexF64, (50, 50)) - B = view(Bbig, 13 .+ (0:14), 3 .+ 5 * (0:6)) - Acopy = tensorcopy(A, 1:4) - Bcopy = tensorcopy(B, 1:2) - α = randn(Float64) - β = randn(Float64) - tensortrace!(B, A, ((2, 3), ()), ((1,), (4,)), false, α, β, b) - Bcopy = β * Bcopy - for i in 1 .+ (0:8) - Bcopy += α * view(A, i, :, :, i) - end - @test B ≈ Bcopy - @test_throws IndexError tensortrace!(B, A, ((1,), ()), ((2,), (3,)), false, α, β, b) - @test_throws DimensionMismatch tensortrace!( - B, A, ((1, 4), ()), ((2,), (3,)), false, α, β, b - ) - @test_throws IndexError tensortrace!( - B, A, ((1, 4), ()), ((1, 1), (4,)), false, α, β, b - ) - @test_throws IndexError tensortrace!( - B, A, ((1, 4), ()), ((1,), (3,)), false, α, β, b - ) - end - - @testset "tensorcontract! with allocator = $allocator" for allocator in - (DefaultAllocator(), ManualAllocator(), BufferAllocator()) - Abig = rand(Float64, (30, 30, 30, 30)) - A = view(Abig, 1 .+ 3 * (0:8), 2 .+ 2 * (0:14), 5 .+ 4 * (0:6), 7 .+ 2 * (0:8)) - Bbig = rand(ComplexF64, (50, 50, 50)) - B = view(Bbig, 3 .+ 5 * (0:6), 7 .+ 2 * (0:7), 13 .+ (0:14)) - Cbig = rand(ComplexF32, (40, 40, 40)) - C = view(Cbig, 3 .+ 2 * (0:8), 13 .+ (0:8), 7 .+ 3 * (0:7)) - Acopy = tensorcopy(A, 1:4) - Bcopy = tensorcopy(B, 1:3) - Ccopy = tensorcopy(C, 1:3) - α = randn(Float64) - β = randn(Float64) - Ccopy = β * Ccopy - for d in 1 .+ (0:8), a in 1 .+ (0:8), e in 1 .+ (0:7) - for b in 1 .+ (0:14), c in 1 .+ (0:6) - Ccopy[d, a, e] += α * A[a, b, c, d] * conj(B[c, e, b]) - end - end - tensorcontract!( - C, A, ((4, 1), (2, 3)), false, B, ((3, 1), (2,)), true, - ((1, 2, 3), ()), α, β, b, allocator - ) - @test C ≈ Ccopy - @test_throws IndexError tensorcontract!( - C, - A, ((4, 1), (2, 4)), false, - B, ((1, 3), (2,)), false, - ((1, 2, 3), ()), α, β, b - ) - @test_throws IndexError tensorcontract!( - C, - A, ((4, 1), (2, 3)), false, - B, ((1, 3), ()), false, - ((1, 2, 3), ()), α, β, b - ) - @test_throws IndexError tensorcontract!( - C, - A, ((4, 1), (2, 3)), false, - B, ((1, 3), (2,)), false, - ((1, 2), ()), α, β, b - ) - @test_throws DimensionMismatch tensorcontract!( - C, - A, ((4, 1), (2, 3)), false, - B, ((1, 3), (2,)), false, - ((1, 3, 2), ()), α, β, b - ) - end -end - -# β = 0 must act as a strong zero: uninitialized output must not leak NaNs -#------------------------------------------------------------------------- -@testset "strong zero β with Diagonal ($T)" for T in (Float64, ComplexF64) - Zero = TensorOperations.VectorInterface.Zero - n = 4 - d = randn(T, n) - S = Diagonal(d) - A = randn(T, n, n, n, n) - F = randn(T, n, n) - poison(dims...) = fill(T(NaN), dims...) - - # C[i,j,k,l] = A[i,j,k,a] * S[a,l] - C1 = poison(n, n, n, n) - tensorcontract!( - C1, A, ((1, 2, 3), (4,)), false, S, ((1,), (2,)), false, - ((1, 2, 3, 4), ()), 1, Zero() - ) - @test C1 ≈ [A[i, j, k, l] * d[l] for i in 1:n, j in 1:n, k in 1:n, l in 1:n] - - # C[i,j] = A[i,j,a,b] * S[a,b] - C2 = poison(n, n) - tensorcontract!( - C2, A, ((1, 2), (3, 4)), false, S, ((1, 2), ()), false, - ((1, 2), ()), 1, Zero() - ) - @test C2 ≈ [sum(A[i, j, a, a] * d[a] for a in 1:n) for i in 1:n, j in 1:n] - - # C[a,c,b,d] = F[a,b] * S[c,d] - C3 = poison(n, n, n, n) - tensorcontract!( - C3, F, ((1, 2), ()), false, S, ((), (1, 2)), false, - ((1, 3, 2, 4), ()), 1, Zero() - ) - @test C3 ≈ [F[a, b] * (c == e ? d[c] : zero(T)) for a in 1:n, c in 1:n, b in 1:n, e in 1:n] - - # C[i,k] = S[i,j] * S[j,k], into a dense matrix and into a Diagonal - C4 = poison(n, n) - tensorcontract!( - C4, S, ((1,), (2,)), false, S, ((1,), (2,)), false, - ((1, 2), ()), 1, Zero() - ) - @test C4 ≈ Array(S * S) - - C5 = Diagonal(poison(n)) - tensorcontract!( - C5, S, ((1,), (2,)), false, S, ((1,), (2,)), false, - ((1, 2), ()), 1, Zero() - ) - @test C5 ≈ S * S -end diff --git a/test/extensions/gpuarrays.jl b/test/extensions/gpuarrays.jl index 4a95e889..c511027a 100644 --- a/test/extensions/gpuarrays.jl +++ b/test/extensions/gpuarrays.jl @@ -1,5 +1,5 @@ using TensorOperations -using TensorOperations: StridedBLAS, StridedNative, linearize, numout +using TensorOperations: StridedBLAS, StridedNative, DefaultBackend, cuTENSORBackend, linearize, numout using Test using Adapt using TupleTools @@ -7,6 +7,7 @@ using JLArrays using VectorInterface using CUDACore, AMDGPU using cuTENSOR # contractions of `CuArray`s require the cuTENSOR extension +using cuBLAS # matrix multiplications of `CuArray`s with `StridedBLAS` require the cuBLAS extension of Strided test_result(a::AbstractArray, b::AbstractArray; kwargs...) = isapprox(collect(a), collect(b); kwargs...) @@ -34,32 +35,30 @@ bufferallocator(::Type{JLArray}; kwargs...) = TensorOperations.JLBufferAllocator bufferallocator(::Type{CuArray}; kwargs...) = TensorOperations.CUDABufferAllocator(; kwargs...) bufferallocator(::Type{ROCArray}; kwargs...) = TensorOperations.AMDBufferAllocator(; kwargs...) -@testset "tensoradd! ($AT)" for AT in ATs - sz = (3, 5, 4, 6) - p = (3, 1, 4, 2) - for backend in backends, T in (Float32, ComplexF32) - A = randn(T, sz) - C = randn(T, TupleTools.getindices(sz, p)) - - @test compare(AT, C, A) do c, a - tensoradd!(c, a, (p, ()), false, One(), Zero(), backend) - end +@isdefined(TestSuite) || include("../testsuite/TestSuite.jl") - α = rand(T) - @test compare(AT, C, A) do c, a - tensoradd!(c, a, (p, ()), false, α, Zero(), backend) - end - - β = rand(T) - @test compare(AT, C, A) do c, a - tensoradd!(c, a, (p, ()), false, α, β, backend) - end +@testset "test suite: $backend ($AT)" for AT in ATs, backend in backends + TestSuite.test_backend(backend; arraytype = AT, eltypes = (Float32, ComplexF32)) +end - T <: Real || @test compare(AT, C, A) do c, a - tensoradd!(c, a, (p, ()), true, α, β, backend) - end +# cuTENSOR requires all arguments to have the same element type, and positive strides; the +# default backend selects cuTENSOR only for equal element types, and Strided otherwise +if CUDACore.functional() + @testset "test suite: cuTENSORBackend() (CuArray)" begin + TestSuite.test_backend( + cuTENSORBackend(); arraytype = CuArray, eltypes = (Float32, ComplexF32), + mixed_eltypes = (), negative_strides = false + ) + end + @testset "test suite: DefaultBackend() (CuArray)" begin + TestSuite.test_backend( + DefaultBackend(); arraytype = CuArray, eltypes = (Float32, ComplexF32), + negative_strides = false + ) end - # test Diagonal special case +end + +@testset "tensoradd! with Diagonal ($AT)" for AT in ATs sz = (8, 8) p = (2, 1) diag_backends = [BaseCopy(), BaseView()] @@ -87,75 +86,6 @@ bufferallocator(::Type{ROCArray}; kwargs...) = TensorOperations.AMDBufferAllocat end end -@testset "tensortrace! ($AT)" for AT in ATs - sz = (2, 4, 3, 2) - p = (2, 3) - q = ((1,), (4,)) - - for backend in backends, T in (Float32, ComplexF32) - A = randn(T, sz) - C = randn(T, TupleTools.getindices(sz, p)) - - @test compare(AT, C, A) do c, a - tensortrace!(c, a, (p, ()), q, false, One(), Zero(), backend) - end - - α = rand(T) - @test compare(AT, C, A) do c, a - tensortrace!(c, a, (p, ()), q, false, α, Zero(), backend) - end - - β = rand(T) - @test compare(AT, C, A) do c, a - tensortrace!(c, a, (p, ()), q, false, α, β, backend) - end - - T <: Real || @test compare(AT, C, A) do c, a - tensortrace!(c, a, (p, ()), q, true, α, β, backend) - end - end -end - -@testset "tensorcontract! ($AT)" for AT in ATs - sz = (2, 4, 3, 4, 2, 5) - pA = ((4, 1), (2, 3)) - pB = ((3, 1), (2,)) - pAB = ((1, 2, 3), ()) - - for backend in backends, T in (Float32, ComplexF32) - A = randn(T, (2, 4, 3, 2)) - B = randn(T, (3, 3, 4)) - C = randn(T, (2, 2, 3)) - - @test compare(AT, C, A, B) do c, a, b - tensorcontract!(c, a, pA, false, b, pB, false, pAB, One(), Zero(), backend) - end - - α = rand(T) - @test compare(AT, C, A, B) do c, a, b - tensorcontract!(c, a, pA, false, b, pB, false, pAB, α, Zero(), backend) - end - - β = rand(T) - @test compare(AT, C, A, B) do c, a, b - tensorcontract!(c, a, pA, false, b, pB, false, pAB, α, β, backend) - end - - if !(T <: Real) - @test compare(AT, C, A, B) do c, a, b - tensorcontract!(c, a, pA, true, b, pB, false, pAB, α, β, backend) - end - @test compare(AT, C, A, B) do c, a, b - tensorcontract!(c, a, pA, false, b, pB, true, pAB, α, β, backend) - end - @test compare(AT, C, A, B) do c, a, b - tensorcontract!(c, a, pA, true, b, pB, true, pAB, α, β, backend) - end - end - end - -end - @testset "BufferAllocator ($AT)" for AT in ATs @testset "tensor network ($T)" for T in (Float32, ComplexF32) D1, D2, D3 = 30, 40, 20 diff --git a/test/extensions/tblis.jl b/test/extensions/tblis.jl index 7c8610f1..721476e3 100644 --- a/test/extensions/tblis.jl +++ b/test/extensions/tblis.jl @@ -11,124 +11,11 @@ const eltypes = (Float32, Float64, ComplexF32, ComplexF64) const tblis = TBLISBackend() const reference = StridedNative() -# `Zero()`/`One()` and plain numbers should all work; poison the output with `NaN` whenever -# `β == 0` so that a kernel that computes `0 * C` instead of ignoring `C` is caught. -poison!(C) = fill!(C, convert(eltype(C), NaN)) +@isdefined(TestSuite) || include("../testsuite/TestSuite.jl") -@testset "tensoradd! (eltype = $T)" for T in eltypes - A = randn(T, (3, 5, 4, 6)) - p = ((3, 1), (4, 2)) - for conjA in (false, true), (α, β) in ((1, 0), (randn(T), 0), (randn(T), randn(T))) - C = randn(T, (4, 3, 6, 5)) - Cref = copy(C) - iszero(β) && (poison!(C); poison!(Cref)) - @test tensoradd!(C, A, p, conjA, α, β, tblis) ≈ - tensoradd!(Cref, A, p, conjA, α, β, reference) - end - - # non-contiguous input and output through views - Aview = view(randn(T, (6, 10, 8, 12)), 1:2:6, 1:2:10, 1:2:8, 1:2:12) - C = randn(T, (4, 3, 6, 5)) - Cref = copy(C) - @test tensoradd!(C, Aview, p, true, 2, 1, tblis) ≈ - tensoradd!(Cref, Aview, p, true, 2, 1, reference) - - # an `Adjoint` input already carries a conjugation, which has to combine with `conjA` - Aadj = adjoint(randn(T, (5, 3))) - for conjA in (false, true) - C = randn(T, (5, 3)) - Cref = copy(C) - @test tensoradd!(C, Aadj, ((2, 1), ()), conjA, 2, 1, tblis) ≈ - tensoradd!(Cref, Aadj, ((2, 1), ()), conjA, 2, 1, reference) - end -end - -@testset "tensortrace! (eltype = $T)" for T in eltypes - A = randn(T, (5, 3, 4, 5, 3, 2)) - p = ((6, 3), ()) - q = ((1, 2), (4, 5)) - for conjA in (false, true), (α, β) in ((1, 0), (randn(T), randn(T))) - C = randn(T, (2, 4)) - Cref = copy(C) - iszero(β) && (poison!(C); poison!(Cref)) - @test tensortrace!(C, A, p, q, conjA, α, β, tblis) ≈ - tensortrace!(Cref, A, p, q, conjA, α, β, reference) - end - - # trace all the way down to a scalar - B = randn(T, (4, 4)) - C = fill(convert(T, NaN)) - Cref = fill(convert(T, NaN)) - @test tensortrace!(C, B, ((), ()), ((1,), (2,)), false, 1, 0, tblis)[] ≈ tr(B) -end - -@testset "tensorcontract! (eltype = $T)" for T in eltypes - # open indices of A are its dimensions 3, 1, 4 (sizes 5, 3, 3), contracted are 2 and 5 - # (sizes 20, 4); open indices of B are its dimensions 4 and 2 (sizes 3, 6) - A = randn(T, (3, 20, 5, 3, 4)) - B = randn(T, (4, 6, 20, 3)) - pA = ((3, 1, 4), (2, 5)) - pB = ((3, 1), (4, 2)) - pAB = ((3, 1, 4), (5, 2)) - - for conjA in (false, true), conjB in (false, true), - (α, β) in ((1, 0), (randn(T), 0), (randn(T), randn(T))) - - C = randn(T, (3, 5, 3, 6, 3)) - Cref = copy(C) - iszero(β) && (poison!(C); poison!(Cref)) - @test tensorcontract!(C, A, pA, conjA, B, pB, conjB, pAB, α, β, tblis) ≈ - tensorcontract!(Cref, A, pA, conjA, B, pB, conjB, pAB, α, β, reference) - end - - # outer product: no contracted indices at all - @testset "outer product" begin - A2 = randn(T, (3, 4)) - B2 = randn(T, (5,)) - pA2, pB2, pAB2 = ((1, 2), ()), ((), (1,)), ((3, 1), (2,)) - C = fill(convert(T, NaN), (5, 3, 4)) - Cref = copy(C) - @test tensorcontract!(C, A2, pA2, true, B2, pB2, false, pAB2, 1, 0, tblis) ≈ - tensorcontract!(Cref, A2, pA2, true, B2, pB2, false, pAB2, 1, 0, reference) - end - - # full contraction: zero-dimensional output - @testset "full contraction" begin - A2 = randn(T, (3, 4)) - B2 = randn(T, (4, 3)) - pA2, pB2, pAB2 = ((), (1, 2)), ((2, 1), ()), ((), ()) - C = fill(convert(T, NaN)) - Cref = fill(convert(T, NaN)) - @test tensorcontract!(C, A2, pA2, true, B2, pB2, true, pAB2, 1, 0, tblis)[] ≈ - tensorcontract!(Cref, A2, pA2, true, B2, pB2, true, pAB2, 1, 0, reference)[] - end - - # non-contiguous factors - @testset "strided views" begin - Av = view(randn(T, (6, 8)), 1:2:6, 1:2:8) - Bv = view(randn(T, (8, 10)), 1:2:8, 1:2:10) - C = fill(convert(T, NaN), (3, 5)) - Cref = copy(C) - pA2, pB2, pAB2 = ((1,), (2,)), ((1,), (2,)), ((1, 2), ()) - @test tensorcontract!(C, Av, pA2, false, Bv, pB2, true, pAB2, 1, 0, tblis) ≈ - tensorcontract!(Cref, Av, pA2, false, Bv, pB2, true, pAB2, 1, 0, reference) - end -end - -@testset "argument checking" begin - A = randn(Float64, (4, 4)) - B = randn(Float64, (4, 4)) - # aliasing must be rejected rather than silently producing garbage - @test_throws ArgumentError tensoradd!(A, A, ((2, 1), ()), false, 1, 0, tblis) - @test_throws ArgumentError tensorcontract!( - A, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1, 0, tblis - ) - # shape errors are reported before anything reaches the library - C = randn(Float64, (4, 3)) - @test_throws DimensionMismatch tensoradd!(C, A, ((1, 2), ()), false, 1, 0, tblis) - @test_throws DimensionMismatch tensorcontract!( - C, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1, 0, tblis - ) +# TBLIS requires all arguments to have the same element type +@testset "test suite" begin + TestSuite.test_backend(tblis; eltypes, mixed_eltypes = (), aliasing = true) end @testset "rejection of unsupported arguments" begin diff --git a/test/runtests.jl b/test/runtests.jl index 84453ad5..c13dcb04 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2,6 +2,8 @@ using ParallelTestRunner using TensorOperations testsuite = find_tests(@__DIR__) +# the backend test suite in `test/testsuite` is used by the tests, but is not a test itself +filter!(!startswith("testsuite/") ∘ first, testsuite) # --fast is passed by the CI on draft PRs, but there are no slow tests to skip args = parse_args(ARGS; custom = ["fast"]) diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl new file mode 100644 index 00000000..f48439aa --- /dev/null +++ b/test/testsuite/TestSuite.jl @@ -0,0 +1,285 @@ +# Based on the design of the GPUArrays.jl and MatrixAlgebraKit.jl test suites + +""" + TestSuite + +Suite of tests for the primitive operations `tensoradd!`, `tensortrace!` and `tensorcontract!` +(and their allocating counterparts) of a TensorOperations backend. + +Every result is compared to that of a reference backend, `BaseCopy()` by default, which is run +on CPU copies of the same inputs. Backends acting on other array types, such as GPU arrays, are +therefore checked against CPU results. `BaseCopy` itself is checked against straightforward +loop implementations through `TestSuite.LoopReference()`. + +Packages providing a backend can include this test suite as follows: + +```julia +using TensorOperations +include(joinpath(pkgdir(TensorOperations), "test", "testsuite", "TestSuite.jl")) +TestSuite.test_backend(MyBackend(); arraytype = MyArray) +``` + +See [`test_backend`](@ref) for the options to describe what the backend supports. +""" +module TestSuite + +using Test +using TensorOperations +using TensorOperations: AbstractBackend, BaseCopy, DefaultAllocator, Index2Tuple, IndexError, + linearize, numind +using LinearAlgebra: Diagonal +using Adapt: adapt +using Random +using VectorInterface: One, Zero + +const tests = Dict{String, Function}() + +""" + @testsuite name (backend, options) -> begin ... end + +Register the test function under `name`, which is then included in [`test_backend`](@ref) and +can also be called separately as `test_(backend; kwargs...)`. +""" +macro testsuite(name, ex) + safe_name = lowercase(replace(name, " " => "_", "/" => "_", "!" => "")) + fn = Symbol("test_", safe_name) + return quote + $(esc(fn))(backend; kwargs...) = $(esc(ex))(backend, Options(; kwargs...)) + @assert !haskey(tests, $name) "testsuite already exists" + tests[$name] = $(esc(fn)) + end +end + +""" + Options(; kwargs...) + +Description of the backend under test, which determines the tests that are run. See the +fields for the keyword arguments and their defaults. +""" +Base.@kwdef struct Options + "Type of arrays the backend acts on, e.g. `CuArray`; the inputs are moved to it with `Adapt.adapt`." + arraytype::Type = Array + "Element types to test." + eltypes::Tuple = (Float32, Float64, ComplexF32, ComplexF64) + """ + Combinations `(TC, TA, TB)` of element types of the output and the inputs, to test + arguments with different element types; `tensoradd!` and `tensortrace!` use `(TC, TA)`. + Pass `()` for backends that require all element types to be equal. + """ + mixed_eltypes::Tuple = ( + (ComplexF64, Float64, ComplexF64), (Float64, Float32, Float64), + (ComplexF32, Float64, Float32), (ComplexF64, ComplexF32, Float64), + ) + "Backend computing the expected results, on `Array`s." + reference = BaseCopy() + "Allocator passed to the backend." + allocator = DefaultAllocator() + "Support for non-contiguous (strided) views of `arraytype`." + views::Bool = true + "Support for views with negative strides." + negative_strides::Bool = true + "Support for `Adjoint` and `Transpose` inputs." + wrappers::Bool = true + "Support for arrays with dimensions of size zero." + zero_size::Bool = true + "Support for `Diagonal` arguments in contractions." + diagonal::Bool = false + "Whether the backend rejects outputs that alias an input with an `ArgumentError`." + aliasing::Bool = false +end + +""" + test_backend(backend; kwargs...) + +Run the entire test suite for `backend`. The keyword arguments describe the backend and its +capabilities, see [`Options`](@ref) for the available options. +""" +function test_backend(backend; kwargs...) + return @testset "$name" for (name, test) in sort!(collect(tests); by = first) + test(backend; kwargs...) + end +end + +const rng = Random.Xoshiro(1234) +seed_rng!(seed) = Random.seed!(rng, seed) + +precision(::Type{T}) where {T <: Number} = sqrt(eps(real(float(T)))) +precision(Ts::Type...) = maximum(precision, Ts) + +randarray(::Type{T}, sz::Dims) where {T} = randn(rng, T, sz) +randarray(::Type{T}, sz::Tuple{}) where {T} = fill(randn(rng, T)) + +# fill the output with `NaN` if `β` is zero, to verify that it is not read +maybe_poison!(C, β) = iszero(β) ? fill!(C, convert(eltype(C), NaN)) : C + +# non-contiguous views of a larger array, possibly with negative strides +function stridedview(::Type{T}, sz::Dims; reversed::Bool = false) where {T} + parent = randarray(T, 2 .* sz .+ 1) + inds = map(n -> reversed ? StepRange(2n, -2, 2) : StepRange(2, 2, 2n), sz) + return view(parent, inds...) +end + +# the different scalar factors, including `Zero()` and `One()`, a numeric zero `β` that also +# needs to be treated as a strong zero, and a zero `α` +function scalars(::Type{T}) where {T} + α, β = randn(rng, T), randn(rng, T) + return ((One(), Zero()), (One(), One()), (α, zero(T)), (α, β), (zero(T), β)) +end + +# Moving data between devices +# --------------------------- +""" + todevice(arraytype, x) + +Copy `x` to an array of type `arraytype`, preserving views and wrappers. +""" +todevice(AT::Type, x) = adapt(AT, deepcopy(x)) + +""" + tocpu(x) + +Copy the values of `x` to a new `Array`. +""" +tocpu(x::AbstractArray) = Array(adapt(Array, x)) + +""" + check_inplace(f!, backend, options, C, args...; rtol) + +Run `f!(C, args..., backend, allocator)` on copies of `C` and `args` on the device, and compare +the result to `f!` with the reference backend on CPU copies. Also verifies that `f!` returns +its output argument, and leaves its inputs untouched. +""" +function check_inplace(f!, backend, options::Options, C, args...; rtol::Real) + Cref = f!(tocpu(C), map(tocpu, args)..., options.reference, DefaultAllocator()) + + Cdev = todevice(options.arraytype, C) + argsdev = map(x -> todevice(options.arraytype, x), args) + Cout = f!(Cdev, argsdev..., backend, options.allocator) + @test Cout === Cdev + @test isapprox(tocpu(Cout), Cref; rtol) + @test all(map((x, y) -> tocpu(x) == tocpu(y), argsdev, args)) + return nothing +end + +""" + check_allocating(f, backend, options, args...; rtol) + +Run `f(args..., backend, allocator)` on copies of `args` on the device, and compare the +result to `f` with the reference backend on CPU copies. Also verifies that the result is +again an array of type `arraytype`. +""" +function check_allocating(f, backend, options::Options, args...; rtol::Real) + Cref = f(map(tocpu, args)..., options.reference, DefaultAllocator()) + + argsdev = map(x -> todevice(options.arraytype, x), args) + Cout = f(argsdev..., backend, options.allocator) + @test Cout isa options.arraytype + @test eltype(Cout) == eltype(Cref) + @test isapprox(tocpu(Cout), Cref; rtol) + return nothing +end + +# Reference implementation +# ------------------------ +""" + LoopReference() + +Backend that implements the primitive operations as loops over all elements, directly +following their definitions. It is slow, but its correctness can be verified by inspection, +and serves to verify the `BaseCopy` backend. +""" +struct LoopReference <: AbstractBackend end + +# Map `Zero()` and `One()` to the strong `Bool` versions +_scalar(x::Number) = x +_scalar(::Zero) = false +_scalar(::One) = true + +function _update(c, x, α, β) + α, β = _scalar(α), _scalar(β) + return iszero(β) ? α * x : β * c + α * x +end + +function TensorOperations.tensoradd!( + C::AbstractArray, A::AbstractArray, pA::Index2Tuple, conjA::Bool, + α::Number, β::Number, ::LoopReference, allocator = DefaultAllocator() + ) + op = conjA ? conj : identity + p = linearize(pA) + for I in CartesianIndices(C) + # dimension p[k] of A becomes dimension k of C + J = zeros(Int, ndims(A)) + for k in 1:ndims(C) + J[p[k]] = I[k] + end + C[I] = _update(C[I], op(A[J...]), α, β) + end + return C +end + +function TensorOperations.tensortrace!( + C::AbstractArray, A::AbstractArray, p::Index2Tuple, q::Index2Tuple, conjA::Bool, + α::Number, β::Number, ::LoopReference, allocator = DefaultAllocator() + ) + op = conjA ? conj : identity + lp, (q1, q2) = linearize(p), q + traced = CartesianIndices(map(i -> size(A, i), q1)) + for I in CartesianIndices(C) + J = zeros(Int, ndims(A)) + for k in 1:ndims(C) + J[lp[k]] = I[k] + end + x = zero(eltype(C)) + for K in traced + for k in eachindex(q1) + J[q1[k]] = J[q2[k]] = K[k] + end + x += op(A[J...]) + end + C[I] = _update(C[I], x, α, β) + end + return C +end + +function TensorOperations.tensorcontract!( + C::AbstractArray, + A::AbstractArray, pA::Index2Tuple, conjA::Bool, + B::AbstractArray, pB::Index2Tuple, conjB::Bool, + pAB::Index2Tuple, α::Number, β::Number, + ::LoopReference, allocator = DefaultAllocator() + ) + opA = conjA ? conj : identity + opB = conjB ? conj : identity + lpAB = linearize(pAB) + nA = length(pA[1]) + contracted = CartesianIndices(map(i -> size(A, i), pA[2])) + for I in CartesianIndices(C) + # the open indices of A followed by those of B; open index lpAB[k] is dimension k of C + O = zeros(Int, numind(pAB)) + for k in 1:ndims(C) + O[lpAB[k]] = I[k] + end + JA, JB = zeros(Int, ndims(A)), zeros(Int, ndims(B)) + for k in eachindex(pA[1]) + JA[pA[1][k]] = O[k] + end + for k in eachindex(pB[2]) + JB[pB[2][k]] = O[nA + k] + end + x = zero(eltype(C)) + for K in contracted + for k in eachindex(pA[2]) + JA[pA[2][k]] = JB[pB[1][k]] = K[k] + end + x += opA(A[JA...]) * opB(B[JB...]) + end + C[I] = _update(C[I], x, α, β) + end + return C +end + +include("tensoradd.jl") +include("tensortrace.jl") +include("tensorcontract.jl") + +end diff --git a/test/testsuite/tensoradd.jl b/test/testsuite/tensoradd.jl new file mode 100644 index 00000000..2ae7f579 --- /dev/null +++ b/test/testsuite/tensoradd.jl @@ -0,0 +1,109 @@ +# sizes of A and the permutation pA +const ADD_CASES = [ + ((3, 4, 5), ((3, 1), (2,))), + ((3, 4), ((1, 2), ())), # identity + ((2, 3, 4), ((), (3, 2, 1))), + ((), ((), ())), # zero-dimensional + ((5,), ((1,), ())), + ((1, 3, 1, 2), ((4, 2), (3, 1))), # dimensions of size one + ((2, 1, 3, 2, 2, 3), ((5, 3, 1), (6, 2, 4))), +] +const ADD_ZERO_SIZE_CASES = [ + ((3, 0, 2), ((2, 3, 1), ())), + ((0,), ((1,), ())), +] + +add_size(szA, pA) = map(i -> szA[i], linearize(pA)) + +function _add!(pA, conjA, α, β) + return (C, A, backend, allocator) -> tensoradd!(C, A, pA, conjA, α, β, backend, allocator) +end + +@testsuite "tensoradd!" (backend, options) -> begin + cases = options.zero_size ? [ADD_CASES; ADD_ZERO_SIZE_CASES] : ADD_CASES + @testset "$T" for T in options.eltypes + rtol = precision(T) + @testset "size $szA, pA = $pA" for (szA, pA) in cases + for conjA in (false, true), (α, β) in scalars(T) + A = randarray(T, szA) + C = maybe_poison!(randarray(T, add_size(szA, pA)), β) + check_inplace(_add!(pA, conjA, α, β), backend, options, C, A; rtol) + end + end + + szA, pA = (3, 4, 5), ((3, 1), (2,)) + α, β = randn(rng, T), randn(rng, T) + szC = add_size(szA, pA) + if options.views + @testset "strided views" for conjA in (false, true), β′ in (β, Zero()) + for (A, C) in ( + (stridedview(T, szA), randarray(T, szC)), + (randarray(T, szA), stridedview(T, szC)), + (stridedview(T, szA), stridedview(T, szC)), + ) + maybe_poison!(C, β′) + check_inplace(_add!(pA, conjA, α, β′), backend, options, C, A; rtol) + end + end + end + if options.views && options.negative_strides + @testset "negative strides" for conjA in (false, true) + A = stridedview(T, szA; reversed = true) + C = stridedview(T, szC; reversed = true) + check_inplace(_add!(pA, conjA, α, β), backend, options, C, A; rtol) + end + end + if options.wrappers + # `Adjoint` already carries a conjugation, which has to combine with `conjA` + @testset "$wrapper input" for wrapper in (adjoint, transpose) + for conjA in (false, true), pA in (((1, 2), ()), ((2,), (1,))) + A = wrapper(randarray(T, (4, 3))) + C = randarray(T, add_size(size(A), pA)) + check_inplace(_add!(pA, conjA, α, β), backend, options, C, A; rtol) + end + end + end + end + + # only the element types of the output and the first input + mixed_eltypes = unique(map(Ts -> Ts[1:2], options.mixed_eltypes)) + @testset "mixed eltypes C::$TC, A::$TA" for (TC, TA) in mixed_eltypes + szA, pA = (3, 4, 5), ((3, 1), (2,)) + rtol = precision(TC, TA) + for conjA in (false, true), (α, β) in scalars(real(TC)) + A = randarray(TA, szA) + C = maybe_poison!(randarray(TC, add_size(szA, pA)), β) + check_inplace(_add!(pA, conjA, α, β), backend, options, C, A; rtol) + end + end + + @testset "allocating" for T in options.eltypes + A = randarray(T, (3, 4, 5)) + B = randarray(T, (5, 3, 4)) + rtol = precision(T) + for conjA in (false, true), conjB in (false, true) + α, β = randn(rng, T), randn(rng, T) + f = (A, B, backend, allocator) -> tensoradd( + A, ((1, 2, 3), ()), conjA, B, ((2, 3, 1), ()), conjB, α, β, backend, allocator + ) + check_allocating(f, backend, options, A, B; rtol) + end + end + + @testset "argument checking" begin + T = first(options.eltypes) + A = todevice(options.arraytype, randarray(T, (3, 4, 5))) + C = todevice(options.arraytype, randarray(T, (5, 3, 4))) + allocator = options.allocator + # wrong number of indices + @test_throws IndexError tensoradd!(C, A, ((3, 1), ()), false, 1, 0, backend, allocator) + # not a permutation + @test_throws IndexError tensoradd!(C, A, ((3, 1, 1), ()), false, 1, 0, backend, allocator) + # wrong sizes + @test_throws DimensionMismatch tensoradd!(C, A, ((1, 2, 3), ()), false, 1, 0, backend, allocator) + if options.aliasing + B = todevice(options.arraytype, randarray(T, (4, 4))) + @test_throws ArgumentError tensoradd!(B, B, ((2, 1), ()), false, 1, 0, backend, allocator) + end + end +end diff --git a/test/testsuite/tensorcontract.jl b/test/testsuite/tensorcontract.jl new file mode 100644 index 00000000..4c49200b --- /dev/null +++ b/test/testsuite/tensorcontract.jl @@ -0,0 +1,177 @@ +# sizes of A and B, and the permutations pA, pB and pAB +const CONTRACT_CASES = [ + ((3, 6, 5, 3, 4), ((3, 1, 4), (2, 5)), (4, 2, 6, 3), ((3, 1), (2, 4)), ((3, 1, 4), (5, 2))), + ((4, 5), ((1,), (2,)), (5, 3), ((1,), (2,)), ((1, 2), ())), # matrix-matrix + ((5, 4), ((2,), (1,)), (3, 5), ((2,), (1,)), ((2, 1), ())), # transposed matrix-matrix + ((4, 5), ((1,), (2,)), (5,), ((1,), ()), ((1,), ())), # matrix-vector + ((5,), ((), (1,)), (5,), ((1,), ()), ((), ())), # inner product + ((3, 4), ((), (1, 2)), (4, 3), ((2, 1), ()), ((), ())), # full contraction + ((3, 4), ((1, 2), ()), (5,), ((), (1,)), ((3, 1), (2,))), # outer product + ((), ((), ()), (), ((), ()), ((), ())), # zero-dimensional + ((), ((), ()), (3, 4), ((), (1, 2)), ((2, 1), ())), # zero-dimensional times matrix + ((1, 3, 1), ((1, 3), (2,)), (3, 1), ((1,), (2,)), ((1, 2, 3), ())), # dimensions of size one + # index layouts that cannot be fused into matrix multiplications without permuting + ((2, 3, 2, 3, 2), ((5, 1, 3), (4, 2)), (3, 2, 2, 3), ((1, 4), (2, 3)), ((2, 4, 1), (5, 3))), +] +const CONTRACT_ZERO_SIZE_CASES = [ + ((3, 0), ((1,), (2,)), (0, 4), ((1,), (2,)), ((1, 2), ())), # contracted dimension of size zero + ((0, 3), ((1,), (2,)), (3, 4), ((1,), (2,)), ((1, 2), ())), +] + +function contract_size(szA, pA, szB, pB, pAB) + open = (map(i -> szA[i], pA[1])..., map(i -> szB[i], pB[2])...) + return map(i -> open[i], linearize(pAB)) +end + +function _contract!(pA, conjA, pB, conjB, pAB, α, β) + return (C, A, B, backend, allocator) -> tensorcontract!( + C, A, pA, conjA, B, pB, conjB, pAB, α, β, backend, allocator + ) +end + +@testsuite "tensorcontract!" (backend, options) -> begin + cases = options.zero_size ? [CONTRACT_CASES; CONTRACT_ZERO_SIZE_CASES] : CONTRACT_CASES + @testset "$T" for T in options.eltypes + rtol = precision(T) + @testset "sizes $szA, $szB, pA = $pA, pB = $pB, pAB = $pAB" for (szA, pA, szB, pB, pAB) in cases + for conjA in (false, true), conjB in (false, true), (α, β) in scalars(T) + A = randarray(T, szA) + B = randarray(T, szB) + C = maybe_poison!(randarray(T, contract_size(szA, pA, szB, pB, pAB)), β) + f! = _contract!(pA, conjA, pB, conjB, pAB, α, β) + check_inplace(f!, backend, options, C, A, B; rtol) + end + end + + szA, pA, szB, pB, pAB = first(CONTRACT_CASES) + szC = contract_size(szA, pA, szB, pB, pAB) + α, β = randn(rng, T), randn(rng, T) + if options.views + @testset "strided views" for conjA in (false, true), β′ in (β, Zero()) + f! = _contract!(pA, conjA, pB, !conjA, pAB, α, β′) + for (A, B, C) in ( + (stridedview(T, szA), randarray(T, szB), randarray(T, szC)), + (randarray(T, szA), stridedview(T, szB), randarray(T, szC)), + (randarray(T, szA), randarray(T, szB), stridedview(T, szC)), + (stridedview(T, szA), stridedview(T, szB), stridedview(T, szC)), + ) + maybe_poison!(C, β′) + check_inplace(f!, backend, options, C, A, B; rtol) + end + end + end + if options.views && options.negative_strides + @testset "negative strides" for conjA in (false, true) + A = stridedview(T, szA; reversed = true) + B = stridedview(T, szB; reversed = true) + C = stridedview(T, szC; reversed = true) + f! = _contract!(pA, conjA, pB, !conjA, pAB, α, β) + check_inplace(f!, backend, options, C, A, B; rtol) + end + end + if options.wrappers + @testset "$wrapperA and $wrapperB inputs" for wrapperA in (adjoint, transpose), + wrapperB in (adjoint, transpose) + + for conjA in (false, true), conjB in (false, true) + A = wrapperA(randarray(T, (5, 4))) + B = wrapperB(randarray(T, (3, 5))) + C = randarray(T, (4, 3)) + f! = _contract!(((1,), (2,)), conjA, ((1,), (2,)), conjB, ((1, 2), ()), α, β) + check_inplace(f!, backend, options, C, A, B; rtol) + end + end + end + if options.diagonal + @testset "Diagonal" for conjA in (false, true), β′ in (β, Zero()) + n = 4 + D, D2 = Diagonal(randarray(T, (n,))), Diagonal(randarray(T, (n,))) + A = randarray(T, (n, n, n)) + # C[i,j,k] = A[i,j,a] * D[a,k] + C = maybe_poison!(randarray(T, (n, n, n)), β′) + f! = _contract!(((1, 2), (3,)), conjA, ((1,), (2,)), !conjA, ((1, 2, 3), ()), α, β′) + check_inplace(f!, backend, options, C, A, D; rtol) + # C[k,i,j] = D[k,a] * A[i,a,j] + C = maybe_poison!(randarray(T, (n, n, n)), β′) + f! = _contract!(((1,), (2,)), conjA, ((2,), (1, 3)), !conjA, ((1, 2, 3), ()), α, β′) + check_inplace(f!, backend, options, C, D, A; rtol) + # C[i,j] = A[i,j,a] * D[a,a]: contracting both indices of the diagonal + A4 = randarray(T, (n, n, n, n)) + C = maybe_poison!(randarray(T, (n, n)), β′) + f! = _contract!(((1, 2), (3, 4)), conjA, ((1, 2), ()), !conjA, ((1, 2), ()), α, β′) + check_inplace(f!, backend, options, C, A4, D; rtol) + # C[a,c,b,d] = M[a,b] * D[c,d]: outer product + M = randarray(T, (n, n)) + C = maybe_poison!(randarray(T, (n, n, n, n)), β′) + f! = _contract!(((1, 2), ()), conjA, ((), (1, 2)), !conjA, ((1, 3, 2, 4), ()), α, β′) + check_inplace(f!, backend, options, C, M, D; rtol) + # diagonal times diagonal, into a dense matrix and into a diagonal + f! = _contract!(((1,), (2,)), conjA, ((1,), (2,)), !conjA, ((1, 2), ()), α, β′) + C = maybe_poison!(randarray(T, (n, n)), β′) + check_inplace(f!, backend, options, C, D, D2; rtol) + C = Diagonal(maybe_poison!(randarray(T, (n,)), β′)) + check_inplace(f!, backend, options, C, D, D2; rtol) + end + end + end + + @testset "mixed eltypes C::$TC, A::$TA, B::$TB" for (TC, TA, TB) in options.mixed_eltypes + szA, pA, szB, pB, pAB = first(CONTRACT_CASES) + rtol = precision(TC, TA, TB) + for conjA in (false, true), (α, β) in scalars(real(TC)) + A = randarray(TA, szA) + B = randarray(TB, szB) + C = maybe_poison!(randarray(TC, contract_size(szA, pA, szB, pB, pAB)), β) + f! = _contract!(pA, conjA, pB, !conjA, pAB, α, β) + check_inplace(f!, backend, options, C, A, B; rtol) + end + end + + @testset "allocating" for T in options.eltypes + rtol = precision(T) + # general contraction and full contraction to a zero-dimensional array + for (szA, pA, szB, pB, pAB) in CONTRACT_CASES[[1, 6]], conjA in (false, true) + A = randarray(T, szA) + B = randarray(T, szB) + α = randn(rng, T) + f = (A, B, backend, allocator) -> tensorcontract( + A, pA, conjA, B, pB, !conjA, pAB, α, backend, allocator + ) + check_allocating(f, backend, options, A, B; rtol) + end + # the value of a full contraction is obtained with `tensorscalar` + szA, pA, szB, pB, pAB = CONTRACT_CASES[6] + A = randarray(T, szA) + B = randarray(T, szB) + Cref = tensorcontract(A, pA, false, B, pB, false, pAB, One(), options.reference) + Adev, Bdev = todevice(options.arraytype, A), todevice(options.arraytype, B) + C = tensorcontract(Adev, pA, false, Bdev, pB, false, pAB, One(), backend, options.allocator) + @test tensorscalar(C) ≈ tensorscalar(Cref) rtol = rtol + end + + @testset "argument checking" begin + T = first(options.eltypes) + A = todevice(options.arraytype, randarray(T, (3, 4, 5))) + B = todevice(options.arraytype, randarray(T, (5, 4, 2))) + C = todevice(options.arraytype, randarray(T, (3, 2))) + allocator = options.allocator + pA, pB, pAB = ((1,), (2, 3)), ((2, 1), (3,)), ((1, 2), ()) + @test tensorcontract!(C, A, pA, false, B, pB, false, pAB, 1, 0, backend, allocator) === C + # wrong number of indices + @test_throws IndexError tensorcontract!(C, A, ((1,), (2,)), false, B, pB, false, pAB, 1, 0, backend, allocator) + @test_throws IndexError tensorcontract!(C, A, pA, false, B, ((2, 1), ()), false, pAB, 1, 0, backend, allocator) + @test_throws IndexError tensorcontract!(C, A, pA, false, B, pB, false, ((1,), ()), 1, 0, backend, allocator) + # different numbers of contracted indices + @test_throws IndexError tensorcontract!(C, A, pA, false, B, ((2,), (1, 3)), false, ((1, 2, 3), ()), 1, 0, backend, allocator) + # contracted dimensions of different size + @test_throws DimensionMismatch tensorcontract!(C, A, pA, false, B, ((1, 2), (3,)), false, pAB, 1, 0, backend, allocator) + # wrong output size + @test_throws DimensionMismatch tensorcontract!(C, A, pA, false, B, pB, false, ((2, 1), ()), 1, 0, backend, allocator) + if options.aliasing + M = todevice(options.arraytype, randarray(T, (4, 4))) + N = todevice(options.arraytype, randarray(T, (4, 4))) + p = ((1,), (2,)) + @test_throws ArgumentError tensorcontract!(M, M, p, false, N, p, false, ((1, 2), ()), 1, 0, backend, allocator) + end + end +end diff --git a/test/testsuite/tensortrace.jl b/test/testsuite/tensortrace.jl new file mode 100644 index 00000000..63629a49 --- /dev/null +++ b/test/testsuite/tensortrace.jl @@ -0,0 +1,103 @@ +# sizes of A, the permutation p of the open indices and the pairs q of traced indices +const TRACE_CASES = [ + ((5, 3, 4, 5, 3, 2), ((6, 3), ()), ((1, 2), (4, 5))), + ((2, 3, 4, 3), ((1,), (3,)), ((2,), (4,))), + ((4, 4), ((), ()), ((1,), (2,))), # full trace + ((3, 2, 3, 2), ((), ()), ((1, 2), (3, 4))), # full trace over two pairs + ((3, 4), ((2, 1), ()), ((), ())), # nothing to trace + ((3, 1, 2, 3, 1), ((3,), ()), ((1, 2), (4, 5))), # dimensions of size one +] +const TRACE_ZERO_SIZE_CASES = [ + ((0, 3, 0), ((2,), ()), ((1,), (3,))), # traced dimensions of size zero + ((2, 0, 2), ((2,), ()), ((1,), (3,))), +] + +trace_size(szA, p) = map(i -> szA[i], linearize(p)) + +function _trace!(p, q, conjA, α, β) + return (C, A, backend, allocator) -> tensortrace!(C, A, p, q, conjA, α, β, backend, allocator) +end + +@testsuite "tensortrace!" (backend, options) -> begin + cases = options.zero_size ? [TRACE_CASES; TRACE_ZERO_SIZE_CASES] : TRACE_CASES + @testset "$T" for T in options.eltypes + rtol = precision(T) + @testset "size $szA, p = $p, q = $q" for (szA, p, q) in cases + for conjA in (false, true), (α, β) in scalars(T) + A = randarray(T, szA) + C = maybe_poison!(randarray(T, trace_size(szA, p)), β) + check_inplace(_trace!(p, q, conjA, α, β), backend, options, C, A; rtol) + end + end + + szA, p, q = (5, 3, 4, 5, 3, 2), ((6, 3), ()), ((1, 2), (4, 5)) + α, β = randn(rng, T), randn(rng, T) + szC = trace_size(szA, p) + if options.views + @testset "strided views" for conjA in (false, true), β′ in (β, Zero()) + for (A, C) in ( + (stridedview(T, szA), randarray(T, szC)), + (randarray(T, szA), stridedview(T, szC)), + (stridedview(T, szA), stridedview(T, szC)), + ) + maybe_poison!(C, β′) + check_inplace(_trace!(p, q, conjA, α, β′), backend, options, C, A; rtol) + end + end + end + if options.views && options.negative_strides + @testset "negative strides" for conjA in (false, true) + A = stridedview(T, szA; reversed = true) + C = stridedview(T, szC; reversed = true) + check_inplace(_trace!(p, q, conjA, α, β), backend, options, C, A; rtol) + end + end + if options.wrappers + @testset "$wrapper input" for wrapper in (adjoint, transpose), conjA in (false, true) + A = wrapper(randarray(T, (4, 4))) + C = randarray(T, ()) + check_inplace(_trace!(((), ()), ((1,), (2,)), conjA, α, β), backend, options, C, A; rtol) + end + end + end + + # only the element types of the output and the first input + mixed_eltypes = unique(map(Ts -> Ts[1:2], options.mixed_eltypes)) + @testset "mixed eltypes C::$TC, A::$TA" for (TC, TA) in mixed_eltypes + szA, p, q = (5, 3, 4, 5, 3, 2), ((6, 3), ()), ((1, 2), (4, 5)) + rtol = precision(TC, TA) + for conjA in (false, true), (α, β) in scalars(real(TC)) + A = randarray(TA, szA) + C = maybe_poison!(randarray(TC, trace_size(szA, p)), β) + check_inplace(_trace!(p, q, conjA, α, β), backend, options, C, A; rtol) + end + end + + @testset "allocating" for T in options.eltypes + A = randarray(T, (5, 3, 4, 5, 3, 2)) + rtol = precision(T) + for conjA in (false, true) + α = randn(rng, T) + f = (A, backend, allocator) -> tensortrace( + A, ((6, 3), ()), ((1, 2), (4, 5)), conjA, α, backend, allocator + ) + check_allocating(f, backend, options, A; rtol) + end + end + + @testset "argument checking" begin + T = first(options.eltypes) + A = todevice(options.arraytype, randarray(T, (5, 3, 4, 5, 3))) + C = todevice(options.arraytype, randarray(T, (4,))) + allocator = options.allocator + # wrong number of indices + @test_throws IndexError tensortrace!(C, A, ((3,), ()), ((1,), (4,)), false, 1, 0, backend, allocator) + # index both open and traced + @test_throws IndexError tensortrace!(C, A, ((3,), ()), ((1, 3), (4, 5)), false, 1, 0, backend, allocator) + # traced dimensions of different size + @test_throws DimensionMismatch tensortrace!(C, A, ((3,), ()), ((1, 2), (5, 4)), false, 1, 0, backend, allocator) + # wrong output size + C3 = todevice(options.arraytype, randarray(T, (3,))) + @test_throws DimensionMismatch tensortrace!(C3, A, ((3,), ()), ((1, 2), (4, 5)), false, 1, 0, backend, allocator) + end +end