better bit-wrangling abstraction using less boilerplate (#367)

Fe-r-oz · Krastanov · web-flow · commit d3f42d23f62d · 2024-11-02T10:35:34.000-04:00
Co-authored-by: Stefan Krastanov &lt;stefan@krastanov.org&gt;
diff --git a/ext/QuantumCliffordGPUExt/apply_noise.jl b/ext/QuantumCliffordGPUExt/apply_noise.jl
@@ -1,16 +1,11 @@
-using QuantumClifford: _div, _mod
+using QuantumClifford: get_bitmask_idxs
 
 #according to https://github.com/JuliaGPU/CUDA.jl/blob/ac1bc29a118e7be56d9edb084a4dea4224c1d707/test/core/device/random.jl#L33
 #CUDA.jl supports calling rand() inside kernel
 function applynoise!(frame::PauliFrameGPU{T},noise::UnbiasedUncorrelatedNoise,i::Int) where {T <: Unsigned}
     p = noise.p
-    lowbit = T(1)
-    ibig = _div(T,i-1)+1
-    ismall = _mod(T,i-1)
-    ismallm = lowbit<<(ismall)
-
-    stab = frame.frame
-    xzs = tab(stab).xzs
+    xzs = tab(frame.frame).xzs
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
     rows = size(stab, 1)
 
     @run_cuda applynoise_kernel(xzs, p, ibig, ismallm, rows) rows
diff --git a/ext/QuantumCliffordGPUExt/pauli_frames.jl b/ext/QuantumCliffordGPUExt/pauli_frames.jl
@@ -1,3 +1,5 @@
+using QuantumClifford: get_bitmask_idxs
+
 ##############################
 # sMZ
 ##############################
@@ -21,10 +23,7 @@ function apply!(frame::PauliFrameGPU{T}, op::QuantumClifford.sMZ) where {T <: Un
     op.bit == 0 && return frame
     i = op.qubit
     xzs = frame.frame.tab.xzs
-    lowbit = T(1)
-    ibig = QuantumClifford._div(T,i-1)+1
-    ismall = QuantumClifford._mod(T,i-1)
-    ismallm = lowbit<<(ismall)
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
     (@run_cuda apply_sMZ_kernel!(xzs, frame.measurements, op, ibig, ismallm, length(frame)) length(frame))
     return frame
 end
@@ -55,10 +54,7 @@ end
 function apply!(frame::PauliFrameGPU{T}, op::QuantumClifford.sMRZ) where {T <: Unsigned} # TODO sMRX, sMRY
     i = op.qubit
     xzs = frame.frame.tab.xzs
-    lowbit = T(1)
-    ibig = QuantumClifford._div(T,i-1)+1
-    ismall = QuantumClifford._mod(T,i-1)
-    ismallm = lowbit<<(ismall)
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
     (@run_cuda apply_sMRZ_kernel!(xzs, frame.measurements, op, ibig, ismallm, length(frame)) length(frame))
     return frame
 end
diff --git a/src/QuantumClifford.jl b/src/QuantumClifford.jl
@@ -912,6 +912,30 @@ function unsafe_bitfindnext_(chunks::AbstractVector{T}, start::Int) where T<:Uns
     return nothing
 end
 
+"""
+$(TYPEDSIGNATURES)
+
+Computes bitmask indices for an unsigned integer at index `i`
+within the binary structure of a `Tableau` or `PauliOperator`.
+
+For `Tableau`, the method operates on the `.xzs` field, while
+for `PauliOperator`, it uses the `.xz` field. It calculates
+the following values based on the index `i`:
+
+- `lowbit`, the lowest bit.
+- `ibig`, the index of the word containing the bit.
+- `ismall`, the position of the bit within the word.
+- `ismallm`, a bitmask isolating the specified bit.
+"""
+@inline function get_bitmask_idxs(xzs::AbstractArray{<:Unsigned}, i::Int)
+    T = eltype(xzs)
+    lowbit = T(1)
+    ibig = _div(T, i-1) + 1
+    ismall = _mod(T, i-1)
+    ismallm = lowbit << ismall
+    return lowbit, ibig, ismall, ismallm
+end
+
 """Permute the qubits (i.e., columns) of the tableau in place."""
 function Base.permute!(s::Tableau, perm::AbstractVector)
     for r in 1:size(s,1)
diff --git a/src/pauli_frames.jl b/src/pauli_frames.jl
@@ -82,12 +82,8 @@ end
 function apply!(frame::PauliFrame, op::sMZ) # TODO sMY, and faster sMX
     op.bit == 0 && return frame
     i = op.qubit
-    xzs = frame.frame.tab.xzs
-    T = eltype(xzs)
-    lowbit = T(1)
-    ibig = _div(T,i-1)+1
-    ismall = _mod(T,i-1)
-    ismallm = lowbit<<(ismall)
+    xzs = tab(frame.frame).xzs
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
 
     @inbounds @simd for f in eachindex(frame)
         should_flip = !iszero(xzs[ibig,f] & ismallm)
@@ -99,12 +95,8 @@ end
 
 function apply!(frame::PauliFrame, op::sMRZ) # TODO sMRY, and faster sMRX
     i = op.qubit
-    xzs = frame.frame.tab.xzs
-    T = eltype(xzs)
-    lowbit = T(1)
-    ibig = _div(T,i-1)+1
-    ismall = _mod(T,i-1)
-    ismallm = lowbit<<(ismall)
+    xzs = tab(frame.frame).xzs
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
 
     if op.bit != 0
         @inbounds @simd for f in eachindex(frame)
@@ -122,45 +114,39 @@ end
 
 function applynoise!(frame::PauliFrame,noise::UnbiasedUncorrelatedNoise,i::Int)
     p = noise.p
-    T = eltype(frame.frame.tab.xzs)
+    xzs = tab(frame.frame).xzs
 
-    lowbit = T(1)
-    ibig = _div(T,i-1)+1
-    ismall = _mod(T,i-1)
-    ismallm = lowbit<<(ismall)
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
     p = p/3
 
     @inbounds @simd for f in eachindex(frame)
         r = rand()
         if  r < p # X error
-            frame.frame.tab.xzs[ibig,f] ⊻= ismallm
+            xzs[ibig,f] ⊻= ismallm
         elseif r < 2p # Z error
-            frame.frame.tab.xzs[end÷2+ibig,f] ⊻= ismallm
+            xzs[end÷2+ibig,f] ⊻= ismallm
         elseif r < 3p # Y error
-            frame.frame.tab.xzs[ibig,f] ⊻= ismallm
-            frame.frame.tab.xzs[end÷2+ibig,f] ⊻= ismallm
+            xzs[ibig,f] ⊻= ismallm
+            xzs[end÷2+ibig,f] ⊻= ismallm
         end
     end
     return frame
 end
 
 function applynoise!(frame::PauliFrame,noise::PauliNoise,i::Int)
-    T = eltype(frame.frame.tab.xzs)
+    xzs = tab(frame.frame).xzs
 
-    lowbit = T(1)
-    ibig = _div(T,i-1)+1
-    ismall = _mod(T,i-1)
-    ismallm = lowbit<<(ismall)
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
 
     @inbounds @simd for f in eachindex(frame)
         r = rand()
         if  r < noise.px # X error
-            frame.frame.tab.xzs[ibig,f] ⊻= ismallm
+            xzs[ibig,f] ⊻= ismallm
         elseif r < noise.px+noise.pz # Z error
-            frame.frame.tab.xzs[end÷2+ibig,f] ⊻= ismallm
+            xzs[end÷2+ibig,f] ⊻= ismallm
         elseif r < noise.px+noise.pz+noise.py # Y error
-            frame.frame.tab.xzs[ibig,f] ⊻= ismallm
-            frame.frame.tab.xzs[end÷2+ibig,f] ⊻= ismallm
+            xzs[ibig,f] ⊻= ismallm
+            xzs[end÷2+ibig,f] ⊻= ismallm
         end
     end
     return frame
diff --git a/src/pauli_operator.jl b/src/pauli_operator.jl
@@ -87,19 +87,23 @@ macro P_str(a)
     quote _P_str($a) end
 end
 
-Base.getindex(p::PauliOperator{Tₚ,Tᵥ}, i::Int) where {Tₚ, Tᵥₑ<:Unsigned, Tᵥ<:AbstractVector{Tᵥₑ}} = ((p.xz[_div(Tᵥₑ, i-1)+1] & Tᵥₑ(0x1)<<_mod(Tᵥₑ,i-1))!=0x0)::Bool, ((p.xz[end÷2+_div(Tᵥₑ,i-1)+1] & Tᵥₑ(0x1)<<_mod(Tᵥₑ,i-1))!=0x0)::Bool
+function Base.getindex(p::PauliOperator{Tₚ,Tᵥ}, i::Int) where {Tₚ, Tᵥₑ<:Unsigned, Tᵥ<:AbstractVector{Tᵥₑ}}
+    _, ibig, _, ismallm = get_bitmask_idxs(p.xz,i)
+    ((p.xz[ibig] & ismallm) != 0x0)::Bool, ((p.xz[end÷2+ibig] & ismallm) != 0x0)::Bool
+end
 Base.getindex(p::PauliOperator{Tₚ,Tᵥ}, r) where {Tₚ, Tᵥₑ<:Unsigned, Tᵥ<:AbstractVector{Tᵥₑ}} = PauliOperator(p.phase[], xbit(p)[r], zbit(p)[r])
 
 function Base.setindex!(p::PauliOperator{Tₚ,Tᵥ}, (x,z)::Tuple{Bool,Bool}, i) where {Tₚ, Tᵥₑ, Tᵥ<:AbstractVector{Tᵥₑ}}
+    _, ibig, _, ismallm = get_bitmask_idxs(p.xz,i)
     if x
-        p.xz[_div(Tᵥₑ,i-1)+1] |= Tᵥₑ(0x1)<<_mod(Tᵥₑ,i-1)
+        p.xz[ibig] |= ismallm
     else
-        p.xz[_div(Tᵥₑ,i-1)+1] &= ~(Tᵥₑ(0x1)<<_mod(Tᵥₑ,i-1))
+        p.xz[ibig] &= ~(ismallm)
     end
     if z
-        p.xz[end÷2+_div(Tᵥₑ,i-1)+1] |= Tᵥₑ(0x1)<<_mod(Tᵥₑ,i-1)
+        p.xz[end÷2+ibig] |= ismallm
     else
-        p.xz[end÷2+_div(Tᵥₑ,i-1)+1] &= ~(Tᵥₑ(0x1)<<_mod(Tᵥₑ,i-1))
+        p.xz[end÷2+ibig] &= ~(ismallm)
     end
     p
 end
diff --git a/src/project_trace_reset.jl b/src/project_trace_reset.jl
@@ -42,17 +42,15 @@ function _generate!(pauli::PauliOperator{Tₚ,Tᵥ}, stabilizer::Stabilizer{Tabl
     xzs = tab(stabilizer).xzs
     xs = @view xzs[1:end÷2,:]
     zs = @view xzs[end÷2+1:end,:]
-    lowbit = Tₘₑ(0x1)
     zerobit = Tₘₑ(0x0)
     px,pz = xview(pauli), zview(pauli)
     used_indices = Int[]
     used = 0
     # remove Xs
     while (i=unsafe_bitfindnext_(px,1); i !== nothing) # TODO awkward notation due to https://github.com/JuliaLang/julia/issues/45499
-        jbig = _div(Tₘₑ,i-1)+1
-        jsmall = lowbit<<_mod(Tₘₑ,i-1)
-        candidate = findfirst(e->e&jsmall!=zerobit, # TODO some form of reinterpret might be faster than equality check
-                              xs[jbig,used+1:end])
+        _, ibig, _, ismallm = get_bitmask_idxs(xzs,i)
+        candidate = findfirst(e->e&ismallm!=zerobit, # TODO some form of reinterpret might be faster than equality check
+                              xs[ibig,used+1:end])
         if isnothing(candidate)
             return nothing
         else
@@ -63,10 +61,9 @@ function _generate!(pauli::PauliOperator{Tₚ,Tᵥ}, stabilizer::Stabilizer{Tabl
     end
     # remove Zs
     while (i=unsafe_bitfindnext_(pz,1); i !== nothing) # TODO awkward notation due to https://github.com/JuliaLang/julia/issues/45499
-        jbig = _div(Tₘₑ,i-1)+1
-        jsmall = lowbit<<_mod(Tₘₑ,i-1)
-        candidate = findfirst(e->e&jsmall!=zerobit, # TODO some form of reinterpret might be faster than equality check
-                              zs[jbig,used+1:end])
+        _, ibig, _, ismallm = get_bitmask_idxs(xzs,i)
+        candidate = findfirst(e->e&ismallm!=zerobit, # TODO some form of reinterpret might be faster than equality check
+                              zs[ibig,used+1:end])
         if isnothing(candidate)
             return nothing
         else
diff --git a/src/symbolic_cliffords.jl b/src/symbolic_cliffords.jl
@@ -415,12 +415,9 @@ LinearAlgebra.inv(op::sInvSQRTZZ) = sSQRTZZ(op.q1, op.q2)
 """Apply a Pauli Z to the `i`-th qubit of state `s`. You should use `apply!(stab,sZ(i))` instead of this."""
 function apply_single_z!(stab::AbstractStabilizer, i)
     s = tab(stab)
-    Tₘₑ = eltype(s.xzs)
-    bigi = _div(Tₘₑ,i-1)+1
-    smalli = _mod(Tₘₑ,i-1)
-    mask = Tₘₑ(0x1)<<smalli
+    _, ibig, _, ismallm = get_bitmask_idxs(s.xzs,i)
     @inbounds @simd for row in 1:size(s.xzs,2)
-        if !iszero(s.xzs[bigi,row] & mask)
+        if !iszero(s.xzs[ibig,row] & ismallm)
             s.phases[row] = (s.phases[row]+0x2)&0x3
         end
     end
@@ -430,12 +427,9 @@ end
 """Apply a Pauli X to the `i`-th qubit of state `s`. You should use `apply!(stab,sX(i))` instead of this."""
 function apply_single_x!(stab::AbstractStabilizer, i)
     s = tab(stab)
-    Tₘₑ = eltype(s.xzs)
-    bigi = _div(Tₘₑ,i-1)+1
-    smalli = _mod(Tₘₑ,i-1)
-    mask = Tₘₑ(0x1)<<smalli
+    _, ibig, _, ismallm = get_bitmask_idxs(s.xzs,i)
     @inbounds @simd for row in 1:size(s.xzs,2)
-        if !iszero(s.xzs[end÷2+bigi,row] & mask)
+        if !iszero(s.xzs[end÷2+ibig,row] & ismallm)
             s.phases[row] = (s.phases[row]+0x2)&0x3
         end
     end
@@ -445,12 +439,9 @@ end
 """Apply a Pauli Y to the `i`-th qubit of state `s`. You should use `apply!(stab,sY(i))` instead of this."""
 function apply_single_y!(stab::AbstractStabilizer, i)
     s = tab(stab)
-    Tₘₑ = eltype(s.xzs)
-    bigi = _div(Tₘₑ,i-1)+1
-    smalli = _mod(Tₘₑ,i-1)
-    mask = Tₘₑ(0x1)<<smalli
+    _, ibig, _, ismallm = get_bitmask_idxs(s.xzs,i)
     @inbounds @simd for row in 1:size(s.xzs,2)
-        if !iszero((s.xzs[bigi,row] & mask) ⊻ (s.xzs[end÷2+bigi,row] & mask))
+        if !iszero((s.xzs[ibig,row] & ismallm) ⊻ (s.xzs[end÷2+ibig,row] & ismallm))
             s.phases[row] = (s.phases[row]+0x2)&0x3
         end
     end
