Use linear scans instead of hashing for contraction labels (#190)

## Summary

Does the dimname/label bookkeeping in the out-of-place contraction path
with linear scans instead of the `Set`- and `Dict`-based
`Base.setdiff`/`intersect`/`indexin`. For the handful of labels a tensor
carries the hashing those build dominates, and on a small contraction it
costs more than the matrix multiply itself. `biperms` and
`contract_labels` use order-preserving `Vector` comprehensions for the
output and contracted labels, and `tuple_indexin` finds label positions
with `findfirst` rather than `Base.indexin`. The old check that every
label appears exactly twice is replaced by a check that the destination
carries exactly the uncontracted labels: the contracted and output
groups partition the operands by construction, so that is the only
consistency left to verify. These all assume the labels within each
group are unique, which holds for a contraction.

On a 4x4 contraction sharing one index, `a1 * a2` drops from 4.6 μs /
124 allocations to 1.9 μs / 60.

Author: mtfishman

Project.toml

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

src/contract/biperms.jl

@@ -1,8 +1,12 @@
-# `Base.indexin` doesn't accept tuples; return the positions of `x` in `y` as a tuple.
-function tuple_indexin(x::Tuple, y::AbstractArray)
-    return Tuple{Vararg{Any, length(x)}}(Base.indexin(x, y))
-end
-tuple_indexin(x::Tuple, y) = tuple_indexin(x, collect(y))
+# `a ∖ b` and `a ∩ b` as a `Vector`, preserving the order of `a`, via a linear
+# scan. For the small collections here `Base.setdiff`/`intersect` are slower
+# because they build a `Set` and hash. Both assume set-like (unique) inputs.
+smallsetdiff(a, b) = [x for x in a if x ∉ b]
+smallintersect(a, b) = [x for x in a if x ∈ b]
+
+# Position of each element of `x` in `y`, as a tuple. Linear scan, no hashing
+# (`Base.indexin` builds a `Dict`), for the small collections here.
+tuple_indexin(x::Tuple, y) = map(v -> findfirst(==(v), y), x)
 
 """
     biperm(t, t1, t2) -> (p1, p2)
@@ -30,14 +34,16 @@ length_codomain(t) = length(t) - length_domain(t)
 
 # codomain <-- domain
 function biperms(::typeof(contract), dimnames_dest, dimnames1, dimnames2)
-    dimnames = collect(Iterators.flatten((dimnames_dest, dimnames1, dimnames2)))
-    for i in unique(dimnames)
-        count(==(i), dimnames) == 2 || throw(ArgumentError("Invalid contraction labels"))
-    end
-
-    codomain = Tuple(setdiff(dimnames1, dimnames2))
-    contracted = Tuple(intersect(dimnames1, dimnames2))
-    domain = Tuple(setdiff(dimnames2, dimnames1))
+    codomain = Tuple(smallsetdiff(dimnames1, dimnames2))
+    contracted = Tuple(smallintersect(dimnames1, dimnames2))
+    domain = Tuple(smallsetdiff(dimnames2, dimnames1))
+
+    # `codomain`/`contracted` and `contracted`/`domain` partition the operands by
+    # construction, so the only label consistency left to check is that the
+    # destination carries exactly the uncontracted labels. `biperm` below then
+    # checks each group lands in the destination.
+    length(codomain) + length(domain) == length(dimnames_dest) ||
+        throw(ArgumentError("Invalid contraction labels"))
 
     perm_codomain_dest, perm_domain_dest = biperm(dimnames_dest, codomain, domain)
     invperm_dest = invperm((perm_codomain_dest..., perm_domain_dest...))

src/contract/contract_labels.jl

@@ -2,7 +2,7 @@ function contract_labels(a1::AbstractArray, labels1, a2::AbstractArray, labels2)
     return contract_labels(labels1, labels2)
 end
 function contract_labels(labels1, labels2)
-    diff1 = setdiff(labels1, labels2)
-    diff2 = setdiff(labels2, labels1)
+    diff1 = smallsetdiff(labels1, labels2)
+    diff2 = smallsetdiff(labels2, labels1)
     return vcat(diff1, diff2)
 end

test/test_setoperations.jl

@@ -0,0 +1,33 @@
+using TensorAlgebra: biperms, contract, smallintersect, smallsetdiff, tuple_indexin
+using Test: @test, @test_throws, @testset
+
+@testset "smallsetdiff/smallintersect" begin
+    # Order-preserving, returning a `Vector`.
+    @test smallsetdiff((:i, :j, :k), (:k, :i)) == [:j]
+    @test smallintersect((:i, :j, :k), (:k, :i)) == [:i, :k]
+    @test smallsetdiff([:i, :j, :k], [:k, :i]) == [:j]
+    @test smallintersect([:i, :j, :k], [:k, :i]) == [:i, :k]
+    # Disjoint and empty cases.
+    @test smallsetdiff((:i, :j), ()) == [:i, :j]
+    @test smallintersect((:i, :j), (:k,)) == []
+end
+
+@testset "tuple_indexin" begin
+    # Position of each element of the first tuple in the second collection.
+    @test tuple_indexin((:c, :b), (:a, :b, :c, :d)) == (3, 2)
+    @test tuple_indexin((:c, :b), [:a, :b, :c, :d]) == (3, 2)
+    @test tuple_indexin((), (:a, :b)) == ()
+end
+
+@testset "biperms rejects an inconsistent destination" begin
+    # The destination must carry exactly the uncontracted labels.
+    @test_throws ArgumentError biperms(contract, (:i, :j, :k), (:i, :j), (:j, :k))
+    @test_throws ArgumentError biperms(contract, (:i,), (:i, :j), (:j, :k))
+end
+
+@testset "biperms with tuple and vector labels agree" begin
+    # A representative multi-index contraction: C_il = A_ijk B_kjl.
+    bp = biperms(contract, (:i, :l), (:i, :j, :k), (:k, :j, :l))
+    @test biperms(contract, [:i, :l], (:i, :j, :k), (:k, :j, :l)) == bp
+    @test biperms(contract, (:i, :l), [:i, :j, :k], [:k, :j, :l]) == bp
+end