some performance optimizations in decoders

ext/QuantumCliffordLDPCDecodersExt/QuantumCliffordLDPCDecodersExt.jl

@@ -74,19 +74,25 @@ parity_checks(d::BeliefPropDecoder) = d.H
 parity_checks(d::BitFlipDecoder) = d.H
 
 function decode(d::BeliefPropDecoder, syndrome_sample)
-    row_x = syndrome_sample[1:d.cx]
-    row_z = syndrome_sample[d.cx+1:d.cx+d.cz]
+    row_x = @view syndrome_sample[1:d.cx]
+    row_z = @view syndrome_sample[d.cx+1:d.cx+d.cz]
     guess_z, success = LDPCDecoders.decode!(d.bpdecoderx, row_x)
     guess_x, success = LDPCDecoders.decode!(d.bpdecoderz, row_z)
-    return vcat(guess_x, guess_z)
+    result = Matrix{Int}(undef, 2, length(guess_x))
+    @inbounds result[1, 1:length(guess_x)] .= guess_x
+    @inbounds result[2, 1:length(guess_z)] .= guess_z
+    return result
 end
 
 function decode(d::BitFlipDecoder, syndrome_sample)
-    row_x = syndrome_sample[1:d.cx]
-    row_z = syndrome_sample[d.cx+1:d.cx+d.cz]
+    row_x = @view syndrome_sample[1:d.cx]
+    row_z = @view syndrome_sample[d.cx+1:d.cx+d.cz]
     guess_z, success = LDPCDecoders.decode!(d.bfdecoderx, row_x)
     guess_x, success = LDPCDecoders.decode!(d.bfdecoderz, row_z)
-    return vcat(guess_x, guess_z)
+    result = Matrix{Int}(undef, 2, length(guess_x))
+    @inbounds result[1, 1:length(guess_x)] .= guess_x
+    @inbounds result[2, 1:length(guess_z)] .= guess_z
+    return result
 end
 
 end

ext/QuantumCliffordPyQDecodersExt/QuantumCliffordPyQDecodersExt.jl

@@ -69,11 +69,14 @@ end
 parity_checks(d::PyBP) = d.H
 
 function decode(d::PyBP, syndrome_sample)
-    row_x = syndrome_sample[1:d.nx] # TODO These copies and indirections might be costly!
-    row_z = syndrome_sample[d.nx+1:end]
+    row_x = @view syndrome_sample[1:d.nx]
+    row_z = @view syndrome_sample[d.nx+1:end]
     guess_z_errors = PythonCall.PyArray(d.pyx.decode(np.array(row_x)))
     guess_x_errors = PythonCall.PyArray(d.pyz.decode(np.array(row_z)))
-    vcat(guess_x_errors, guess_z_errors)
+    result = Matrix{Int}(undef, 2, length(guess_x_errors))
+    @inbounds result[1, 1:length(guess_x_errors)] .= guess_x_errors
+    @inbounds result[2, 1:length(guess_z_errors)] .= guess_z_errors
+    return result
 end
 
 struct PyMatchingDecoder <: AbstractSyndromeDecoder # TODO all these decoders have the same fields, maybe we can factor out a common type
@@ -106,19 +109,26 @@ end
 parity_checks(d::PyMatchingDecoder) = d.H
 
 function decode(d::PyMatchingDecoder, syndrome_sample)
-    row_x = syndrome_sample[1:d.nx] # TODO This copy is costly!
-    row_z = syndrome_sample[d.nx+1:end]
+    row_x = @view syndrome_sample[1:d.nx]
+    row_z = @view syndrome_sample[d.nx+1:end]
     guess_z_errors = PythonCall.PyArray(d.pyx.decode(row_x))
     guess_x_errors = PythonCall.PyArray(d.pyz.decode(row_z))
-    vcat(guess_x_errors, guess_z_errors)
+    result = Matrix{Int}(undef, 2, length(guess_x_errors))
+    @inbounds result[1, 1:length(guess_x_errors)] .= guess_x_errors
+    @inbounds result[2, 1:length(guess_z_errors)] .= guess_z_errors
+    return result
 end
 
 function batchdecode(d::PyMatchingDecoder, syndrome_samples)
-    row_x = syndrome_samples[:,1:d.nx] # TODO This copy is costly!
-    row_z = syndrome_samples[:,d.nx+1:end]
+    row_x = @view syndrome_samples[:,1:d.nx]
+    row_z = @view syndrome_samples[:,d.nx+1:end]
     guess_z_errors = PythonCall.PyArray(d.pyx.decode_batch(row_x))
     guess_x_errors = PythonCall.PyArray(d.pyz.decode_batch(row_z))
-    hcat(guess_x_errors, guess_z_errors)
+    n_cols_x = size(guess_x_errors, 2)
+    result = Matrix{Int}(undef, size(guess_x_errors, 1), n_cols_x + size(guess_z_errors, 2))
+    @inbounds result[:,1:n_cols_x] .= guess_x_errors
+    @inbounds result[:,n_cols_x+1:end] .= guess_z_errors
+    return result
 end
 
 end

src/ecc/decoder_pipeline.jl

@@ -16,7 +16,8 @@ function batchdecode(d::AbstractSyndromeDecoder, syndrome_samples)
     samples, _s = size(syndrome_samples)
     s == _s || throw(ArgumentError(lazy"The syndromes given to `batchdecode` have the wrong dimensions. The syndrome length is $(_s) while it should be $(s)"))
     results = falses(samples, 2n)
-    for (i,syndrome_sample) in enumerate(eachrow(syndrome_samples))
+    @inbounds for i in 1:samples
+        syndrome_sample = @view syndrome_samples[i,:]
         guess = decode(d, syndrome_sample)# TODO use `decode!`
         isnothing(guess) || (results[i,:] = guess)
     end
@@ -171,10 +172,21 @@ end
 
 function evaluate_guesses(measured_faults, guesses, faults_matrix)
     nsamples = size(guesses, 1)
-    guess_faults = (faults_matrix * guesses') .% 2 # TODO this can be faster and non-allocating by turning it into a loop
     decoded = 0
-    for i in 1:nsamples # TODO this can be faster and non-allocating by having the loop and the matrix multiplication on the line above work together and not store anything
-        (@view guess_faults[:,i]) == (@view measured_faults[i,:]) && (decoded += 1)
+    for i in 1:nsamples
+        is_decoded = true
+        for j in 1:size(faults_matrix, 1)
+            sum_mod = 0
+            @inbounds @simd for k in 1:size(faults_matrix, 2)
+                sum_mod += faults_matrix[j, k] * guesses[i, k]
+            end
+            sum_mod %= 2
+            if sum_mod != measured_faults[i, j]
+                is_decoded = false
+                break
+            end
+        end
+        decoded += is_decoded
     end
     return (nsamples - decoded) / nsamples
 end
@@ -253,8 +265,8 @@ function create_lookup_table(code::Stabilizer)
     lookup_table
 end;
 
-function decode(d::TableDecoder, syndrome_sample)
-    d.lookup_buffer .= syndrome_sample # TODO have this work without data copying, by supporting the correct types, especially in the batch decode case
+function decode(d::TableDecoder, syndrome_sample::AbstractVector{Bool})
+    @inbounds copyto!(d.lookup_buffer, syndrome_sample)
     return get(d.lookup_table, d.lookup_buffer, nothing)
 end