Alternative Allocators

src/implementation/functions.jl

@@ -1,12 +1,31 @@
 # methods/simple.jl
 #
 # Method-based access to tensor operations using simple definitions.
+
+# utility function
+function _kwargs2args(; kwargs...)
+    for k in keys(kwargs)
+        if !(k in (:backend, :allocator))
+            throw(ArgumentError("unknown keyword argument: $k"))
+        end
+    end
+    if haskey(kwargs, :backend) && haskey(kwargs, :allocator)
+        return (kwargs[:backend], kwargs[:allocator])
+    elseif haskey(kwargs, :allocator)
+        return (DefaultBackend(), kwargs[:allocator])
+    elseif haskey(kwargs, :backend)
+        return (kwargs[:backend],)
+    else
+        return ()
+    end
+end
+
 # ------------------------------------------------------------------------------------------
 # tensorcopy
 # ------------------------------------------------------------------------------------------
 """
-    tensorcopy([IC=IA], A, IA, [conjA=false, [α=1]])
-    tensorcopy(A, pA::Index2Tuple, conjA, α, [backend]) # expert mode
+    tensorcopy([IC=IA], A, IA, [conjA=false, [α=1]]; [backend=..., allocator=...])
+    tensorcopy(A, pA::Index2Tuple, conjA, α, [backend, allocator]) # expert mode
 
 Create a copy of `A`, where the dimensions of `A` are assigned indices from the
 iterable `IA` and the indices of the copy are contained in `IC`. Both iterables
@@ -16,20 +35,23 @@ The result of this method is equivalent to `α * permutedims(A, pA)` where `pA`
 permutation such that `IC = IA[pA]`. The implementation of `tensorcopy` is however more
 efficient on average, especially if `Threads.nthreads() > 1`.
 
-Optionally, the flag `conjA` can be used to specify whether the input tensor should be
-conjugated (`true`) or not (`false`).
+The optional argument `conjA` can be used to specify whether the input tensor should be
+conjugated (`true`) or not (`false`), whereas `α` can be used to scale the result. It is
+also optional to specify a backend implementation to use, and an allocator to be used if
+temporary tensor objects are needed.
 
 See also [`tensorcopy!`](@ref).
 """
 function tensorcopy end
 
-function tensorcopy(IC::Labels, A, IA::Labels, conjA::Bool=false, α::Number=One())
+function tensorcopy(IC::Labels, A, IA::Labels, conjA::Bool=false, α::Number=One();
+                    kwargs...)
     pA = add_indices(IA, IC)
-    return tensorcopy(A, pA, conjA, α)
+    return tensorcopy(A, pA, conjA, α, _kwargs2args(; kwargs...)...)
 end
 # default `IC`
