add codereview suggestions

QuantumSavory · Sep 28, 2024 · 18bcec7 · 18bcec7
1 parent c37cc77
commit 18bcec7
Show file tree

Hide file tree

Showing 7 changed files with 71 additions and 63 deletions.
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}
+    xzs = frame.frame.tab.xzs
     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
+    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
@@ -26,7 +26,7 @@ export
     nqubits,
     stabilizerview, destabilizerview, logicalxview, logicalzview, phases,
     fastcolumn, fastrow,
-    bitview, quantumstate, tab,
+    bitview, quantumstate, tab, get_bitmask_idxs,
     BadDataStructure,
     affectedqubits, #TODO move to QuantumInterface?
     # GF2
@@ -910,6 +910,34 @@ 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.
+
+Internal bitwise operations `_div`, `_mod`, and `_mask` are
+used to compute the word and bit positions based on the bit
+size of the type `T`.
+"""
+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
@@ -79,22 +79,14 @@ function apply!(frame::PauliFrame, op::sMRX) # TODO implement a faster direct ve
     apply!(frame, sMRZ(op.qubit, op.bit))
 end
 
-function _get_bitmask_idxs(frame::PauliFrame, i::Int)
-    T = eltype(frame.frame.tab.xzs)
-    lowbit = T(1)
-    ibig = _div(T, i-1) + 1
-    ismall = _mod(T, i-1)
-    ismallm = lowbit << ismall
-    return ibig, ismall, ismallm
-end
-
 function apply!(frame::PauliFrame, op::sMZ) # TODO sMY, and faster sMX
     op.bit == 0 && return frame
     i = op.qubit
-    ibig, ismall, ismallm = _get_bitmask_idxs(frame,i)
+    xzs = frame.frame.tab.xzs
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
 
     @inbounds @simd for f in eachindex(frame)
-        should_flip = !iszero(frame.frame.tab.xzs[ibig,f] & ismallm)
+        should_flip = !iszero(xzs[ibig,f] & ismallm)
         frame.measurements[f,op.bit] = should_flip
     end
 
@@ -103,25 +95,28 @@ end
 
 function apply!(frame::PauliFrame, op::sMRZ) # TODO sMRY, and faster sMRX
     i = op.qubit
-    ibig, ismall, ismallm = _get_bitmask_idxs(frame,i)
+    xzs = frame.frame.tab.xzs
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
 
     if op.bit != 0
         @inbounds @simd for f in eachindex(frame)
-            should_flip = !iszero(frame.frame.tab.xzs[ibig,f] & ismallm)
+            should_flip = !iszero(xzs[ibig,f] & ismallm)
             frame.measurements[f,op.bit] = should_flip
         end
     end
     @inbounds @simd for f in eachindex(frame)
-        frame.frame.tab.xzs[ibig,f] &= ~ismallm
-        rand(Bool) && (frame.frame.tab.xzs[end÷2+ibig,f] ⊻= ismallm)
+        xzs[ibig,f] &= ~ismallm
+        rand(Bool) && (xzs[end÷2+ibig,f] ⊻= ismallm)
     end
 
     return frame
 end
 
 function applynoise!(frame::PauliFrame,noise::UnbiasedUncorrelatedNoise,i::Int)
     p = noise.p
-    ibig, ismall, ismallm = _get_bitmask_idxs(frame,i)
+    xzs = frame.frame.tab.xzs
+
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
     p = p/3
 
     @inbounds @simd for f in eachindex(frame)
@@ -139,7 +134,9 @@ function applynoise!(frame::PauliFrame,noise::UnbiasedUncorrelatedNoise,i::Int)
 end
 
 function applynoise!(frame::PauliFrame,noise::PauliNoise,i::Int)
-    ibig, ismall, ismallm = _get_bitmask_idxs(frame,i)
+    xzs = frame.frame.tab.xzs
+
+    lowbit, ibig, ismall, ismallm = get_bitmask_idxs(xzs,i)
 
     @inbounds @simd for f in eachindex(frame)
         r = rand()

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, (p.xz[end÷2+_div(Tᵥₑ,i-1)+1] & Tᵥₑ(0x1)<<_mod(Tᵥₑ,i-1))!=0x0
+Base.getindex(p::PauliOperator{Tₚ,Tᵥ}, i::Int) where {Tₚ, Tᵥₑ<:Unsigned, Tᵥ<:AbstractVector{Tᵥₑ}} = begin
+    _, ibig, _, ismallm = get_bitmask_idxs(p.xz,i)
+    (p.xz[ibig] & ismallm) != 0x0, (p.xz[end÷2+ibig] & ismallm) != 0x0
+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
@@ -394,12 +394,9 @@ LinearAlgebra.inv(op::sInvZCrY) = sZCrY(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
@@ -409,12 +406,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
@@ -424,12 +418,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