[WIP] Add a lazy ConjArray absorbed by permutation, matricize and contraction

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.23.3"
+version = "0.23.4"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl

@@ -31,20 +31,16 @@ function TA.contractpermopadd!(
         α::Number, β::Number
     )
     TA.check_input(
-        TA.contract!,
+        TA.contractpermopadd!,
         a_dest, perm_dest_codomain, perm_dest_domain,
-        a1, perm1_codomain, perm1_domain,
-        a2, perm2_codomain, perm2_domain
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain
     )
     permblocks1 = Tuple.((perm1_codomain, perm1_domain))
     permblocks2 = Tuple.((perm2_codomain, perm2_domain))
     permblocks_dest = Tuple.((perm_dest_codomain, perm_dest_domain))
-    conj1 = op1 === conj
-    conj2 = op2 === conj
-    a1′ = (op1 === identity || op1 === conj) ? a1 : op1.(a1)
-    a2′ = (op2 === identity || op2 === conj) ? a2 : op2.(a2)
     return TO.tensorcontract!(
-        a_dest, a1′, permblocks1, conj1, a2′, permblocks2, conj2,
+        a_dest, a1, permblocks1, op1 === conj, a2, permblocks2, op2 === conj,
         permblocks_dest, α, β, backend(algorithm), allocator(algorithm)
     )
 end

src/TensorAlgebra.jl

@@ -5,7 +5,7 @@ export contract, contract!, contractalign, dual, isdual, MatrixAlgebra
 if VERSION >= v"1.11.0-DEV.469"
     eval(
         Meta.parse(
-            "public AbstractAlgorithm, add!, AddBroadcasted, addends, allocate_output, allocate_project, arguments, axes, bipartition, bipartition_axes, biperm, bipermutedims, bipermutedims!, bipermutedimsopadd!, cat_axis, cat_similar, check_input, concatenate, concatenate!, ConjBroadcasted, contractadd!, ContractAlgorithm, contractopadd!, contractperm, contractperm!, contractpermadd!, contractpermalign, contractpermopadd!, data, datatype, default_algorithm, dims2cat, directsum, dot, dotperm, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc, eigh_vals, fill_map, flattenlinear, has_bipartition, infer_aux_space, invsqrth_safe, is_output_view, is_projected, isidentitybiperm, label_type, left_null, left_orth, left_polar, LinearBroadcasted, linearbroadcasted, lq_compact, lq_full, matricize, MatricizeContract, matricizeop, matricizeop!, matricizeopcopy, matricizeopview, MATRIX_FUNCTIONS, ndims, ndims_codomain, ndims_domain, one, ones_map, operation, output_axes, PermutedDims, permuteddims, permutedims, permutedims!, permutedimsadd!, permutedimsop, permutedimsopadd!, project, project!, project_aux, project_hermitian, projectto!, qr_compact, qr_full, rand_map, randn_map, right_null, right_orth, right_polar, scalar, scale!, scaleadd!, ScaledBroadcasted, select_algorithm, similar_map, size, sqrth_invsqrth_safe, sqrth_safe, sum, svd_compact, svd_full, svd_trunc, svd_vals, TensorOperationsContract, to_range, tr, trivialrange, tryflattenlinear, tryproject, tryproject_aux, twist!, unchecked_project, unchecked_project_aux, ungrade, unmatricize, unmatricize!, unmatricize_factors, unmatricizeadd!, unproject, unscaled, zero!, zeros_map"
+            "public AbstractAlgorithm, add!, AddBroadcasted, addends, allocate_output, allocate_project, arguments, axes, bipartition, bipartition_axes, biperm, bipermutedims, bipermutedims!, bipermutedimsopadd!, cat_axis, cat_similar, check_input, concatenate, concatenate!, ConjArray, ConjBroadcasted, conjed, contractadd!, ContractAlgorithm, contractopadd!, contractperm, contractperm!, contractpermadd!, contractpermalign, contractpermop, contractpermopadd!, data, datatype, default_algorithm, dims2cat, directsum, dot, dotperm, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc, eigh_vals, fill_map, flattenlinear, has_bipartition, infer_aux_space, invsqrth_safe, is_output_view, is_projected, isidentitybiperm, label_type, left_null, left_orth, left_polar, LinearBroadcasted, linearbroadcasted, lq_compact, lq_full, matricize, MatricizeContract, matricizeop, matricizeop!, matricizeopcopy, matricizeopview, MATRIX_FUNCTIONS, ndims, ndims_codomain, ndims_domain, one, ones_map, operation, output_axes, PermutedDims, permuteddims, permutedims, permutedims!, permutedimsadd!, permutedimsop, permutedimsopadd!, project, project!, project_aux, project_hermitian, projectto!, qr_compact, qr_full, rand_map, randn_map, right_null, right_orth, right_polar, scalar, scale!, scaleadd!, ScaledBroadcasted, select_algorithm, similar_map, size, sqrth_invsqrth_safe, sqrth_safe, sum, svd_compact, svd_full, svd_trunc, svd_vals, TensorOperationsContract, to_range, tr, trivialrange, tryflattenlinear, tryproject, tryproject_aux, twist!, unchecked_project, unchecked_project_aux, ungrade, unmatricize, unmatricize!, unmatricize_factors, unmatricizeadd!, unproject, unscaled, zero!, zeros_map"
         )
     )
 end