-function tensorcopy(A, IA::Labels, conjA::Bool=false, α::Number=One())
-    return tensorcopy(IA, A, IA, conjA, α)
+function tensorcopy(A, IA::Labels, conjA::Bool=false, α::Number=One(); kwargs...)
+    return tensorcopy(IA, A, IA, conjA, α; kwargs...)
 end
 # expert mode
 function tensorcopy(A, pA::Index2Tuple, conjA::Bool=false, α::Number=One(),
@@ -40,15 +62,15 @@ function tensorcopy(A, pA::Index2Tuple, conjA::Bool=false, α::Number=One(),
 end
 
 """
-    tensorcopy!(C, A, pA::Index2Tuple, conjA=false, α=1, [backend])
+    tensorcopy!(C, A, pA::Index2Tuple, conjA=false, α=1, [backend, allocator])
 
 Copy the contents of tensor `A` into `C`, where the dimensions `A` are permuted according to
 the permutation and repartition `pA`.
 
 The result of this method is equivalent to `α * permutedims!(C, A, pA)`.
 
 Optionally, the flag `conjA` can be used to specify whether the input tensor should be
-conjugated (`true`) or not (`false`).
+conjugated (`true`) or not (`false`). 
 
 !!! warning 
     The object `C` must not be aliased with `A`.
@@ -64,8 +86,8 @@ end
 # tensoradd
 # ------------------------------------------------------------------------------------------
 """
-    tensoradd([IC=IA], A, IA, [conjA], B, IB, [conjB], [α=1, [β=1]])
-    tensoradd(A, pA::Index2Tuple, conjA, B, pB::Index2Tuple, conjB, α=1, β=1, [backend]) # expert mode
+    tensoradd([IC=IA], A, IA, [conjA], B, IB, [conjB], [α=1, [β=1]]; [backend=..., allocator=...])
+    tensoradd(A, pA::Index2Tuple, conjA, B, pB::Index2Tuple, conjB, α=1, β=1, [backend, allocator]) # expert mode
 
 Return the result of adding arrays `A` and `B` where the iterables `IA` and `IB`
 denote how the array data should be permuted in order to be added. More specifically,
@@ -81,23 +103,26 @@ See also [`tensoradd!`](@ref).
 """
 function tensoradd end
 
-function tensoradd(IC::Labels, A, IA::Labels, conjA::Bool, B, IB::Labels,
-                   conjB::Bool, α::Number=One(), β::Number=One())
-    return tensoradd(A, add_indices(IA, IC), conjA, B, add_indices(IB, IC), conjB, α, β)
+function tensoradd(IC::Labels, A, IA::Labels, conjA::Bool, B, IB::Labels, conjB::Bool,
+                   α::Number=One(), β::Number=One(); kwargs...)
+    pA = add_indices(IA, IC)
+    pB = add_indices(IB, IC)
+    return tensoradd(A, pA, conjA, B, pB, conjB, α, β, _kwargs2args(; kwargs...)...)
 end
 # default `IC`
 function tensoradd(A, IA::Labels, conjA::Bool, B, IB::Labels, conjB::Bool,
-                   α::Number=One(), β::Number=One())
-    return tensoradd(IA, A, IA, conjA, B, IB, conjB, α, β)
+                   α::Number=One(), β::Number=One(); kwargs...)
+    return tensoradd(IA, A, IA, conjA, B, IB, conjB, α, β; kwargs...)
 end
 # default `conjA` and `conjB`
-function tensoradd(IC::Labels, A, IA::Labels, B, IB::Labels, α::Number=One(),
-                   β::Number=One())
+function tensoradd(IC::Labels, A, IA::Labels, B, IB::Labels,
+                   α::Number=One(), β::Number=One(); kwargs...)
     return tensoradd(IC, A, IA, false, B, IB, false, α, β)
 end
 # default `IC`, `conjA` and `conjB`
-function tensoradd(A, IA::Labels, B, IB::Labels, α::Number=One(), β::Number=One())
-    return tensoradd(IA, A, IA, B, IB, α, β)
+function tensoradd(A, IA::Labels, B, IB::Labels, α::Number=One(), β::Number=One();
+                   kwargs...)
+    return tensoradd(IA, A, IA, B, IB, α, β; kwargs...)
 end
 # expert mode
 function tensoradd(A, pA::Index2Tuple, conjA::Bool,
@@ -114,8 +139,8 @@ end
 # tensortrace
 # ------------------------------------------------------------------------------------------
 """
-    tensortrace([IC], A, IA, [conjA], [α=1])
-    tensortrace(A, p::Index2Tuple, q::Index2Tuple, conjA, α=1, [backend]) # expert mode
+    tensortrace([IC], A, IA, [conjA], [α=1]; [backend=..., allocator=...])
+    tensortrace(A, p::Index2Tuple, q::Index2Tuple, conjA, α=1, [backend, allocator]) # expert mode
 
 Trace or contract pairs of indices of tensor `A`, by assigning them identical indices in the
 iterable `IA`. The untraced indices, which are assigned a unique index, can be reordered
@@ -130,22 +155,21 @@ See also [`tensortrace!`](@ref).
 """
 function tensortrace end
 
+function tensortrace(IC::Labels, A, IA::Labels, conjA::Bool, α::Number=One(); kwargs...)
+    p, q = trace_indices(IA, IC)
+    return tensortrace(A, p, q, conjA, α, _kwargs2args(; kwargs...)...)
+end
 # default `IC`
-function tensortrace(A, IA::Labels, conjA::Bool, α::Number=One())
-    return tensortrace(unique2(IA), A, IA, conjA, α)
+function tensortrace(A, IA::Labels, conjA::Bool, α::Number=One(); kwargs...)
+    return tensortrace(unique2(IA), A, IA, conjA, α; kwargs...)
 end
 # default `conjA`
-function tensortrace(IC::Labels, A, IA::Labels, α::Number=One())
-    return tensortrace(IC, A, IA, false, α)
+function tensortrace(IC::Labels, A, IA::Labels, α::Number=One(); kwargs...)
+    return tensortrace(IC, A, IA, false, α; kwargs...)
 end
 # default `IC` and `conjA`
-function tensortrace(A, IA::Labels, α::Number=One())
-    return tensortrace(unique2(IA), A, IA, false, α)
-end
-# labels to indices
-function tensortrace(IC::Labels, A, IA::Labels, conjA::Bool, α::Number=One())
-    p, q = trace_indices(IA, IC)
-    return tensortrace(A, p, q, conjA, α)
+function tensortrace(A, IA::Labels, α::Number=One(); kwargs...)
+    return tensortrace(unique2(IA), A, IA, false, α; kwargs...)
 end
 # expert mode
 function tensortrace(A, p::Index2Tuple, q::Index2Tuple, conjA::Bool, α::Number=One(),
@@ -158,10 +182,9 @@ end
 # ------------------------------------------------------------------------------------------
 # tensorcontract
 # ------------------------------------------------------------------------------------------
-
 """
-    tensorcontract([IC], A, IA, [conjA], B, IB, [conjB], [α=1])
-    tensorcontract(A, pA::Index2Tuple, conjA, B, pB::Index2Tuple, conjB, pAB::Index2Tuple, α=1, [backend]) # expert mode
+    tensorcontract([IC], A, IA, [conjA], B, IB, [conjB], [α=1]; [backend=..., allocator=...])
+    tensorcontract(A, pA::Index2Tuple, conjA, B, pB::Index2Tuple, conjB, pAB::Index2Tuple, α=1, [backend, allocator]) # expert mode
 
 Contract indices of tensor `A` with corresponding indices in tensor `B` by assigning
 them identical labels in the iterables `IA` and `IB`. The indices of the resulting
@@ -180,22 +203,23 @@ See also [`tensorcontract!`](@ref).
 function tensorcontract end
 
 function tensorcontract(IC::Labels, A, IA::Labels, conjA::Bool, B, IB::Labels, conjB::Bool,
-                        α::Number=One())
+                        α::Number=One(); kwargs...)
     pA, pB, pAB = contract_indices(IA, IB, IC)
-    return tensorcontract(A, pA, conjA, B, pB, conjB, pAB, α)
+    return tensorcontract(A, pA, conjA, B, pB, conjB, pAB, α, _kwargs2args(; kwargs...)...)
 end
 # default `IC`
 function tensorcontract(A, IA::Labels, conjA::Bool, B, IB::Labels, conjB::Bool,
-                        α::Number=One())
-    return tensorcontract(symdiff(IA, IB), A, IA, conjA, B, IB, conjB, α)
+                        α::Number=One(); kwargs...)
+    return tensorcontract(symdiff(IA, IB), A, IA, conjA, B, IB, conjB, α; kwargs...)
 end
 # default `conjA` and `conjB`
