From 48c74ca937f6dce7493cb26705e381e03329da80 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 21 Oct 2026 13:44:54 +0200 Subject: [PATCH 1/3] hipTENSOR support in TensorOperations --- Project.toml | 2 +- ext/TensorOperationsAMDGPUExt.jl | 314 ++++++++++++++++ src/TensorOperations.jl | 2 +- src/backends.jl | 11 + src/indexnotation/tensormacros.jl | 26 ++ test/hiptensor.jl | 587 ++++++++++++++++++++++++++++++ test/runtests.jl | 3 + 7 files changed, 943 insertions(+), 2 deletions(-) create mode 100644 test/hiptensor.jl diff --git a/Project.toml b/Project.toml index f82ce567..89de922a 100644 --- a/Project.toml +++ b/Project.toml @@ -41,7 +41,7 @@ TensorOperationscuTENSORExt = "cuTENSOR" TensorOperationsJLArraysExt = "JLArrays" [compat] -AMDGPU = "2" +AMDGPU = "2.8" Aqua = "0.6, 0.7, 0.8" Adapt = "4" Bumper = "0.6, 0.7" diff --git a/ext/TensorOperationsAMDGPUExt.jl b/ext/TensorOperationsAMDGPUExt.jl index bb64aab1..70463e22 100644 --- a/ext/TensorOperationsAMDGPUExt.jl +++ b/ext/TensorOperationsAMDGPUExt.jl @@ -1,8 +1,28 @@ module TensorOperationsAMDGPUExt using AMDGPU +using AMDGPU.hipTENSOR using TensorOperations using TensorOperations: TensorOperations as TO +using TensorOperations: hipTENSORBackend, AMDAllocator, DefaultAllocator +using TensorOperations: isstrided + +using Strided +using TupleTools: TupleTools as TT + +hiptensor_working = AMDGPU.hipTENSOR.has_hiptensor() + +using AMDGPU.hipTENSOR: OP_IDENTITY, OP_CONJ, OP_ADD +using AMDGPU.hipTENSOR: is_unary, is_binary +using AMDGPU.hipTENSOR: hiptensorWorksizePreference_t, hiptensorAlgo_t, hiptensorOperationDescriptor_t, + hiptensorOperator_t, hiptensorJitMode_t, hiptensorPlanPreference_t +using AMDGPU.hipTENSOR: WORKSPACE_DEFAULT, ALGO_DEFAULT, JIT_MODE_NONE +using AMDGPU.hipTENSOR: hiptensorCreatePlanPreference, hiptensorPlan, hipTensorPlan, + hipTensorDescriptor, ModeType + +using AMDGPU.hipTENSOR: elementwise_binary_execute!, permute!, contract!, reduce! +using AMDGPU.hipTENSOR: reduction_compute_types, hiptensorCreateReduction +using AMDGPU.Adapt: adapt #------------------------------------------------------------------------------------------- # Allocator @@ -100,4 +120,298 @@ function Base.resize!(buffer::ROCBufferAllocator{B}, n::Integer) where {B} return buffer end +# Literal copy of the StridedViewsAMDGPUExt module +const ROCStridedView{T, N, A <: ROCArray{T}} = StridedView{T, N, A} + +#------------------------------------------------------------------------------------------- +# @hiptensor macro +#------------------------------------------------------------------------------------------- +function TensorOperations._hiptensor(src, ex...) + # TODO: there is no check for doubled tensor kwargs + return Expr( + :macrocall, GlobalRef(TensorOperations, Symbol("@tensor")), + src, + Expr(:(=), :backend, Expr(:call, GlobalRef(TensorOperations, :hipTENSORBackend))), + Expr(:(=), :allocator, Expr(:call, GlobalRef(TensorOperations, :AMDAllocator))), + ex... + ) +end + +#------------------------------------------------------------------------------------------- +# Backend selection and passing +#------------------------------------------------------------------------------------------- +# A Base wrapper over `ROCArray` 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 `ROCStridedView` here. TODO handle the case where hipTensor isn't +# functional. +function TO.select_backend(::typeof(TO.tensoradd!), ::ROCStridedView, ::ROCStridedView) + return hipTENSORBackend() +end +function TO.select_backend(::typeof(TO.tensortrace!), ::ROCStridedView, ::ROCStridedView) + return hipTENSORBackend() +end +function TO.select_backend( + ::typeof(TO.tensorcontract!), ::ROCStridedView, ::ROCStridedView, ::ROCStridedView + ) + return hipTENSORBackend() +end + +# Make sure that if the `hipTensorBackend` is specified, arrays are converted to ROCArrays +function TO.tensoradd!( + C::AbstractArray, + A::AbstractArray, pA::Index2Tuple, conjA::Bool, + α::Number, β::Number, + backend::hipTENSORBackend, allocator = AMDAllocator() + ) + C_hip, isview = _hipstrided(C, allocator) + A_hip, = _hipstrided(A, allocator) + tensoradd!(C_hip, A_hip, pA, conjA, α, β, backend, allocator) + isview || copy!(C, C_hip.parent) + return C +end +function TO.tensorcontract!( + C::AbstractArray, + A::AbstractArray, pA::Index2Tuple, conjA::Bool, + B::AbstractArray, pB::Index2Tuple, conjB::Bool, + pAB::Index2Tuple, + α::Number, β::Number, + backend::hipTENSORBackend, allocator = AMDAllocator() + ) + C_hip, isview = _hipstrided(C, allocator) + A_hip, = _hipstrided(A, allocator) + B_hip, = _hipstrided(B, allocator) + tensorcontract!( + C_hip, A_hip, pA, conjA, B_hip, pB, conjB, pAB, α, β, backend, allocator + ) + isview || copy!(C, C_hip.parent) + return C +end +function TO.tensortrace!( + C::AbstractArray, + A::AbstractArray, p::Index2Tuple, q::Index2Tuple, conjA::Bool, + α::Number, β::Number, + backend::hipTENSORBackend, allocator = AMDAllocator() + ) + C_hip, isview = _hipstrided(C, allocator) + A_hip, = _hipstrided(A, allocator) + tensortrace!(C_hip, A_hip, p, q, conjA, α, β, backend, allocator) + isview || copy!(C, C_hip.parent) + return C +end + +_hipstrided(A::AbstractArray, ::DefaultAllocator) = _hipstrided(A, AMDAllocator()) +# input tensors are marshalled to the device in the same way as for the default allocator: +# the buffer is reserved for the temporaries that are created by the tensor operations +_hipstrided(A::AbstractArray, ::TO.BufferAllocator) = _hipstrided(A, AMDAllocator()) +function _hipstrided(A::AbstractArray, allocator::AMDAllocator) + if isstrided(A) + return _hipstrided(StridedView(A), allocator) + else + return StridedView(ROCArray{eltype(A), ndims(A)}(A)), false + end +end +function _hipstrided( + A::StridedView, allocator::AMDAllocator + ) + P = A.parent + if P isa ROCArray + return A, true + else + P_hip = ROCArray{eltype(P), ndims(P)}(P) + return StridedView(P_hip, A.size, A.strides, A.offset, A.op), false + end +end + +#------------------------------------------------------------------------------------------- +# Implementation +#------------------------------------------------------------------------------------------- +function TO.tensorscalar(C::ROCStridedView) + return ndims(C) == 0 ? AMDGPU.@allowscalar(C[]) : throw(DimensionMismatch()) +end + +function tensorop(A::ROCStridedView, conjA::Bool = false) + return (eltype(A) <: Real || !xor(conjA, A.op === conj)) ? OP_IDENTITY : OP_CONJ +end + +isdense(A::ROCStridedView) = strides(A) == Base.size_to_strides(1, size(A)...) + +# hipTENSOR ignores `OP_CONJ` in contractions, so conjugated operands are materialized +function materialize_conj(A::ROCStridedView, conjA::Bool) + tensorop(A, conjA) == OP_IDENTITY && return A + A′ = StridedView(ROCArray{eltype(A)}(undef, size(A))) + A′ .= conjA ? conj.(A) : A + return A′ +end + +function TO.tensoradd!( + C::ROCStridedView, + A::ROCStridedView, pA::Index2Tuple, conjA::Bool, + α::Number, β::Number, + backend::hipTENSORBackend, allocator + ) + # hipTENSOR does not validate sizes itself + TO.argcheck_tensoradd(C, A, TO.linearize(pA)) + TO.dimcheck_tensoradd(C, A, TO.linearize(pA)) + + # hipTENSOR's permutations and elementwise operations ignore the layout of the output, + # so go through a dense temporary instead + if !isdense(C) + C′ = StridedView(ROCArray{eltype(C)}(undef, size(C))) + iszero(β) || (C′ .= C) + TO.tensoradd!(C′, A, pA, conjA, α, β, backend, allocator) + C .= C′ + AMDGPU.unsafe_free!(C′.parent) + return C + end + + # convert arguments + Ainds, Cinds = collect.(TO.add_labels(pA)) + opA = tensorop(A, conjA) + + # dispatch to hipTENSOR + return if iszero(β) + permute!(α, A, Ainds, opA, C, Cinds) + else + elementwise_binary_execute!( + α, + A, Ainds, opA, + β, + C, Cinds, OP_IDENTITY, + C, Cinds, + OP_ADD + ) + end +end + +function TO.tensorcontract!( + C::ROCStridedView, + A::ROCStridedView, pA::Index2Tuple, conjA::Bool, + B::ROCStridedView, pB::Index2Tuple, conjB::Bool, + pAB::Index2Tuple, + α::Number, β::Number, + backend::hipTENSORBackend, allocator + ) + # hipTENSOR does not validate sizes itself + TO.argcheck_tensorcontract(C, A, pA, B, pB, pAB) + TO.dimcheck_tensorcontract(C, A, pA, B, pB, pAB) + + # convert arguments + Ainds, Binds, Cinds = collect.(TO.contract_labels(pA, pB, pAB)) + A′ = materialize_conj(A, conjA) + B′ = materialize_conj(B, conjB) + + # hipTENSOR reads C even if β is zero, so NaNs in uninitialized memory would propagate + iszero(β) && (C .= zero(eltype(C))) + + # dispatch to hipTENSOR + contract!( + α, + A′, Ainds, OP_IDENTITY, + B′, Binds, OP_IDENTITY, + β, + C, Cinds, OP_IDENTITY, + OP_IDENTITY + ) + A′ === A || AMDGPU.unsafe_free!(A′.parent) + B′ === B || AMDGPU.unsafe_free!(B′.parent) + return C +end + +function TO.tensortrace!( + C::ROCStridedView, + A::ROCStridedView, p::Index2Tuple, q::Index2Tuple, conjA::Bool, + α::Number, β::Number, + backend::hipTENSORBackend, allocator + ) + # hipTENSOR does not validate sizes itself + TO.argcheck_tensortrace(C, A, TO.linearize(p), q) + TO.dimcheck_tensortrace(C, A, TO.linearize(p), q) + + # convert arguments + Ainds, Cinds = collect.(TO.trace_labels(p, q)) + opA = tensorop(A, conjA) + + # map to reduction operation + plan = plan_trace(A, Ainds, opA, C, Cinds, OP_IDENTITY, OP_ADD) + reduce!(plan, α, A, β, C) + AMDGPU.unsafe_free!(plan) # the plan-taking `reduce!` does not free the plan + return C +end + +function hipTENSOR.hipTensorDescriptor( + a::ROCStridedView; + size = size(a), strides = strides(a), eltype = eltype(a) + ) + sz = collect(Int64, size) + st = collect(Int64, strides) + alignment = UInt32(find_alignment(a)) + return hipTENSOR.hipTensorDescriptor(sz, st, eltype, alignment) +end + +const MAX_ALIGNMENT = UInt(1) # This is the largest alignment of AMD memory +"find the alignment of the first element of the view" +find_alignment(A::ROCStridedView) = gcd(MAX_ALIGNMENT, convert(UInt, pointer(A))) + +# trace! +# ------ +# not actually part of hipTENSOR, just a special case of reduce +function plan_trace( + @nospecialize(A::AbstractArray), Ainds::ModeType, opA::hiptensorOperator_t, + @nospecialize(C::AbstractArray), Cinds::ModeType, opC::hiptensorOperator_t, + opReduce::hiptensorOperator_t; + jit::hiptensorJitMode_t = JIT_MODE_NONE, + workspace::hiptensorWorksizePreference_t = WORKSPACE_DEFAULT, + algo::hiptensorAlgo_t = ALGO_DEFAULT, + compute_type::Union{DataType, Nothing} = nothing + ) + !is_unary(opA) && throw(ArgumentError("opA must be a unary op!")) + !is_unary(opC) && throw(ArgumentError("opC must be a unary op!")) + !is_binary(opReduce) && throw(ArgumentError("opReduce must be a binary op!")) + + # TODO: check if this can be avoided, available in caller + # TODO: hipTENSOR will allocate sizes and strides anyways, could use that here + p, q = TO.trace_indices(tuple(Ainds...), tuple(Cinds...)) + qsorted = TT.sort(q[2]) + # add strides of cindA2 to strides of cindA1 -> selects diagonal + stA = strides(A) + for (i, j) in zip(q...) + stA = Base.setindex(stA, stA[i] + stA[j], i) + end + szA = TT.deleteat(size(A), qsorted) + stA′ = TT.deleteat(stA, qsorted) + + # hipTENSOR needs A's modes in the same order as C's, followed by the traced ones + modesA = deleteat!(Ainds, qsorted) + order = vcat(Cinds, filter(∉(Cinds), modesA)) + perm = map(m -> findfirst(==(m), modesA), order) + + descA = hipTensorDescriptor(A; size = szA[perm], strides = stA′[perm]) + descC = hipTensorDescriptor(C) + + modeA = collect(Cint, order) + modeC = collect(Cint, Cinds) + + actual_compute_type = if compute_type === nothing + reduction_compute_types[(eltype(A), eltype(C))] + else + compute_type + end + + desc = Ref{hiptensorOperationDescriptor_t}() + hiptensorCreateReduction( + hipTENSOR.handle(), + desc, + descA, modeA, opA, + descC, modeC, opC, + descC, modeC, opReduce, + actual_compute_type + ) + + plan_pref = Ref{hiptensorPlanPreference_t}() + hiptensorCreatePlanPreference(hipTENSOR.handle(), plan_pref, algo, jit) + + # the plan takes ownership of `desc` and `plan_pref` and destroys them when it is freed + return hipTensorPlan(desc[], plan_pref[], actual_compute_type, Any[descA, descC, modeA, modeC]; workspacePref = workspace) +end + end diff --git a/src/TensorOperations.jl b/src/TensorOperations.jl index 499dbe30..057159ab 100644 --- a/src/TensorOperations.jl +++ b/src/TensorOperations.jl @@ -16,7 +16,7 @@ using Base.Meta: isexpr #--------- # export macro API export @tensor, @tensoropt, @tensoropt_verbose, @optimalcontractiontree, @notensor, @ncon -export @cutensor, @butensor +export @cutensor, @hiptensor, @butensor # export function based API export ncon diff --git a/src/backends.jl b/src/backends.jl index a49a0a65..b8ec4e9e 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -97,6 +97,7 @@ struct TBLISBackend <: AbstractBackend end #----------------- """ cuTENSORBackend(; workspace = :default) + Backend for tensor operations that is based on the NVIDIA cuTENSOR library. `workspace` is cuTENSOR's workspace preference for contractions: `:min`, `:default` or `:max`. """ @@ -113,3 +114,13 @@ cuTENSOR's workspace preference for contractions: `:min`, `:default` or `:max`. return new(workspace) end end + + +# hipTENSOR backend +#----------------- +""" + hipTENSORBackend() + +Backend for tensor operations that is based on the AMD hipTENSOR library. +""" +struct hipTENSORBackend <: AbstractBackend end diff --git a/src/indexnotation/tensormacros.jl b/src/indexnotation/tensormacros.jl index 31407708..f7c8d689 100644 --- a/src/indexnotation/tensormacros.jl +++ b/src/indexnotation/tensormacros.jl @@ -319,6 +319,32 @@ end function _cutensor end +""" + @hiptensor tensor_expr + +Use the GPU to perform all tensor operations, through the use of the hipTENSOR library. This +will transfer all arrays to the GPU before performing the requested operations. If the +output is an existing host array, the result will be transferred back. If a new array is +created (i.e. using `:=`), it will remain on the GPU device and it is up to the user to +transfer it back. This macro is equivalent to + + @tensor backend = hipTENSORBackend() allocator = AMDAllocator() tensor_expr + +!!! note + This macro requires the hipTENSOR library to be installed and loaded. This can be + achieved by running `using AMDGPU.hipTENSOR` or `import AMDGPU.hipTENSOR` before using the macro. + Not every AMDGPU supports hipTENSOR, so you should also run `hipTENSOR.functional()` to check + whether the library has fully loaded. +""" +macro hiptensor(ex...) + if length(methods(_hiptensor)) == 0 + throw(ArgumentError("AMDGPU.hipTENSOR not loaded: add `using AMDGPU.hipTENSOR` or `import AMDGPU.hipTENSOR` before using `@hiptensor`")) + end + return esc(_hiptensor(__source__, ex...)) +end + +function _hiptensor end + """ @butensor tensor_expr diff --git a/test/hiptensor.jl b/test/hiptensor.jl new file mode 100644 index 00000000..35925bfd --- /dev/null +++ b/test/hiptensor.jl @@ -0,0 +1,587 @@ +@testset "@hiptensor dependency check" begin + @test_throws ArgumentError begin + ex = :(@hiptensor A[a, b, c, d] := B[a, b, c, d]) + macroexpand(Main, ex) + end +end + +using AMDGPU, AMDGPU.hipTENSOR +using AMDGPU.hipTENSOR: hipTENSOR +if hipTENSOR.has_hiptensor() + using LinearAlgebra: norm + using TensorOperations: IndexError + using TensorOperations: hipTENSORBackend, AMDAllocator + using TensorOperations: BufferAllocator, AMDBufferAllocator + using TensorOperations: tensoralloc, tensorfree! + using TensorOperations: allocator_checkpoint!, allocator_reset! + + @testset "elementary operations" verbose = true begin + @testset "tensorcopy" begin + A = randn(Float32, (3, 5, 4, 6)) + @tensor C1[4, 1, 3, 2] := A[1, 2, 3, 4] + @tensor C2[4, 1, 3, 2] := ROCArray(A)[1, 2, 3, 4] + @test C2 isa ROCArray + @test collect(C2) ≈ C1 + end + + @testset "tensoradd" begin + A = randn(Float32, (5, 6, 3, 4)) + B = randn(Float32, (5, 6, 3, 4)) + α = randn(Float32) + @tensor C1[a, b, c, d] := A[a, b, c, d] + α * B[a, b, c, d] + @tensor C2[a, b, c, d] := ROCArray(A)[a, b, c, d] + α * ROCArray(B)[a, b, c, d] + @test collect(C2) ≈ C1 + + # complex doens't work yet + #=C = randn(ComplexF32, (5, 6, 3, 4)) + D = randn(ComplexF32, (5, 3, 4, 6)) + β = randn(ComplexF32) + @tensor E1[a, b, c, d] := C[a, b, c, d] + β * conj(D[a, c, d, b]) + @tensor E2[a, b, c, d] := ROCArray(C)[a, b, c, d] + + β * conj(ROCArray(D)[a, c, d, b]) + @test collect(E2) ≈ E1=# + end + + @testset "tensortrace" begin + A = randn(Float32, (5, 10, 10)) + @tensor B1[a] := A[a, b′, b′] + @tensor B2[a] := ROCArray(A)[a, b′, b′] + @test collect(B2) ≈ B1 + + # not supported yet + #=C = randn(ComplexF32, (3, 20, 5, 3, 20, 4, 5)) + @tensor D1[e, a, d] := C[a, b, c, d, b, e, c] + @tensor D2[e, a, d] := ROCArray(C)[a, b, c, d, b, e, c] + @test collect(D2) ≈ D1 + + @tensor D3[a, e, d] := conj(C[a, b, c, d, b, e, c]) + @tensor D4[a, e, d] := conj(ROCArray(C)[a, b, c, d, b, e, c]) + @test collect(D4) ≈ D3 + + α = randn(ComplexF32) + @tensor D5[d, e, a] := α * C[a, b, c, d, b, e, c] + @tensor D6[d, e, a] := α * ROCArray(C)[a, b, c, d, b, e, c] + @test collect(D6) ≈ D5=# + end + + @testset "tensorcontract" begin + A = randn(Float32, (3, 20, 5, 3, 4)) + B = randn(Float32, (5, 6, 20, 3)) + @tensor C1[a, g, e, d, f] := A[a, b, c, d, e] * B[c, f, b, g] + @tensor C2[a, g, e, d, f] := ROCArray(A)[a, b, c, d, e] * ROCArray(B)[c, f, b, g] + @test collect(C2) ≈ C1 + + D = randn(ComplexF32, (5, 5, 5)) + E = rand(ComplexF32, (5, 5, 5)) + @tensor F1[a, b, c, d, e, f] := D[a, b, c] * conj(E[d, e, f]) + @tensor F2[a, b, c, d, e, f] := ROCArray(D)[a, b, c] * conj(ROCArray(E)[d, e, f]) + @test collect(F2) ≈ F1 + # not yet supported + #=D = randn(Float64, (5, 5, 5)) + E = rand(ComplexF64, (5, 5, 5)) + @tensor F1[a, b, c, d, e, f] := D[a, b, c] * conj(E[d, e, f]) + @tensor F2[a, b, c, d, e, f] := ROCArray(D)[a, b, c] * conj(ROCArray(E)[d, e, f]) + @test collect(F2) ≈ F1=# + end + end + + @testset "more complicated expressions" verbose = true begin + Da, Db, Dc, Dd, De, Df, Dg, Dh = 10, 15, 4, 8, 6, 7, 3, 2 + A = rand(Float32, (Dc, Da, Df, Da, De, Db, Db, Dg)) + B = rand(Float32, (Dc, Dh, Dg, De, Dd)) + C = rand(Float32, (Dd, Dh, Df)) + + @tensor D1[d, f, h] := A[c, a, f, a, e, b, b, g] * B[c, h, g, e, d] + + Float32(0.5) * C[d, h, f] + @tensor D2[d, f, h] := ROCArray(A)[c, a, f, a, e, b, b, g] * + ROCArray(B)[c, h, g, e, d] + Float32(0.5) * ROCArray(C)[d, h, f] + @test collect(D2) ≈ D1 + + @test norm(vec(D1)) ≈ sqrt(abs(@tensor D1[d, f, h] * conj(D1[d, f, h]))) + @test norm(D2) ≈ sqrt(abs(@tensor D2[d, f, h] * conj(D2[d, f, h]))) + + @testset "readme example" begin + α = randn() + A = randn(5, 5, 5, 5, 5, 5) + B = randn(5, 5, 5) + C = randn(5, 5, 5) + D = zeros(5, 5, 5) + D2 = ROCArray(D) + @tensor begin + D[a, b, c] = A[a, e, f, c, f, g] * B[g, b, e] + α * C[c, a, b] + E[a, b, c] := A[a, e, f, c, f, g] * B[g, b, e] + α * C[c, a, b] + end + @tensor begin + D2[a, b, c] = ROCArray(A)[a, e, f, c, f, g] * ROCArray(B)[g, b, e] + + α * ROCArray(C)[c, a, b] + E2[a, b, c] := ROCArray(A)[a, e, f, c, f, g] * ROCArray(B)[g, b, e] + + α * ROCArray(C)[c, a, b] + end + @test collect(D2) ≈ D + @test collect(E2) ≈ E + end + + @testset "tensor network examples ($T)" for T in + (Float32, Float64, ComplexF32, ComplexF64) + D1, D2, D3 = 30, 40, 20 + d1, d2 = 2, 3 + + A1 = randn(T, D1, d1, D2) + A2 = randn(T, D2, d2, D3) + ρₗ = randn(T, D1, D1) + ρᵣ = randn(T, D3, D3) + H = randn(T, d1, d2, d1, d2) + + @tensor begin + HRAA1[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + @tensor begin + HRAA2[a, s1, s2, c] := ROCArray(ρₗ)[a, a'] * ROCArray(A1)[a', t1, b] * + ROCArray(A2)[b, t2, c'] * ROCArray(ρᵣ)[c', c] * ROCArray(H)[s1, s2, t1, t2] + end + @test HRAA2 isa ROCArray{T} + @test collect(HRAA2) ≈ HRAA1 + + allocator = AMDAllocator() + @tensor backend = hipTENSORBackend() allocator = allocator begin + HRAA3[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + @test HRAA3 isa ROCArray{T, 4} + @test collect(HRAA3) ≈ HRAA1 + + @tensor begin + E1 = ρₗ[a', a] * A1[a, s, b] * A2[b, s', c] * ρᵣ[c, c'] * H[t, t', s, s'] * + conj(A1[a', t, b']) * conj(A2[b', t', c']) + E2 = ROCArray(ρₗ)[a', a] * ROCArray(A1)[a, s, b] * ROCArray(A2)[b, s', c] * + ROCArray(ρᵣ)[c, c'] * ROCArray(H)[t, t', s, s'] * + conj(ROCArray(A1)[a', t, b']) * conj(ROCArray(A2)[b', t', c']) + end + @tensor backend = hipTENSORBackend() allocator = AMDAllocator() begin + E3 = ρₗ[a', a] * A1[a, s, b] * A2[b, s', c] * ρᵣ[c, c'] * H[t, t', s, s'] * + conj(A1[a', t, b']) * conj(A2[b', t', c']) + end + @test E1 ≈ E2 ≈ E3 + end + end + + @testset "BufferAllocator" verbose = true begin + # is the memory of `A` taken from `buffer`? + function isbufferbacked(A, buffer) + iszero(length(buffer)) && return false + base = UInt(pointer(buffer)) + return base ≤ UInt(pointer(A)) < base + length(buffer) + end + + @testset "Constructor and basic properties" begin + buffer = AMDBufferAllocator(; sizehint = 1024) + @test buffer isa BufferAllocator{ROCArray{UInt8, 1, AMDGPU.Mem.HIPBuffer}} + @test length(buffer) == 1024 + @test isempty(buffer) + @test buffer.offset == 0 + + # resizing frees the old buffer and allocates a new one + resize!(buffer, 3000) + @test length(buffer) == 4096 + @test isempty(buffer) + buffer.offset = 100 + @test_throws ErrorException resize!(buffer, 8192) + empty!(buffer) + @test length(resize!(buffer, 8192)) == 8192 + end + + @testset "tensoralloc" begin + buffer = AMDBufferAllocator(; sizehint = 4096) + + # temporaries are taken from the buffer, in the memory space of the buffer + C1 = tensoralloc(ROCArray{Float32, 2}, (8, 8), Val(true), buffer) + @test C1 isa ROCArray{Float32, 2, AMDGPU.Mem.HIPBuffer} + @test size(C1) == (8, 8) + @test isbufferbacked(C1, buffer) + @test buffer.offset == 8 * 8 * sizeof(Float32) + + # non-temporaries are not + offset = buffer.offset + C2 = tensoralloc(ROCArray{Float32, 2}, (8, 8), Val(false), buffer) + @test C2 isa ROCArray{Float32, 2} + @test !isbufferbacked(C2, buffer) + @test buffer.offset == offset + + # freeing a buffer-backed tensor does not invalidate the buffer + ptr1 = pointer(C1) + tensorfree!(C1, buffer) + allocator_reset!(buffer, 0) + C3 = tensoralloc(ROCArray{Float32, 2}, (8, 8), Val(true), buffer) + @test isbufferbacked(C3, buffer) + @test pointer(C3) == ptr1 + fill!(C3, 1.0f0) + @test all(isone, collect(C3)) + + end + + @testset "storage mismatch falls back" begin + # a host buffer cannot back ROCArrays + hostbuffer = BufferAllocator(; sizehint = 4096) + C1 = tensoralloc(ROCArray{Float32, 2}, (8, 8), Val(true), hostbuffer) + @test C1 isa ROCArray{Float32, 2} + @test hostbuffer.offset == 0 + + # a device buffer cannot back Arrays + cubuffer = AMDBufferAllocator(; sizehint = 4096) + C2 = tensoralloc(Array{Float64, 2}, (8, 8), Val(true), cubuffer) + @test C2 isa Array{Float64, 2} + @test cubuffer.offset == 0 + end + + @testset "alignment" begin + # cuTENSOR is sensitive to this: it only selects its fastest kernels for + # 256-byte aligned data, so every temporary has to be padded to that + buffer = AMDBufferAllocator(; sizehint = 8192) + @test iszero(UInt(pointer(buffer)) % 256) + + # a deliberately misaligning allocation of 3 bytes + C1 = tensoralloc(ROCArray{UInt8, 1}, (3,), Val(true), buffer) + @test isbufferbacked(C1, buffer) + @test buffer.offset == 3 + for T in (Float32, Float64, ComplexF32, ComplexF64) + C2 = tensoralloc(ROCArray{T, 1}, (4,), Val(true), buffer) + @test isbufferbacked(C2, buffer) + @test iszero(UInt(pointer(C2)) % 256) + end + + # host buffers stay at the alignment Julia actually guarantees + @test TensorOperations.buffer_alignment(BufferAllocator()) == 16 + end + + @testset "checkpoint and reset" begin + buffer = AMDBufferAllocator(; sizehint = 4096) + cp0 = allocator_checkpoint!(buffer) + @test cp0 == 0 + + C1 = tensoralloc(ROCArray{Float32, 2}, (8, 8), Val(true), buffer) + cp1 = allocator_checkpoint!(buffer) + @test cp1 > cp0 + C2 = tensoralloc(ROCArray{Float32, 2}, (8, 8), Val(true), buffer) + @test pointer(C2) != pointer(C1) + + allocator_reset!(buffer, cp1) + @test buffer.offset == cp1 + @test_throws ArgumentError allocator_reset!(buffer, cp1 + 10) + + allocator_reset!(buffer, cp0) + @test isempty(buffer) + end + + @testset "tensor network ($T)" for T in + (Float32, Float64, ComplexF32, ComplexF64) + D1, D2, D3 = 30, 40, 20 + d1, d2 = 2, 3 + + A1 = ROCArray(randn(T, D1, d1, D2)) + A2 = ROCArray(randn(T, D2, d2, D3)) + ρₗ = ROCArray(randn(T, D1, D1)) + ρᵣ = ROCArray(randn(T, D3, D3)) + H = ROCArray(randn(T, d1, d2, d1, d2)) + + @tensor begin + HRAA1[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + + buffer = AMDBufferAllocator() + @tensor allocator = buffer begin + HRAA2[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + @test HRAA2 isa ROCArray{T, 4} + @test collect(HRAA2) ≈ collect(HRAA1) + + # all temporaries were reclaimed, and the buffer was actually used + @test buffer.offset == 0 + @test buffer.max_offset > 0 + + # The high-water mark only counts the temporaries that actually fit in the + # buffer, so it may still grow while the buffer is warming up, but it has to + # converge to a fixed size after a couple of contractions. + max0 = buffer.max_offset + for _ in 1:5 + @tensor allocator = buffer begin + HRAA3[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + @test collect(HRAA3) ≈ collect(HRAA1) + end + max1 = buffer.max_offset + @test max1 ≥ max0 + @test length(buffer) ≥ max1 + + for _ in 1:5 + @tensor allocator = buffer begin + HRAA3[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + @test collect(HRAA3) ≈ collect(HRAA1) + end + @test buffer.offset == 0 + @test buffer.max_offset == max1 + + # scalar output + @tensor begin + E1 = ρₗ[a', a] * A1[a, s, b] * A2[b, s', c] * ρᵣ[c, c'] * + H[t, t', s, s'] * conj(A1[a', t, b']) * conj(A2[b', t', c']) + end + @tensor allocator = buffer begin + E2 = ρₗ[a', a] * A1[a, s, b] * A2[b, s', c] * ρᵣ[c, c'] * + H[t, t', s, s'] * conj(A1[a', t, b']) * conj(A2[b', t', c']) + end + @test E1 ≈ E2 + @test buffer.offset == 0 + end + + @testset "host inputs are promoted to ROCArray" begin + D1, D2, D3 = 30, 40, 20 + d1, d2 = 2, 3 + T = Float64 + + A1 = randn(T, D1, d1, D2) + A2 = randn(T, D2, d2, D3) + ρₗ = randn(T, D1, D1) + ρᵣ = randn(T, D3, D3) + H = randn(T, d1, d2, d1, d2) + + @tensor begin + HRAA1[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + + buffer = AMDBufferAllocator() + @tensor backend = hipTENSORBackend() allocator = buffer begin + HRAA2[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + @test HRAA2 isa ROCArray{T, 4} + @test collect(HRAA2) ≈ HRAA1 + @test buffer.offset == 0 + @test buffer.max_offset > 0 + end + + @testset "ncon" begin + A = ROCArray(randn(Float32, 5, 5)) + B = ROCArray(randn(Float32, 5, 5)) + C = ROCArray(randn(Float32, 5, 5)) + buffer = AMDBufferAllocator() + + R = ncon([A, B, C], [[-1, 1], [1, 2], [2, -2]]; allocator = buffer) + @test R isa ROCArray{Float32, 2} + @test collect(R) ≈ collect(A) * collect(B) * collect(C) + @test buffer.offset == 0 + + max1 = buffer.max_offset + for _ in 1:5 + ncon([A, B, C], [[-1, 1], [1, 2], [2, -2]]; allocator = buffer) + end + @test buffer.offset == 0 + @test buffer.max_offset == max1 + end + end + + @testset "@hiptensor" verbose = true begin + @testset "tensorcontract 1" begin + A = randn(Float64, (3, 5, 4, 6)) + # Float64 permutations are not supported yet + #=@tensor C1[4, 1, 3, 2] := A[1, 2, 3, 4] + @hiptensor C2[4, 1, 3, 2] := A[1, 2, 3, 4] + @test C1 ≈ collect(C2)=# + @test_throws IndexError begin + @hiptensor C[1, 2, 3, 4] := A[1, 2, 3] + end + @test_throws IndexError begin + @hiptensor C[1, 2, 3, 4] := A[1, 2, 2, 4] + end + + B = randn(Float64, (5, 6, 3, 4)) + p = [3, 1, 4, 2] + # Float64 permutations are not supported yet + #=@tensor C1[3, 1, 4, 2] := A[3, 1, 4, 2] + B[1, 2, 3, 4] + @hiptensor C2[3, 1, 4, 2] := A[3, 1, 4, 2] + B[1, 2, 3, 4] + @test C1 ≈ collect(C2)=# + @test_throws DimensionMismatch begin + @hiptensor C[1, 2, 3, 4] := A[1, 2, 3, 4] + B[1, 2, 3, 4] + end + + A = randn(Float64, (50, 100, 100)) + @tensor C1[a] := A[a, b', b'] + @hiptensor C2[a] := A[a, b', b'] + @test C1 ≈ collect(C2) + + A = randn(Float64, (3, 20, 5, 3, 20, 4, 5)) + @tensor C1[e, a, d] := A[a, b, c, d, b, e, c] + @hiptensor C2[e, a, d] := A[a, b, c, d, b, e, c] + @test C1 ≈ collect(C2) + + A = randn(Float64, (3, 20, 5, 3, 4)) + B = randn(Float64, (5, 6, 20, 3)) + @tensor C1[a, g, e, d, f] := A[a, b, c, d, e] * B[c, f, b, g] + @hiptensor C2[a, g, e, d, f] := A[a, b, c, d, e] * B[c, f, b, g] + @test C1 ≈ collect(C2) + @test_throws IndexError begin + @hiptensor A[a, b, c, d] * B[c, f, b, g] + end + end + + @testset "tensorcontract 2" begin + A = randn(Float64, (5, 5, 5, 5)) + B = rand(ComplexF64, (5, 5, 5, 5)) + # mixed real/complex contractions are not supported yet + #=@tensor C1[1, 2, 5, 6, 3, 4, 7, 8] := A[1, 2, 3, 4] * B[5, 6, 7, 8] + @hiptensor C2[1, 2, 5, 6, 3, 4, 7, 8] := A[1, 2, 3, 4] * B[5, 6, 7, 8] + @test C1 ≈ collect(C2)=# + @test_throws IndexError begin + @hiptensor C[a, b, c, d, e, f, g, i] := A[a, b, c, d] * B[e, f, g, h] + end + end + + # complex traces are not supported yet + #=@testset "tensorcontract 3" begin + Da, Db, Dc, Dd, De, Df, Dg, Dh = 10, 15, 4, 8, 6, 7, 3, 2 + A = rand(ComplexF64, (Da, Dc, Df, Da, De, Db, Db, Dg)) + B = rand(ComplexF64, (Dc, Dh, Dg, De, Dd)) + C = rand(ComplexF64, (Dd, Dh, Df)) + @tensor D1[d, f, h] := A[a, c, f, a, e, b, b, g] * B[c, h, g, e, d] + + 0.5 * C[d, h, f] + @hiptensor D2[d, f, h] := A[a, c, f, a, e, b, b, g] * B[c, h, g, e, d] + + 0.5 * C[d, h, f] + @test D1 ≈ collect(D2) + E1 = sqrt(abs((@tensor tensorscalar(D1[d, f, h] * conj(D1[d, f, h]))))) + E2 = sqrt(abs((@hiptensor tensorscalar(D2[d, f, h] * conj(D2[d, f, h]))))) + @test E1 ≈ E2 + end=# + + @testset "views" begin + p = [3, 1, 4, 2] + Abig = AMDGPU.randn(Float32, (30, 30, 30, 30)) + A = view( + Abig, 1 .+ 3 .* (0:9), 2 .+ 2 .* (0:6), 5 .+ 4 .* (0:6), 4 .+ 3 .* (0:8) + ) + Cbig = AMDGPU.zeros(Float32, (50, 50, 50, 50)) + C = view(Cbig, 13 .+ (0:6), 11 .+ 4 .* (0:9), 15 .+ 4 .* (0:8), 4 .+ 3 .* (0:6)) + Acopy = copy(A) + Ccopy = copy(C) + @tensor C[3, 1, 4, 2] = A[1, 2, 3, 4] + @tensor Ccopy[3, 1, 4, 2] = Acopy[1, 2, 3, 4] + @test copy(C) ≈ Ccopy + @test_throws TensorOperations.IndexError begin + @tensor C[3, 1, 4, 2] = A[1, 2, 3] + end + @test_throws DimensionMismatch begin + @tensor C[3, 1, 4, 2] = A[3, 1, 4, 2] + end + @test_throws TensorOperations.IndexError begin + @tensor C[1, 1, 2, 3] = A[1, 2, 3, 4] + end + end + + @testset "views 2" begin + p = [3, 1, 4, 2] + Abig = AMDGPU.randn(Float32, (30, 30, 30, 30)) + A = view( + Abig, 1 .+ 3 .* (0:9), 2 .+ 2 .* (0:6), 5 .+ 4 .* (0:6), 4 .+ 3 .* (0:8) + ) + Cbig = AMDGPU.zeros(Float32, (50, 50, 50, 50)) + C = view(Cbig, 13 .+ (0:6), 11 .+ 4 .* (0:9), 15 .+ 4 .* (0:8), 4 .+ 3 .* (0:6)) + Acopy = permutedims(copy(A), p) + Ccopy = copy(C) + α = randn(Float32) + β = randn(Float32) + @tensor C[3, 1, 4, 2] = β * C[3, 1, 4, 2] + α * A[1, 2, 3, 4] + Ccopy = β * Ccopy + α * Acopy + @test copy(C) ≈ Ccopy + @test_throws ArgumentError @macroexpand( + @tensor C[3, 1, 4, 2] = 0.5 * C[3, 1, 4, 2] + 1.2 * A[1, 2, 3] + ) + @test_throws DimensionMismatch begin + @tensor C[3, 1, 4, 2] = 0.5 * C[3, 1, 4, 2] + 1.2 * A[3, 1, 2, 4] + end + @test_throws ArgumentError @macroexpand( + @tensor C[1, 1, 2, 3] = 0.5 * C[1, 1, 2, 3] + 1.2 * A[1, 2, 3, 4] + ) + end + + @testset "views 3" begin + Abig = AMDGPU.rand(Float32, (30, 30, 30, 30)) + A = view( + Abig, 1 .+ 3 .* (0:8), 2 .+ 2 .* (0:14), 5 .+ 4 .* (0:6), 7 .+ 2 .* (0:8) + ) + Bbig = AMDGPU.rand(Float32, (50, 50)) + B = view(Bbig, 13 .+ (0:14), 3 .+ 5 .* (0:6)) + Acopy = copy(A) + Bcopy = copy(B) + α = randn(Float32) + @tensor B[b, c] += α * A[a, b, c, a] + @tensor Bcopy[b, c] += α * Acopy[a, b, c, a] + @test copy(B) ≈ Bcopy + @test_throws IndexError begin + @tensor B[b, c] += α * A[a, b, c] + end + @test_throws DimensionMismatch begin + @tensor B[c, b] += α * A[a, b, c, a] + end + @test_throws ArgumentError @macroexpand(@tensor B[c, b] += α * A[a, b, a, a]) + @test_throws DimensionMismatch begin + @tensor B[c, b] += α * A[a, b, a, c] + end + end + + @testset "views 4" begin + Abig = AMDGPU.rand(Float32, (30, 30, 30, 30)) + A = view( + Abig, 1 .+ 3 .* (0:8), 2 .+ 2 .* (0:14), 5 .+ 4 .* (0:6), + 7 .+ 2 .* (0:8) + ) + Bbig = AMDGPU.rand(Float32, (50, 50, 50)) + B = view(Bbig, 3 .+ 5 .* (0:6), 7 .+ 2 .* (0:7), 13 .+ (0:14)) + Cbig = AMDGPU.rand(Float32, (40, 40, 40)) + C = view(Cbig, 3 .+ 2 .* (0:8), 13 .+ (0:8), 7 .+ 3 .* (0:7)) + Acopy = copy(A) + Bcopy = copy(B) + Ccopy = copy(C) + α = randn(Float32) + @tensor C[d, a, e] -= α * A[a, b, c, d] * conj(B[c, e, b]) + @tensor Ccopy[d, a, e] -= α * Acopy[a, b, c, d] * conj(Bcopy[c, e, b]) + @test copy(C) ≈ Ccopy + end + end + + @testset "tensortrace! does not leak plans" begin + # each leaked plan holds on to a 128 KiB reduction workspace + A = ROCArray(randn(Float32, 64, 64, 64, 64)) + C = AMDGPU.zeros(Float32, 64, 64) + @tensor C[a, b] = A[a, c, b, c] # warm up + + GC.gc(true) + AMDGPU.reclaim() + AMDGPU.synchronize() + GC.enable(false) + try + live0 = AMDGPU.memory_stats().live + for _ in 1:1000 + @tensor C[a, b] = A[a, c, b, c] + end + AMDGPU.synchronize() + @test AMDGPU.memory_stats().live - live0 < 2^20 # would be 125 MiB if leaking + finally + GC.enable(true) + end + end + + @testset "Issues" verbose = true begin + @testset "Issue PR #186" begin + # https://github.com/Jutho/TensorOperations.jl/pull/186 + A = randn(Float32, (5, 5, 5, 5)) + Atr = @tensor A[a, b, b, a] + Atr2 = @hiptensor A[a, b, b, a] + @test Atr ≈ Atr2 + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 10f2728f..e7b1b490 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -70,6 +70,9 @@ end @testset "cuTENSOR extension" verbose = true begin include("cutensor.jl") end +@testset "hipTENSOR extension" verbose = true begin + include("hiptensor.jl") +end @testset "GPUArrays" verbose = true begin include("gpu.jl") end From 8ec4c5486a227ce2624cb284120f05f1964e659c Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 21 Oct 2026 16:00:47 +0200 Subject: [PATCH 2/3] Doc updates --- docs/src/index.md | 4 ++++ docs/src/man/indexnotation.md | 19 ++++++++++--------- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index ddb2a6db..507f56a5 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -60,6 +60,10 @@ The TensorOperations.jl package is centered around the following features: [CUDA.jl and cuTENSOR.jl](https://github.com/JuliaGPU/CUDA.jl), by relying on (and thus providing a high level interface into) NVidia's [cuTENSOR](https://developer.nvidia.com/cutensor) library. + - Support for `ROCArray` objects if used together with + [AMDGPU.jl and hipTENSOR](https://github.com/JuliaGPU/AMDGPU.jl), by relying on (and thus + providing a high level interface into) AMD's + [hipTENSOR](https://rocm.docs.amd.com/projects/hipTensor/en/latest/index.html) library. - Fall-back support for `AbstractArray` instances which are not strided, by using a general implementation using only constructions from Julia's `Base` and `LinearAlgebra` modules (`reshape`, `permutedims`, `mul!`, etc.) and some indexing in the case of diff --git a/docs/src/man/indexnotation.md b/docs/src/man/indexnotation.md index 9d7171e0..e3942204 100644 --- a/docs/src/man/indexnotation.md +++ b/docs/src/man/indexnotation.md @@ -471,16 +471,17 @@ simple wrapper over the cuTENSOR library of NVidia, and will only be loaded when `cuTENSOR.jl` package is loaded. The `@tensor` macro will then automatically work for operations between GPU arrays. -Mixed operations between host arrays (e.g. `Array`) and device arrays (e.g. `CuArray`) will -fail. However, if one wants to harness the computing power of the GPU to perform all tensor -operations, there is a dedicated macro `@cutensor`. This will transfer all host arrays to -the GPU before performing the requested operations. If the output is an existing host array, -the result will be copied back. If a new result array is created (i.e. using `:=`), it will -remain on the GPU device and it is up to the user to transfer it back. Arrays are transfered -to the GPU just before they are first used, and in a complicated tensor expression, this -might have the benefit that transer of the later arrays overlaps with computation of earlier -operations. +Mixed operations between host arrays (e.g. `Array`) and device arrays (e.g. `CuArray` or `ROCArray`) +will fail. However, if one wants to harness the computing power of the GPU to perform all tensor +operations, there is a dedicated macro `@cutensor` for CUDA devices, or `@hiptensor` for AMD. +This will transfer all host arrays to the GPU before performing the requested operations. +If the output is an existing host array, the result will be copied back. +If a new result array is created (i.e. using `:=`), it will remain on the GPU device and it +is up to the user to transfer it back. Arrays are transfered to the GPU just before they are +first used, and in a complicated tensor expression, this might have the benefit that transer +of the later arrays overlaps with computation of earlier operations. ```@docs @cutensor +@hiptensor ``` From 95b02b3b12faed0ee0b43dace9edcf64544b04b6 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Mon, 26 Oct 2026 13:30:43 +0100 Subject: [PATCH 3/3] Fall back to Strided if hipTensor won't work --- ext/TensorOperationsAMDGPUExt.jl | 54 ++++++++++++++++++------ test/hiptensor.jl | 71 ++++++++++++++++++++++---------- 2 files changed, 92 insertions(+), 33 deletions(-) diff --git a/ext/TensorOperationsAMDGPUExt.jl b/ext/TensorOperationsAMDGPUExt.jl index 70463e22..d1886931 100644 --- a/ext/TensorOperationsAMDGPUExt.jl +++ b/ext/TensorOperationsAMDGPUExt.jl @@ -5,13 +5,12 @@ using AMDGPU.hipTENSOR using TensorOperations using TensorOperations: TensorOperations as TO using TensorOperations: hipTENSORBackend, AMDAllocator, DefaultAllocator +using TensorOperations: StridedNative using TensorOperations: isstrided using Strided using TupleTools: TupleTools as TT -hiptensor_working = AMDGPU.hipTENSOR.has_hiptensor() - using AMDGPU.hipTENSOR: OP_IDENTITY, OP_CONJ, OP_ADD using AMDGPU.hipTENSOR: is_unary, is_binary using AMDGPU.hipTENSOR: hiptensorWorksizePreference_t, hiptensorAlgo_t, hiptensorOperationDescriptor_t, @@ -22,6 +21,8 @@ using AMDGPU.hipTENSOR: hiptensorCreatePlanPreference, hiptensorPlan, hipTensorP using AMDGPU.hipTENSOR: elementwise_binary_execute!, permute!, contract!, reduce! using AMDGPU.hipTENSOR: reduction_compute_types, hiptensorCreateReduction +using AMDGPU.hipTENSOR: permutation_compute_types, elementwise_binary_compute_types, + contraction_compute_types using AMDGPU.Adapt: adapt #------------------------------------------------------------------------------------------- @@ -140,20 +141,38 @@ end #------------------------------------------------------------------------------------------- # Backend selection and passing #------------------------------------------------------------------------------------------- -# A Base wrapper over `ROCArray` 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 `ROCStridedView` here. TODO handle the case where hipTensor isn't -# functional. -function TO.select_backend(::typeof(TO.tensoradd!), ::ROCStridedView, ::ROCStridedView) - return hipTENSORBackend() +# hipTENSOR only implements a handful of element type combinations for each operation (in +# particular, no complex permutations, elementwise operations or reductions), and might not +# be functional at all on the current device. Everything else is handled by the Strided +# implementation, which also runs on `ROCArray`s. Support for double precision types is also +# pretty spotty so that may fall back as well. +hiptensor_supports(table, eltypes) = has_hiptensor() && haskey(table, eltypes) + +# `tensoradd!` maps onto a permutation if `β` is zero, and onto an elementwise binary +# operation otherwise +supports_add(C, A) = supports_add(C, A, false) || supports_add(C, A, true) +supports_add(C, A, β) = hiptensor_supports( + iszero(β) ? permutation_compute_types : elementwise_binary_compute_types, + (eltype(A), eltype(C)) +) +supports_trace(C, A) = hiptensor_supports(reduction_compute_types, (eltype(A), eltype(C))) +supports_contract(C, A, B) = + hiptensor_supports(contraction_compute_types, (eltype(A), eltype(B), eltype(C))) + +strided_contract_backend(C, A, B) = @invoke TO.select_backend( + TO.tensorcontract!::typeof(TO.tensorcontract!), C::StridedView, A::StridedView, B::StridedView +) + +function TO.select_backend(::typeof(TO.tensoradd!), C::ROCStridedView, A::ROCStridedView) + return supports_add(C, A) ? hipTENSORBackend() : StridedNative() end -function TO.select_backend(::typeof(TO.tensortrace!), ::ROCStridedView, ::ROCStridedView) - return hipTENSORBackend() +function TO.select_backend(::typeof(TO.tensortrace!), C::ROCStridedView, A::ROCStridedView) + return supports_trace(C, A) ? hipTENSORBackend() : StridedNative() end function TO.select_backend( - ::typeof(TO.tensorcontract!), ::ROCStridedView, ::ROCStridedView, ::ROCStridedView + ::typeof(TO.tensorcontract!), C::ROCStridedView, A::ROCStridedView, B::ROCStridedView ) - return hipTENSORBackend() + return supports_contract(C, A, B) ? hipTENSORBackend() : strided_contract_backend(C, A, B) end # Make sure that if the `hipTensorBackend` is specified, arrays are converted to ROCArrays @@ -249,6 +268,9 @@ function TO.tensoradd!( α::Number, β::Number, backend::hipTENSORBackend, allocator ) + supports_add(C, A, β) || + return TO.tensoradd!(C, A, pA, conjA, α, β, StridedNative(), allocator) + # hipTENSOR does not validate sizes itself TO.argcheck_tensoradd(C, A, TO.linearize(pA)) TO.dimcheck_tensoradd(C, A, TO.linearize(pA)) @@ -291,6 +313,11 @@ function TO.tensorcontract!( α::Number, β::Number, backend::hipTENSORBackend, allocator ) + supports_contract(C, A, B) || return TO.tensorcontract!( + C, A, pA, conjA, B, pB, conjB, pAB, α, β, + strided_contract_backend(C, A, B), allocator + ) + # hipTENSOR does not validate sizes itself TO.argcheck_tensorcontract(C, A, pA, B, pB, pAB) TO.dimcheck_tensorcontract(C, A, pA, B, pB, pAB) @@ -323,6 +350,9 @@ function TO.tensortrace!( α::Number, β::Number, backend::hipTENSORBackend, allocator ) + supports_trace(C, A) || + return TO.tensortrace!(C, A, p, q, conjA, α, β, StridedNative(), allocator) + # hipTENSOR does not validate sizes itself TO.argcheck_tensortrace(C, A, TO.linearize(p), q) TO.dimcheck_tensortrace(C, A, TO.linearize(p), q) diff --git a/test/hiptensor.jl b/test/hiptensor.jl index 35925bfd..d0528bbf 100644 --- a/test/hiptensor.jl +++ b/test/hiptensor.jl @@ -32,14 +32,13 @@ if hipTENSOR.has_hiptensor() @tensor C2[a, b, c, d] := ROCArray(A)[a, b, c, d] + α * ROCArray(B)[a, b, c, d] @test collect(C2) ≈ C1 - # complex doens't work yet - #=C = randn(ComplexF32, (5, 6, 3, 4)) + C = randn(ComplexF32, (5, 6, 3, 4)) D = randn(ComplexF32, (5, 3, 4, 6)) β = randn(ComplexF32) @tensor E1[a, b, c, d] := C[a, b, c, d] + β * conj(D[a, c, d, b]) @tensor E2[a, b, c, d] := ROCArray(C)[a, b, c, d] + β * conj(ROCArray(D)[a, c, d, b]) - @test collect(E2) ≈ E1=# + @test collect(E2) ≈ E1 end @testset "tensortrace" begin @@ -48,8 +47,7 @@ if hipTENSOR.has_hiptensor() @tensor B2[a] := ROCArray(A)[a, b′, b′] @test collect(B2) ≈ B1 - # not supported yet - #=C = randn(ComplexF32, (3, 20, 5, 3, 20, 4, 5)) + C = randn(ComplexF32, (3, 20, 5, 3, 20, 4, 5)) @tensor D1[e, a, d] := C[a, b, c, d, b, e, c] @tensor D2[e, a, d] := ROCArray(C)[a, b, c, d, b, e, c] @test collect(D2) ≈ D1 @@ -61,7 +59,7 @@ if hipTENSOR.has_hiptensor() α = randn(ComplexF32) @tensor D5[d, e, a] := α * C[a, b, c, d, b, e, c] @tensor D6[d, e, a] := α * ROCArray(C)[a, b, c, d, b, e, c] - @test collect(D6) ≈ D5=# + @test collect(D6) ≈ D5 end @testset "tensorcontract" begin @@ -76,12 +74,11 @@ if hipTENSOR.has_hiptensor() @tensor F1[a, b, c, d, e, f] := D[a, b, c] * conj(E[d, e, f]) @tensor F2[a, b, c, d, e, f] := ROCArray(D)[a, b, c] * conj(ROCArray(E)[d, e, f]) @test collect(F2) ≈ F1 - # not yet supported - #=D = randn(Float64, (5, 5, 5)) + D = randn(Float64, (5, 5, 5)) E = rand(ComplexF64, (5, 5, 5)) @tensor F1[a, b, c, d, e, f] := D[a, b, c] * conj(E[d, e, f]) @tensor F2[a, b, c, d, e, f] := ROCArray(D)[a, b, c] * conj(ROCArray(E)[d, e, f]) - @test collect(F2) ≈ F1=# + @test collect(F2) ≈ F1 end end @@ -389,10 +386,9 @@ if hipTENSOR.has_hiptensor() @testset "@hiptensor" verbose = true begin @testset "tensorcontract 1" begin A = randn(Float64, (3, 5, 4, 6)) - # Float64 permutations are not supported yet - #=@tensor C1[4, 1, 3, 2] := A[1, 2, 3, 4] + @tensor C1[4, 1, 3, 2] := A[1, 2, 3, 4] @hiptensor C2[4, 1, 3, 2] := A[1, 2, 3, 4] - @test C1 ≈ collect(C2)=# + @test C1 ≈ collect(C2) @test_throws IndexError begin @hiptensor C[1, 2, 3, 4] := A[1, 2, 3] end @@ -402,10 +398,9 @@ if hipTENSOR.has_hiptensor() B = randn(Float64, (5, 6, 3, 4)) p = [3, 1, 4, 2] - # Float64 permutations are not supported yet - #=@tensor C1[3, 1, 4, 2] := A[3, 1, 4, 2] + B[1, 2, 3, 4] + @tensor C1[3, 1, 4, 2] := A[3, 1, 4, 2] + B[1, 2, 3, 4] @hiptensor C2[3, 1, 4, 2] := A[3, 1, 4, 2] + B[1, 2, 3, 4] - @test C1 ≈ collect(C2)=# + @test C1 ≈ collect(C2) @test_throws DimensionMismatch begin @hiptensor C[1, 2, 3, 4] := A[1, 2, 3, 4] + B[1, 2, 3, 4] end @@ -433,17 +428,15 @@ if hipTENSOR.has_hiptensor() @testset "tensorcontract 2" begin A = randn(Float64, (5, 5, 5, 5)) B = rand(ComplexF64, (5, 5, 5, 5)) - # mixed real/complex contractions are not supported yet - #=@tensor C1[1, 2, 5, 6, 3, 4, 7, 8] := A[1, 2, 3, 4] * B[5, 6, 7, 8] + @tensor C1[1, 2, 5, 6, 3, 4, 7, 8] := A[1, 2, 3, 4] * B[5, 6, 7, 8] @hiptensor C2[1, 2, 5, 6, 3, 4, 7, 8] := A[1, 2, 3, 4] * B[5, 6, 7, 8] - @test C1 ≈ collect(C2)=# + @test C1 ≈ collect(C2) @test_throws IndexError begin @hiptensor C[a, b, c, d, e, f, g, i] := A[a, b, c, d] * B[e, f, g, h] end end - # complex traces are not supported yet - #=@testset "tensorcontract 3" begin + @testset "tensorcontract 3" begin Da, Db, Dc, Dd, De, Df, Dg, Dh = 10, 15, 4, 8, 6, 7, 3, 2 A = rand(ComplexF64, (Da, Dc, Df, Da, De, Db, Db, Dg)) B = rand(ComplexF64, (Dc, Dh, Dg, De, Dd)) @@ -456,7 +449,7 @@ if hipTENSOR.has_hiptensor() E1 = sqrt(abs((@tensor tensorscalar(D1[d, f, h] * conj(D1[d, f, h]))))) E2 = sqrt(abs((@hiptensor tensorscalar(D2[d, f, h] * conj(D2[d, f, h]))))) @test E1 ≈ E2 - end=# + end @testset "views" begin p = [3, 1, 4, 2] @@ -553,6 +546,42 @@ if hipTENSOR.has_hiptensor() end end + @testset "unsupported types fall back to Strided" begin + using TensorOperations: select_backend, StridedNative, StridedBLAS + using Strided: StridedView + sv(T, dims...) = StridedView(ROCArray(rand(T, dims...))) + + A32, C32 = sv(Float32, 4, 4), sv(Float32, 4, 4) + @test select_backend(tensoradd!, C32, A32) == hipTENSORBackend() + @test select_backend(tensortrace!, sv(Float32, 4), sv(Float32, 4, 4, 4)) == + hipTENSORBackend() + @test select_backend(tensorcontract!, C32, A32, A32) == hipTENSORBackend() + + Az, Cz = sv(ComplexF32, 4, 4), sv(ComplexF32, 4, 4) + @test select_backend(tensoradd!, Cz, Az) == StridedNative() + @test select_backend(tensortrace!, sv(ComplexF32, 4), sv(ComplexF32, 4, 4, 4)) == + StridedNative() + @test select_backend(tensorcontract!, Cz, Az, Az) == hipTENSORBackend() + @test select_backend(tensorcontract!, Cz, A32, Az) == StridedBLAS() + Ai = sv(Int, 4, 4) + @test select_backend(tensorcontract!, Ai, Ai, Ai) == StridedNative() + + # results stay on the device and match the host + A = randn(ComplexF64, (3, 4, 5)) + @tensor C1[c, a, b] := 2 * conj(A[a, b, c]) + C2 = similar(ROCArray(A), (5, 3, 4)) + @tensor backend = hipTENSORBackend() C2[c, a, b] = 2 * conj(ROCArray(A)[a, b, c]) + @test C2 isa ROCArray + @test collect(C2) ≈ C1 + + A = rand(-5:5, (3, 4)) + B = rand(-5:5, (4, 5)) + @tensor C1[a, c] := A[a, b] * B[b, c] + @tensor C2[a, c] := ROCArray(A)[a, b] * ROCArray(B)[b, c] + @test C2 isa ROCArray + @test collect(C2) == C1 + end + @testset "tensortrace! does not leak plans" begin # each leaked plan holds on to a 128 KiB reduction workspace A = ROCArray(randn(Float32, 64, 64, 64, 64))