@@ -35,5 +35,6 @@ include("matrixfunctions.jl")
 include("similar_map.jl")
 include("projectto.jl")
 include("linearbroadcasted.jl")
+include("conjarray.jl")
 
 end

src/conjarray.jl

@@ -0,0 +1,156 @@
+"""
+    ConjArray(a::AbstractArray)
+
+Lazy conjugate of `a`: an `AbstractArray` whose axes are the conjugated (dualized) parent axes
+and whose elements are the conjugated parent elements. Constructed by [`conjed`](@ref).
+Permutation, matricization and contraction absorb it by unwrapping to the parent and folding
+`conj` into their `op`, so it is never materialized on its way into those operations. Whole-array
+reads, writes and broadcasting go through the parent as well, as do reductions for a dense parent.
+Scalar indexing is disallowed, so an operation without such a path errors instead of silently
+looping over elements.
+"""
+struct ConjArray{T, N, P <: AbstractArray{T, N}} <: AbstractArray{T, N}
+    parent::P
+end
+
+"""
+    conjed(a::AbstractArray)
+
+Lazy conjugate of `a`, returning a [`ConjArray`](@ref). The lazy companion to the eager
+`conj(a)`, mirroring `adjoint`/`Adjoint`. `conjed(conjed(a))` returns `a`.
+"""
+conjed(a::AbstractArray) = ConjArray(a)
+conjed(a::ConjArray) = parent(a)
+
+# Most methods below either go through an op-aware kernel (`permutedimsopadd!`, `matricizeop`,
+# contraction), which owns what `conj` means for its array type, or use identities that hold for any
+# antiunitary involution (`norm`, `iszero`, `dot` of two lazy conjugates). `fill!`, `==`, `mapreduce`
+# and `dot` with one lazy operand instead assume `conj` acts elementwise on a flat memory layout, so
+# they are only defined for a dense parent, and other array types opt in.
+const ConjDenseArray{T, N} = ConjArray{T, N, <:DenseArray{T, N}}
+
+Base.parent(a::ConjArray) = a.parent
+
+# Conjugating an axis is identity for plain integer ranges and dualizes graded axes (where
+# `conj` is overloaded as `dual` downstream).
+Base.axes(a::ConjArray) = map(conj, axes(parent(a)))
+Base.size(a::ConjArray) = size(parent(a))
+
+function Base.getindex(a::ConjArray, I::Int...)
+    return throw(
+        ArgumentError(
+            "scalar indexing of a lazy `ConjArray` is disallowed, materialize it with `copy` first"
+        )
+    )
+end
+function Base.setindex!(a::ConjArray, value, I::Int...)
+    return throw(
+        ArgumentError(
+            "scalar indexing of a lazy `ConjArray` is disallowed, write the whole array with `copyto!`"
+        )
+    )
+end
+
+# Allocate like the parent, so a graded parent yields a graded array.
+function Base.similar(a::ConjArray, ::Type{T}, dims::Dims) where {T}
+    return similar(parent(a), T, dims)
+end
+
+# Materialize through the eager `conj.(a)` path, so a lazy and an eager conjugate agree
+# (including any parent-specific behavior such as a graded fermion sign).
+Base.copy(a::ConjArray) = copy(ConjBroadcasted(parent(a)))
+Base.Array{T, N}(a::ConjArray{<:Any, N}) where {T, N} = convert(Array{T, N}, copy(a))
+# Base `collect` iterates elements, which scalar indexing disallows.
+Base.collect(a::ConjArray) = Array(a)
+# Like `conj(::Adjoint)` in LinearAlgebra, this returns the parent rather than a copy.
+Base.conj(a::ConjArray) = parent(a)
+
+function Base.copyto!(dest::AbstractArray, src::ConjArray)
+    return permutedimsopadd!(
+        dest,
+        conj,
+        parent(src),
+        ntuple(identity, ndims(src)),
+        true,
+        false
+    )
+end
+# Write into the parent the conjugate of what the lazy conjugate should hold. `conj` is an
+# involution, fermion signs included, so this is exact.
+function Base.copyto!(dest::ConjArray, src::AbstractArray)
+    permutedimsopadd!(parent(dest), conj, src, ntuple(identity, ndims(src)), true, false)
+    return dest
+end
+# Also resolves the ambiguity between the two methods above.
+function Base.copyto!(dest::ConjArray, src::ConjArray)
+    copyto!(parent(dest), parent(src))
+    return dest
+end
+function Base.fill!(a::ConjDenseArray, value)
+    fill!(parent(a), conj(value))
+    return a
+end
+
+Base.:(==)(a::ConjArray, b::ConjArray) = parent(a) == parent(b)
+# A fused two-array reduction rather than a copy, so it maps to a single kernel on GPU.
+function Base.:(==)(a::ConjDenseArray, b::DenseArray)
+    axes(a) == axes(b) || return false
+    return mapreduce((x, y) -> conj(x) == y, &, parent(a), b; init = true)
+end
+Base.:(==)(a::DenseArray, b::ConjDenseArray) = b == a
+# Conjugation, and any sign it carries, preserves every p-norm. `isapprox` goes through `norm`.
+LinearAlgebra.norm(a::ConjArray, p::Real = 2) = LinearAlgebra.norm(parent(a), p)
+Base.iszero(a::ConjArray) = iszero(parent(a))
+# `conj` is antiunitary, fermion signs included, so a pair of lazy conjugates reduces to the parents.
+function LinearAlgebra.dot(a::ConjArray, b::ConjArray)
+    return conj(LinearAlgebra.dot(parent(a), parent(b)))
+end
+# Against a dense array, the lazy conjugate cancels the conjugation in `dot`, leaving the bilinear
+# pairing of the parent with `b` (BLAS `dotu`).
+LinearAlgebra.dot(a::ConjDenseArray, b::DenseArray) = transpose(vec(parent(a))) * vec(b)
+LinearAlgebra.dot(a::DenseArray, b::ConjDenseArray) = conj(LinearAlgebra.dot(b, a))
+function Base.mapreduce(f, op, a::ConjDenseArray; kwargs...)
+    return mapreduce(f ∘ conj, op, parent(a); kwargs...)
+end
+Base.permutedims(a::ConjArray, perm) = permutedims(a, perm)
+
+Base.show(io::IO, a::ConjArray) = print(io, "conjed(", parent(a), ")")
+function Base.show(io::IO, mime::MIME"text/plain", a::ConjArray)
+    print(io, "conjed of ")
+    return show(io, mime, copy(a))
+end
+
+# A broadcast over a lazy conjugate reads it as `conj.(parent)`, so it never indexes the wrapper
+# and a linear one lowers to a `ConjBroadcasted` leaf.
+Base.Broadcast.broadcastable(a::ConjArray) = Base.Broadcast.broadcasted(conj, parent(a))
+function Base.Broadcast.BroadcastStyle(::Type{<:ConjArray{<:Any, <:Any, P}}) where {P}
+    return Base.Broadcast.BroadcastStyle(P)
+end
+
+linearbroadcasted(::typeof(conj), a::ConjArray) = parent(a)
+
+# Each operation unwraps a `ConjArray` to its parent and folds `conj` into `op`
+# (`identity -> conj`, `conj -> identity`), before any backend dispatch.
+function bipermutedimsopadd!(
+        dest, op, src::ConjArray,
+        perm_codomain, perm_domain,
+        α::Number, β::Number
+    )
+    return bipermutedimsopadd!(
+        dest, _compose_op(op, conj), parent(src),
+        perm_codomain, perm_domain,
+        α, β
+    )
+end
+function permutedimsop(op, src::ConjArray, perm_codomain, perm_domain)
+    return permutedimsop(_compose_op(op, conj), parent(src), perm_codomain, perm_domain)
+end
+function matricizeop(op, a::ConjArray, perm_codomain, perm_domain)
+    return matricizeop(_compose_op(op, conj), parent(a), perm_codomain, perm_domain)
+end
+function matricizeopcopy(op, a::ConjArray, perm_codomain, perm_domain)
+    return matricizeopcopy(_compose_op(op, conj), parent(a), perm_codomain, perm_domain)
+end
+
+_unwrap_conj(op, a) = (op, a)
+_unwrap_conj(op, a::ConjArray) = (_compose_op(op, conj), parent(a))