-function tensorcontract(IC::Labels, A, IA::Labels, B, IB::Labels, α::Number=One())
-    return tensorcontract(IC, A, IA, false, B, IB, false, α)
+function tensorcontract(IC::Labels, A, IA::Labels, B, IB::Labels, α::Number=One();
+                        kwargs...)
+    return tensorcontract(IC, A, IA, false, B, IB, false, α; kwargs...)
 end
 # default `IC`, `conjA` and `conjB`
-function tensorcontract(A, IA::Labels, B, IB::Labels, α::Number=One())
-    return tensorcontract(symdiff(IA, IB), A, IA, false, B, IB, false, α)
+function tensorcontract(A, IA::Labels, B, IB::Labels, α::Number=One(); kwargs...)
+    return tensorcontract(symdiff(IA, IB), A, IA, false, B, IB, false, α; kwargs...)
 end
 # expert mode
 function tensorcontract(A, pA::Index2Tuple, conjA::Bool,
@@ -213,8 +237,8 @@ end
 # ------------------------------------------------------------------------------------------
 
 """
-    tensorproduct([IC], A, IA, [conjA], B, IB, [conjB], [α=1])
-    tensorproduct(A, pA::Index2Tuple, conjA, B, pB::Index2Tuple, conjB, pAB::Index2Tuple, α=1, [backend]) # expert mode
+    tensorproduct([IC], A, IA, [conjA], B, IB, [conjB], [α=1]; [backend=..., allocator=...])
+    tensorproduct(A, pA::Index2Tuple, conjA, B, pB::Index2Tuple, conjB, pAB::Index2Tuple, α=1, [backend, allocator]) # expert mode
 
 Compute the tensor product (outer product) of two tensors `A` and `B`, i.e. returns a new
 tensor `C` with `ndims(C) = ndims(A) + ndims(B)`. The indices of the output tensor are
@@ -231,22 +255,22 @@ See also [`tensorproduct!`](@ref) and [`tensorcontract`](@ref).
 function tensorproduct end
 
 function tensorproduct(IC::Labels, A, IA::Labels, conjA::Bool, B, IB::Labels, conjB::Bool,
-                       α::Number=One())
+                       α::Number=One(); kwargs...)
     pA, pB, pAB = contract_indices(IA, IB, IC)
-    return tensorproduct(A, pA, conjA, B, pB, conjB, pAB, α)
+    return tensorproduct(A, pA, conjA, B, pB, conjB, pAB, α, _kwargs2args(; kwargs...)...)
 end
 # default `IC`
 function tensorproduct(A, IA::Labels, conjA::Bool, B, IB::Labels, conjB::Bool,
-                       α::Number=One())
-    return tensorproduct(vcat(IA, IB), A, IA, conjA, B, IB, conjB, α)
+                       α::Number=One(); kwargs...)
+    return tensorproduct(vcat(IA, IB), A, IA, conjA, B, IB, conjB, α; kwargs...)
 end
 # default `conjA` and `conjB`
-function tensorproduct(IC::Labels, A, IA::Labels, B, IB::Labels, α::Number=One())
-    return tensorproduct(IC, A, IA, false, B, IB, false, α)
+function tensorproduct(IC::Labels, A, IA::Labels, B, IB::Labels, α::Number=One(); kwargs...)
+    return tensorproduct(IC, A, IA, false, B, IB, false, α; kwargs...)
 end
 # default `IC`, `conjA` and `conjB`
-function tensorproduct(A, IA::Labels, B, IB::Labels, α::Number=One())
-    return tensorproduct(vcat(IA, IB), A, IA, false, B, IB, false, α)
+function tensorproduct(A, IA::Labels, B, IB::Labels, α::Number=One(); kwargs...)
+    return tensorproduct(vcat(IA, IB), A, IA, false, B, IB, false, α; kwargs...)
 end
 # expert mode
 function tensorproduct(A, pA::Index2Tuple, conjA::Bool,
@@ -259,11 +283,14 @@ function tensorproduct(A, pA::Index2Tuple, conjA::Bool,
 end
 
 """
-    tensorproduct!(C, A, pA::Index2Tuple, conjA, B, pB::Index2Tuple, conjB, pAB::Index2Tuple, α=1, β=0)
+    tensorproduct!(C, A, pA::Index2Tuple, conjA, B, pB::Index2Tuple, conjB, pAB::Index2Tuple, α=1, β=0, [backend, allocator])
 
 Compute the tensor product (outer product) of two tensors `A` and `B`, i.e. a wrapper of
 [`tensorcontract!`](@ref) with no indices being contracted over. This method checks whether
-the indices indeed specify a tensor product instead of a genuine contraction.
+the indices indeed specify a tensor product instead of a genuine contraction.  Optionally 
+specify a backend implementation to use, and an allocator to be used if temporary tensor
+objects are needed.
+
 
 !!! warning 
     The object `C` must not be aliased with `A` or `B`.

src/implementation/ncon.jl

@@ -1,5 +1,5 @@
 """
-    ncon(tensorlist, indexlist, [conjlist, sym]; order = ..., output = ...)
+    ncon(tensorlist, indexlist, [conjlist, sym]; order = ..., output = ..., backend = ..., allocator = ...)
 
 Contract the tensors in `tensorlist` (of type `Vector` or `Tuple`) according to the network
 as specified by `indexlist`. Here, `indexlist` is a list (i.e. a `Vector` or `Tuple`) with
@@ -24,39 +24,39 @@ See also the macro version [`@ncon`](@ref).
 """
 function ncon(tensors, network,
               conjlist=fill(false, length(tensors));
-              order=nothing, output=nothing)
+              order=nothing, output=nothing, kwargs...)
     length(tensors) == length(network) == length(conjlist) ||
         throw(ArgumentError("number of tensors and of index lists should be the same"))
     isnconstyle(network) || throw(ArgumentError("invalid NCON network: $network"))
     output′ = nconoutput(network, output)
 
     if length(tensors) == 1
         if length(output′) == length(network[1])
