support compactifications of more sophisticated types (#264)

Now compactification works for more types. In particular see these examples of things that previously could not be compactified:

using Revise
using QuantumClifford
using BenchmarkTools

function x_diag_circuit_noisy_measurement(csize)
    circuit = []
    for i in 1:csize
        push!(circuit, PauliError(i, 0.1))
        push!(circuit, sHadamard(i))
        push!(circuit, sCNOT(i, csize+1))
        push!(circuit, sMZ(csize+1,i))
        push!(circuit, ClassicalXOR(1:(i%6+2),i))
    end
    return circuit
end

@benchmark pftrajectories(state,circuit) setup=(state=PauliFrame(1000, 1001, 1001); circuit=x_diag_circuit_noisy_measurement(1000)) evals=1

BenchmarkTools.Trial: 1126 samples with 1 evaluation.
 Range (min … max):  3.532 ms …  4.088 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     3.573 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):   3.577 ms ± 27.682 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

              ▁▆▄▅▇▄▇█▆▆▃▆▄▅▅▄▂▁▂▁
  ▃▂▁▃▃▁▂▄▄▄▄▇█████████████████████▇▇▆▆▆▆▄▆▃▃▅▅▄▄▃▂▃▃▂▃▂▂▃▂▃ ▅
  3.53 ms        Histogram: frequency by time        3.63 ms <

 Memory estimate: 281.25 KiB, allocs estimate: 6000.

@benchmark pftrajectories(state,circuit) setup=(state=PauliFrame(1000, 1001, 1001); circuit=compactify_circuit(x_diag_circuit_noisy_measurement(1000))) evals=1

Before:
BenchmarkTools.Trial: 53 samples with 1 evaluation.
 Range (min … max):  3.495 ms …   5.890 ms  ┊ GC (min … max): 0.00% … 38.90%
 Time  (median):     3.623 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   3.711 ms ± 370.663 μs  ┊ GC (mean ± σ):  1.16% ±  5.34%

     █▃▂  ▂▂
  ▄▄████▇▇██▄▇▇▁▁▄▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▄▁▁▁▄ ▁
  3.5 ms          Histogram: frequency by time        4.71 ms <

 Memory estimate: 421.97 KiB, allocs estimate: 9002.

After:
BenchmarkTools.Trial: 1116 samples with 1 evaluation.
 Range (min … max):  3.221 ms …   4.887 ms  ┊ GC (min … max): 0.00% … 28.58%
 Time  (median):     3.388 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   3.357 ms ± 150.766 μs  ┊ GC (mean ± σ):  0.36% ±  2.70%

                                      ▄█▆▂
  ▂▂▃▆▆▇▇▅▇▆▆▅▄▄▄▄▃▃▃▃▂▂▃▃▂▂▁▂▂▂▁▁▁▃▃▅█████▇▆▃▄▂▂▂▂▁▁▂▃▃▃▃▁▁▂ ▃
  3.22 ms         Histogram: frequency by time        3.49 ms <

 Memory estimate: 187.50 KiB, allocs estimate: 4000.

Author: Krastanov

CHANGELOG.md

@@ -9,6 +9,7 @@
 
 - Gate errors are now conveniently supported by the various ECC benchmark setups in the `ECC` module.
 - Remove printing of spurious debug info from the PyBP decoder. 
+- Significant improvements to the low-level circuit compiler (the sumtype compactifier), leading to faster Pauli frame simulation of noisy circuits.
 
 ## v0.9.3 - 2024-04-10
 

benchmark/benchmarks.jl

@@ -158,4 +158,25 @@ for (cs, c) in [("shor",Shor9()), ("toric8",Toric(8,8))]
     end
 end
 
+
+if V > v"0.9.0"
+
+function x_diag_circuit_noisy_measurement(csize)
+    circuit = []
+    for i in 1:csize
+        push!(circuit, PauliError(i, 0.1))
+        push!(circuit, sHadamard(i))
+        push!(circuit, sCNOT(i, csize+1))
+        push!(circuit, sMZ(csize+1,i))
+        push!(circuit, ClassicalXOR(1:(i%6+2),i))
+    end
+    return circuit
+end
+
+SUITE["circuitsim"]["compactification"] = BenchmarkGroup(["compactification"])
+SUITE["circuitsim"]["compactification"]["no_compact"] = @benchmarkable pftrajectories(state,circuit) setup=(state=PauliFrame(1000, 1001, 1001); circuit=x_diag_circuit_noisy_measurement(1000)) evals=1
+SUITE["circuitsim"]["compactification"]["compact"] = @benchmarkable pftrajectories(state,circuit) setup=(state=PauliFrame(1000, 1001, 1001); circuit=compactify_circuit(x_diag_circuit_noisy_measurement(1000))) evals=1
+
+end
+
 end