src/contract/allocate_output.jl

@@ -6,33 +6,39 @@ function check_biperm(a, perm_codomain, perm_domain)
     return nothing
 end
 
-function check_input(
-        ::typeof(contract),
-        a1, perm1_codomain, perm1_domain,
-        a2, perm2_codomain, perm2_domain
-    )
-    check_biperm(a1, perm1_codomain, perm1_domain)
-    check_biperm(a2, perm2_codomain, perm2_domain)
-    # The contracted axes of `a1` (its domain) pair with the contracted axes of `a2` (its codomain),
-    # and each pair must be a dual pair: `dual(ax1) == ax2`. `dual` falls back to the identity for
-    # plain ranges, so this reduces to `ax1 == ax2` for non-graded arrays.
+# The contracted axes of `op1(a1)` (its domain) pair with the contracted axes of `op2(a2)` (its
+# codomain), and each pair must be a dual pair: `dual(ax1) == ax2`. `op` is `identity` or `conj`,
+# and `conj` dualizes a graded axis. `dual` and `conj` fall back to the identity for plain ranges,
+# so this reduces to `ax1 == ax2` for non-graded arrays.
+function check_contracted_axes(op1, a1, perm1_domain, op2, a2, perm2_codomain)
     length(perm1_domain) == length(perm2_codomain) || throw(
         ArgumentError(
             "Number of contracted axes do not match: `a1` has $(length(perm1_domain)), `a2` has $(length(perm2_codomain))"
         )
     )
     for (i, j) in zip(perm1_domain, perm2_codomain)