-            return tensorcopy(output′, tensors[1], network[1], conjlist[1])
+            return tensorcopy(output′, tensors[1], network[1], conjlist[1]; kwargs...)
         else
-            return tensortrace(output′, tensors[1], network[1], conjlist[1])
+            return tensortrace(output′, tensors[1], network[1], conjlist[1]; kwargs...)
         end
     end
 
     (tensors, network) = resolve_traces(tensors, network)
     tree = order === nothing ? ncontree(network) : indexordertree(network, order)
 
-    A, IA, CA = contracttree(tensors, network, conjlist, tree[1])
-    B, IB, CB = contracttree(tensors, network, conjlist, tree[2])
+    A, IA, CA = contracttree(tensors, network, conjlist, tree[1]; kwargs...)
+    B, IB, CB = contracttree(tensors, network, conjlist, tree[2]; kwargs...)
     IC = tuple(output′...)
 
-    return tensorcontract(IC, A, IA, CA, B, IB, CB)
+    return tensorcontract(IC, A, IA, CA, B, IB, CB; kwargs...)
 end
 
-function contracttree(tensors, network, conjlist, tree)
+function contracttree(tensors, network, conjlist, tree; kwargs...)
     @nospecialize
     if tree isa Int
         return tensors[tree], tuple(network[tree]...), (conjlist[tree])
     end
-    A, IA, CA = contracttree(tensors, network, conjlist, tree[1])
-    B, IB, CB = contracttree(tensors, network, conjlist, tree[2])
+    A, IA, CA = contracttree(tensors, network, conjlist, tree[1]; kwargs...)
+    B, IB, CB = contracttree(tensors, network, conjlist, tree[2]; kwargs...)
     IC = tuple(symdiff(IA, IB)...)
-    C = tensorcontract(IC, A, IA, CA, B, IB, CB)
+    C = tensorcontract(IC, A, IA, CA, B, IB, CB; kwargs...)
     return C, IC, false
 end