ext/QuantumCliffordQuantikzExt/QuantumCliffordQuantikzExt.jl

@@ -4,7 +4,6 @@ import Quantikz
 using QuantumClifford
 using QuantumClifford.Experimental.NoisyCircuits
 using QuantumClifford: AbstractOperation
-using QuantumClifford: ClassicalXORConcreteWorkaround
 
 function Quantikz.QuantikzOp(op::SparseGate)
     g = op.cliff
@@ -73,7 +72,7 @@ function Quantikz.QuantikzOp(op::Reset) # TODO This is complicated because quant
 end
 Quantikz.QuantikzOp(op::NoiseOp) = Quantikz.Noise(collect(op.indices))
 Quantikz.QuantikzOp(op::NoiseOpAll) = Quantikz.NoiseAll()
-Quantikz.QuantikzOp(op::ClassicalXORConcreteWorkaround) = Quantikz.ClassicalDecision(sort([op.store, op.bits...]))
+Quantikz.QuantikzOp(op::ClassicalXOR) = Quantikz.ClassicalDecision(sort([op.store, op.bits...]))
 
 function lstring(pauli::PauliOperator)
     v = join(("\\mathtt{$(o)}" for o in replace(string(pauli)[3:end],"_"=>"I")),"\\\\")

src/affectedqubits.jl

@@ -12,9 +12,9 @@ affectedqubits(p::PauliOperator) = 1:length(p)
 affectedqubits(m::Union{AbstractMeasurement,sMRX,sMRY,sMRZ}) = (m.qubit,)
 affectedqubits(v::VerifyOp) = v.indices
 affectedqubits(c::CliffordOperator) = 1:nqubits(c)
-affectedqubits(c::ClassicalXORConcreteWorkaround) = ()
+affectedqubits(c::ClassicalXOR) = ()
 
 affectedbits(o) = ()
 affectedbits(m::sMRZ) = (m.bit,)
 affectedbits(m::sMZ) = (m.bit,)
-affectedbits(c::ClassicalXORConcreteWorkaround) = (c.bits..., c.store)
+affectedbits(c::ClassicalXOR) = (c.bits..., c.store)

src/misc_ops.jl

@@ -127,34 +127,15 @@ end
 
 operatordeterminism(::Type{VerifyOp}) = DeterministicOperatorTrait()
 
-abstract type ClassicalXORConcreteWorkaround <: AbstractOperation end  # See below for more of this abomination - replace everywhere by ClassicalXOR when compactification is fixed
 """Applies an XOR gate to classical bits. Currently only implemented for functionality with pauli frames."""
-struct ClassicalXOR{N} <: ClassicalXORConcreteWorkaround
+struct ClassicalXOR{N} <: AbstractOperation
     "The indices of the classical bits to be xor-ed"
     bits::NTuple{N,Int}
     "The index of the classical bit that will store the results"
     store::Int
-    function ClassicalXOR(bits, store) # See below for more of this abomination
+    function ClassicalXOR(bits, store)
         tbits = tuple(bits...)
-        n = length(bits)
-        if n <= 15
-            return eval(Symbol("ClassicalXOR",string(n)))(tbits, store)
-        else
-            return new{n}(tbits, store)
-        end
+        n = length(tbits)
+        return new{n}(tbits, store)
     end
 end
-#ClassicalXOR(bits::Vector, store::Int) = ClassicalXOR(tuple(bits...),store)
-# XXX TODO remove this abomination
-# Workaround for not being able to compactify non-concrete types
-for n in 2:15
-    name = Symbol("ClassicalXOR",string(n))
-    eval(
-        quote
-            struct $name <: ClassicalXORConcreteWorkaround
-                bits::NTuple{$n,Int}
-                store::Int
-            end
-        end
-    )
-end

src/pauli_frames.jl

@@ -57,7 +57,7 @@ function apply!(f::PauliFrame, op::AbstractCliffordOperator)
     return f
 end
 
-function apply!(frame::PauliFrame, xor::ClassicalXORConcreteWorkaround)
+function apply!(frame::PauliFrame, xor::ClassicalXOR)
     for f in eachindex(frame)
         value = frame.measurements[f,xor.bits[1]]
         for i in xor.bits[2:end]

src/sumtypes.jl

@@ -1,14 +1,33 @@
+# Here be dragons...
+
 using SumTypes
 using InteractiveUtils: subtypes
 
