From 0aac2d6ec9eb4d3a4c757add374b395bbf4f46f7 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 2 Nov 2026 11:00:16 -0500 Subject: [PATCH 1/5] Move the test dependencies to test/Project.toml Co-Authored-By: Claude Opus 5.5 --- Project.toml | 34 ++--------------------------- test/Project.toml | 55 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 32 deletions(-) create mode 100644 test/Project.toml diff --git a/Project.toml b/Project.toml index 31c63a38..c6d2d6d9 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" 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/test/Project.toml b/test/Project.toml new file mode 100644 index 00000000..4f0133b9 --- /dev/null +++ b/test/Project.toml @@ -0,0 +1,55 @@ +name = "TensorOperationsTests" + +[deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" +CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" +cuRAND = "20fd9a0b-12d5-4c2f-a8af-7c34e9e60431" +cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1" +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" + +[sources] +TensorOperations = {path = ".."} + +[compat] +Adapt = "4" +AMDGPU = "2" +Aqua = "0.6, 0.7, 0.8" +Bumper = "0.6, 0.7" +ChainRulesCore = "1" +ChainRulesTestUtils = "1" +CUDACore = "6" +cuRAND = "6" +cuTENSOR = "6" +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" +TBLIS = "0.3" +Test = "1" +TupleTools = "1.6" +VectorInterface = "0.4.1, 0.5, 0.6, 0.7" From 1fb7fed79a193e3d5eba64735e75b4ceac4fb6fc Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 2 Nov 2026 11:00:16 -0500 Subject: [PATCH 2/5] Fix tensortrace! with traced dimensions of size zero for the Base backends Co-Authored-By: Claude Opus 5.5 --- src/implementation/base.jl | 4 ++++ src/utils.jl | 4 ++++ 2 files changed, 8 insertions(+) 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/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) From a1014248dc3fc4b9686dac1ac69f6840f9958aa9 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 2 Nov 2026 12:10:50 -0500 Subject: [PATCH 3/5] Fix edge cases of the cuTENSOR, TBLIS and StridedNative backends - cuTENSOR: check the indices and sizes of the arguments, and handle empty outputs, contractions and traces over dimensions of size zero, and traces without traced indices - TBLIS: handle traces over dimensions of size zero - StridedNative: reject outputs that alias an input of a contraction Co-Authored-By: Claude Opus 5.5 --- ext/TensorOperationsTBLISExt.jl | 2 ++ ext/TensorOperationscuTENSORExt.jl | 26 +++++++++++++++++++++----- src/implementation/strided.jl | 3 +++ 3 files changed, 26 insertions(+), 5 deletions(-) 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/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, β) From 92888bcc62c34b053dbc6a0f1109db0323d8b844 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 2 Nov 2026 12:10:50 -0500 Subject: [PATCH 4/5] Add a reusable test suite for the backends Every backend is compared to BaseCopy on CPU copies of the inputs, which is itself compared to plain loop implementations. Backends of other array types, e.g. GPU arrays, are therefore compared to CPU results. The test suite replaces the per-backend tests of the primitive operations, and adds many more edge cases. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 1 + test/Project.toml | 18 +- test/core/backends/basecopy.jl | 15 ++ test/core/backends/baseview.jl | 9 + test/core/backends/default.jl | 9 + test/core/backends/strided.jl | 20 +++ test/core/methods.jl | 172 ------------------- test/extensions/gpuarrays.jl | 116 +++---------- test/extensions/tblis.jl | 121 +------------ test/runtests.jl | 2 + test/testsuite/TestSuite.jl | 285 +++++++++++++++++++++++++++++++ test/testsuite/tensoradd.jl | 109 ++++++++++++ test/testsuite/tensorcontract.jl | 177 +++++++++++++++++++ test/testsuite/tensortrace.jl | 103 +++++++++++ 14 files changed, 767 insertions(+), 390 deletions(-) create mode 100644 test/core/backends/basecopy.jl create mode 100644 test/core/backends/baseview.jl create mode 100644 test/core/backends/default.jl create mode 100644 test/core/backends/strided.jl create mode 100644 test/testsuite/TestSuite.jl create mode 100644 test/testsuite/tensoradd.jl create mode 100644 test/testsuite/tensorcontract.jl create mode 100644 test/testsuite/tensortrace.jl 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/test/Project.toml b/test/Project.toml index 4f0133b9..9a063dde 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,15 +1,13 @@ name = "TensorOperationsTests" [deps] -Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" 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" -CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" -cuRAND = "20fd9a0b-12d5-4c2f-a8af-7c34e9e60431" -cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1" DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a" @@ -25,20 +23,21 @@ 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] -Adapt = "4" AMDGPU = "2" +Adapt = "4" Aqua = "0.6, 0.7, 0.8" Bumper = "0.6, 0.7" +CUDACore = "6" ChainRulesCore = "1" ChainRulesTestUtils = "1" -CUDACore = "6" -cuRAND = "6" -cuTENSOR = "6" DynamicPolynomials = "0.5, 0.6" Enzyme = "0.13.183" EnzymeTestUtils = "0.2" @@ -53,3 +52,6 @@ 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 From 57b749d865c50db5e1fc0fefec022a3a65604f69 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 3 Nov 2026 17:45:33 -0500 Subject: [PATCH 5/5] Require Strided v2.6.7 for strong-zero empty products Co-Authored-By: Claude Opus 5.5 --- Project.toml | 2 +- test/Project.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index c6d2d6d9..d242ad43 100644 --- a/Project.toml +++ b/Project.toml @@ -55,7 +55,7 @@ PackageExtensionCompat = "1" PrecompileTools = "1.1" Preferences = "1.4" PtrArrays = "1.2" -Strided = "2.6" +Strided = "2.6.7" StridedViews = "0.5" TBLIS = "0.3" TupleTools = "1.6" diff --git a/test/Project.toml b/test/Project.toml index 9a063dde..3dc75059 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -47,7 +47,7 @@ Logging = "1.6" Mooncake = "0.5" ParallelTestRunner = "2" Random = "1" -Strided = "2.6" +Strided = "2.6.7" TBLIS = "0.3" Test = "1" TupleTools = "1.6"