-        ax1 = axes(a1, i)
-        ax2 = axes(a2, j)
+        ax1 = op1(axes(a1, i))
+        ax2 = op2(axes(a2, j))
         dual(ax1) == ax2 || throw(
             ArgumentError(
-                "Contracted axes do not match: `axes(a1, $i) = $ax1` and `axes(a2, $j) = $ax2`"
+                "Contracted axes do not match: `$op1(axes(a1, $i)) = $ax1` and `$op2(axes(a2, $j)) = $ax2`"
             )
         )
     end
     return nothing
 end
 
+function check_input(
+        ::typeof(contract),
+        a1, perm1_codomain, perm1_domain,
+        a2, perm2_codomain, perm2_domain
+    )
+    check_biperm(a1, perm1_codomain, perm1_domain)
+    check_biperm(a2, perm2_codomain, perm2_domain)
+    check_contracted_axes(identity, a1, perm1_domain, identity, a2, perm2_codomain)
+    return nothing
+end
+
 function check_input(
         ::typeof(contract!),
         a_dest, perm_dest_codomain, perm_dest_domain,
@@ -53,16 +59,62 @@ function check_input(
     return nothing
 end
 
-# TODO: Use `ArrayLayouts`-like `MulAdd` object,
-# i.e. `ContractAdd`?
+function check_input(
+        ::typeof(contractpermop),
+        perm_dest_codomain, perm_dest_domain,
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain
+    )
+    for op in (op1, op2)
+        op === identity || op === conj ||
+            throw(ArgumentError("`op` must be `identity` or `conj`, got `$op`"))
+    end
+    check_biperm(a1, perm1_codomain, perm1_domain)
+    check_biperm(a2, perm2_codomain, perm2_domain)
+    check_contracted_axes(op1, a1, perm1_domain, op2, a2, perm2_codomain)
+    return nothing
+end
+
+function check_input(
+        ::typeof(contractpermopadd!),
+        a_dest, perm_dest_codomain, perm_dest_domain,
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain
+    )
+    # TODO: FIXME: Check that uncontracted axes match.
+    check_input(
+        contractpermop,
+        perm_dest_codomain, perm_dest_domain,
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain
+    )
+    check_biperm(a_dest, perm_dest_codomain, perm_dest_domain)
+    return nothing
+end
+
 function output_axes(
         ::typeof(contract),
         perm_dest_codomain, perm_dest_domain,
         a1, perm1_codomain, perm1_domain,
         a2, perm2_codomain, perm2_domain
     )
-    axes_codomain, _ = bipartition(axes(a1), perm1_codomain, perm1_domain)
-    _, axes_domain = bipartition(axes(a2), perm2_codomain, perm2_domain)
+    return output_axes(
+        contractpermop,
+        perm_dest_codomain, perm_dest_domain,
+        identity, a1, perm1_codomain, perm1_domain,
+        identity, a2, perm2_codomain, perm2_domain
+    )
+end
+# TODO: Use `ArrayLayouts`-like `MulAdd` object,
+# i.e. `ContractAdd`?
+function output_axes(
+        ::typeof(contractpermop),
+        perm_dest_codomain, perm_dest_domain,
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain
+    )
+    axes_codomain = map(i -> op1(axes(a1, i)), perm1_codomain)
+    axes_domain = map(i -> op2(axes(a2, i)), perm2_domain)
     axes_uncontracted = (axes_codomain..., axes_domain...)
     axes_codomain_dest, axes_domain_dest = bipartition(
         axes_uncontracted, perm_dest_codomain, perm_dest_domain
@@ -72,32 +124,42 @@ function output_axes(
     return axes_codomain_dest, conj.(axes_domain_dest)
 end
 
-# TODO: Use `ArrayLayouts`-like `MulAdd` object,
-# i.e. `ContractAdd`?
-# The destination `contract` writes into. A structured operand type overloads this directly, deriving
-# the axes and element type from `output_axes` and `Base.promote_op` as below; the permutations are
-# part of the signature because the contraction pattern is not recoverable from the destination leg
-# counts alone.
 function allocate_output(
         ::typeof(contract),
         perm_dest_codomain, perm_dest_domain,
         a1, perm1_codomain, perm1_domain,
         a2, perm2_codomain, perm2_domain
     )
+    return allocate_output(
+        contractpermop,
+        perm_dest_codomain, perm_dest_domain,
+        identity, a1, perm1_codomain, perm1_domain,
+        identity, a2, perm2_codomain, perm2_domain
+    )
+end
+# TODO: Use `ArrayLayouts`-like `MulAdd` object,
+# i.e. `ContractAdd`?
+# The destination `contractpermop` writes into. A structured operand type overloads this directly,
+# deriving the axes and element type from `output_axes` and `Base.promote_op` as below; the
+# permutations are part of the signature because the contraction pattern is not recoverable from the
+# destination leg counts alone.
+function allocate_output(
+        ::typeof(contractpermop),
+        perm_dest_codomain, perm_dest_domain,
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain
+    )
     check_input(
-        contract,
-        a1,
-        perm1_codomain,
-        perm1_domain,
-        a2,
-        perm2_codomain,
-        perm2_domain
+        contractpermop,
+        perm_dest_codomain, perm_dest_domain,
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain
     )
     axes_codomain_dest, axes_domain_dest = output_axes(
-        contract,
+        contractpermop,
         perm_dest_codomain, perm_dest_domain,
-        a1, perm1_codomain, perm1_domain,
-        a2, perm2_codomain, perm2_domain
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain
     )
     T = Base.promote_op(matprod, eltype(a1), eltype(a2))
     return zero!(similar_map(a1, T, axes_codomain_dest, axes_domain_dest))