+"""An intermediary when we want to create a new concrete type in a macro."""
+struct SymbolicDataType
+    name::Symbol
+    types#::Core.SimpleVector
+    fieldnames
+    originaltype
+end
+_header(s) = s
+_header(s::SymbolicDataType) = s.name
+_symbol(s) = Symbol(s)
+_symbol(s::SymbolicDataType) = s.name
+_types(s) = s.types
+_fieldnames(s) = fieldnames(s)
+_fieldnames(s::SymbolicDataType) = s.fieldnames
+_originaltype(s) = s
+_originaltype(s::SymbolicDataType) = s.originaltype
+
 """
 ```
 julia> make_variant(sCNOT)
 :(sCNOT(::Int64, ::Int64))
 ```
 """
-function make_variant(type::DataType)
-    Expr(:call, Symbol(type), [:(::$t) for t in type.types]...)
+function make_variant(type::Union{DataType,SymbolicDataType})
+    Expr(:call, _symbol(type), [:(::$t) for t in _types(type)]...)
 end
 
 """
@@ -17,9 +36,9 @@ julia> make_variant_deconstruct(sCNOT, :apply!, (:s,))
 :(sCNOT(q1, q2) => apply!(s, sCNOT(q1, q2)))
 ```
 """
-function make_variant_deconstruct(type::DataType, call, preargs=(), postargs=())
-    variant = Expr(:call, Symbol(type), fieldnames(type)...)
-    original = :(($type)($(fieldnames(type)...)))
+function make_variant_deconstruct(type::Union{DataType,SymbolicDataType}, call, preargs=(), postargs=())
+    variant = Expr(:call, _symbol(type), _fieldnames(type)...)
+    original = :(($(_originaltype(type)))($(_fieldnames(type)...)))
     :($variant => $(Expr(:call, call, preargs..., original, postargs...)))
 end
 
@@ -36,7 +55,7 @@ end
 function make_sumtype(concrete_types)
     return quote
         @sum_type CompactifiedGate :hidden begin
-            $([make_variant(t) for t in concrete_types if isa(t, DataType)]...)
+            $([make_variant(t) for t in concrete_types if isa(t, DataType) || isa(t, SymbolicDataType)]...)
         end
     end
 end
@@ -56,7 +75,8 @@ function make_sumtype_method(concrete_types, call, preargs=(), postargs=())
     return quote
         function QuantumClifford.$call($(preargs...), g::CompactifiedGate, $(postargs...))
             @cases g begin
-                $([make_variant_deconstruct(t, call, preargs, postargs) for t in concrete_types if isa(t, DataType)]...)
+                $([make_variant_deconstruct(t, call, preargs, postargs) for t in concrete_types if isa(t, DataType) || isa(t, SymbolicDataType)]...)
+                #_ => @error "something wrong is happening when working with $(g) -- you are probably getting wrong results, please report this as a bug" # this being present ruins some safety guarantees, but it is useful for debugging
             end
         end
     end
@@ -71,40 +91,89 @@ end)
 ```
 """
 function make_sumtype_variant_constructor(type)
-    if isa(type, DataType)
-        return :( CompactifiedGate(g::$(type)) = CompactifiedGate'.$(Symbol(type))($([:(g.$n) for n in fieldnames(type)]...)) )
+    if isa(type, DataType) || isa(type, SymbolicDataType)
+        return :( CompactifiedGate(g::$(_header(type))) = CompactifiedGate'.$(_symbol(type))($([:(g.$n) for n in _fieldnames(type)]...)) )
     else
-        return :( CompactifiedGate(g::$(type)) = (@warn "The operation is of a type that can not be unified, defaulting to slower runtime dispatch" typeof(g); return g) )
+        #return :( CompactifiedGate(g::$(_header(type))) = (@warn "The operation is of a type that can not be unified, defaulting to slower runtime dispatch" typeof(g); return g) )
+        return :()
     end
 end
 
+genericsupertypeparams(t) = :body ∈ propertynames(t) ? genericsupertypeparams(t.body) : t
+
+"""Returns a tuple of all concrete subtypes and all UnionAll non-abstract subtypes of a given type."""
+function get_all_subtypes(type)
+    if !isabstracttype(type)
+        if isa(type, DataType)
+            isbitstype(type) || @debug "$type will be problematic during compactification"
+            return [type], []
+        elseif isa(type, UnionAll)
+            return [], [type]
+        else
+            @error "The gate compiler has encountered a type that it can not handle: $type. The QuantumClifford library should continue functioning, but potentially at degraded performance. Please report this as a performance bug."
+        end
+    else
+        return Iterators.flatten.(zip(get_all_subtypes.(subtypes(type))...))
+    end
+end
 
-function make_all_sumtype_infrastructure_expr(concrete_types, callsigs)
+module_of_type(t::UnionAll) = genericsupertypeparams(t).name.module
+module_of_type(t::DataType) = t.name.module
+
+function make_all_sumtype_infrastructure_expr(t::DataType, callsigs)
+    concrete_types, unionall_types = get_all_subtypes(t)
+    concrete_types = collect(Any, concrete_types)
+    concrete_types = Any[t for t in concrete_types if module_of_type(t)==QuantumClifford]
+    unionall_types = Any[t for t in unionall_types if module_of_type(t)==QuantumClifford]
+    concretifier_workarounds_types = [] # e.g. var"ClassicalXOR_{2}" generated as a workaround for providing a concrete type for ClassicalXOR{N}
+    concretifier_additional_constructors = [] # e.g. CompactifiedGate(g::ClassicalXOR{2}) = CompactifiedGate'.var"ClassicalXOR_{2}"(g.bits, g.store)
+    for ut in unionall_types
+        names, generated_concretetypes, generated_variant_constructors = concretifier(ut)
+        append!(concretifier_workarounds_types, generated_concretetypes)
+        append!(concrete_types, names)
+        append!(concretifier_additional_constructors, generated_variant_constructors)
+        push!(concrete_types, ut) # fallback
+    end
     sumtype = make_sumtype(concrete_types)
     constructors = make_sumtype_variant_constructor.(concrete_types)
     methods = [make_sumtype_method(concrete_types, call, preargs, postargs) for (call, preargs, postargs) in callsigs]
     return quote
-        $(sumtype.args...)
-        $(constructors...)
+        $(concretifier_workarounds_types...)
+        $(sumtype.args...) # defining the sum type
+        $(constructors...) # creating constructors for the sumtype which turn our concrete types into instance of the sum type
+        $(concretifier_additional_constructors...) # creating constructors for the newly generated "workaround" concrete types
+        :( CompactifiedGate(g::AbstractOperation) = (@warn "The operation is of a type that can not be unified, defaulting to slower runtime dispatch" typeof(g); return g) )
         $(methods...)
     end
 end
 
-function get_all_concrete_subtypes(type)
-    if !isabstracttype(type)
-        return [type]
-    else
-        return vcat(get_all_concrete_subtypes.(subtypes(type))...)
-    end
+function concrete_typeparams(t)
+    @debug "The gate compiler is not able to concretify the type $t. Define a `concrete_typeparams` method for this type to improve performance."
+    return ()
 end
 
-module_of_type(t::UnionAll) = module_of_type(t.body)
-module_of_type(t::DataType) = t.name.module
+function concretifier(t)
+    names = []
+    generated_concretetypes = []
+    generated_variant_constructors = []
+    for typeparams in concrete_typeparams(t)
+        name = Symbol(t,"{",typeparams,"}")
+        parameterized_type = t{typeparams...}
+        ftypes = parameterized_type.types
+        fnames = fieldnames(t)
+        push!(names, SymbolicDataType(name, ftypes, fnames, t))
+        push!(generated_concretetypes, :(
+            struct $(name)
+                $([:($n::$t) for (n,t) in zip(fnames,ftypes)]...)
+            end
+        ))
+        push!(generated_variant_constructors, make_concretifier_sumtype_variant_constructor(parameterized_type, name))
+    end
+    return names, generated_concretetypes, generated_variant_constructors
+end
 
-function make_all_sumtype_infrastructure_expr(t::DataType, callsigs)
-    concrete_types = get_all_concrete_subtypes(t)
-    non_experimental_concrete_types = [t for t in concrete_types if module_of_type(t)==QuantumClifford]
-    make_all_sumtype_infrastructure_expr(non_experimental_concrete_types, callsigs)
+function make_concretifier_sumtype_variant_constructor(parameterized_type, variant_name)
+    return :( CompactifiedGate(g::$(parameterized_type)) = CompactifiedGate'.$(variant_name)($([:(g.$n) for n in _fieldnames(parameterized_type)]...)) )
 end
 
 function make_all_sumtype_infrastructure()
@@ -120,8 +189,6 @@ function make_all_sumtype_infrastructure()
     ) |> eval
 end
 
-make_all_sumtype_infrastructure()
-
 """
 Convert a list of gates to a more optimized "sum type" format which permits faster dispatch.
 
@@ -130,3 +197,21 @@ Generally, this should be called on a circuit before it is used in a simulation.
 function compactify_circuit(circuit)
     return CompactifiedGate.(circuit)
 end
+
+
+##
+# `concrete_typeparams` annotations for the parameteric types we care about
+##
+
+function concrete_typeparams(t::Type{ClassicalXOR})
+    return 2:16
+end
+
+function concrete_typeparams(t::Type{NoiseOp})
+    return [(UnbiasedUncorrelatedNoise{Float64}, i) for i in 1:8]
+end
+
+
+# XXX This has to happen after defining all the `concrete_typeparams` methods
+
+make_all_sumtype_infrastructure()