fix: Improvements to precompilation and adjustments to new Quasar api (#61)

* change: Updates for latest Quasar parsing

* fix: More cleanup, improved precompilation

Author: kshyatt-aws

Project.toml

@@ -1,7 +1,7 @@
 name = "BraketSimulator"
 uuid = "76d27892-9a0b-406c-98e4-7c178e9b3dff"
 authors = ["Katharine Hyatt <hyatkath@amazon.com> and contributors"]
-version = "0.0.6"
+version = "0.0.7"
 
 [deps]
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
@@ -39,7 +39,7 @@ Logging = "1.6"
 OrderedCollections = "=1.6.3"
 PrecompileTools = "=1.2.1"
 PythonCall = "=0.9.23"
-Quasar = "0.0.1"
+Quasar = "0.0.2"
 Random = "1.6"
 StaticArrays = "1.9"
 StatsBase = "0.34"

src/BraketSimulator.jl

@@ -78,9 +78,11 @@ include("noise_kernels.jl")
 
 function __init__()
     Quasar.builtin_gates[] = builtin_gates
+    Quasar.parse_pragma[]  = parse_pragma
+    Quasar.visit_pragma[]  = visit_pragma
 end
 
-const LOG2_CHUNK_SIZE = 10
+const LOG2_CHUNK_SIZE = 16
 const CHUNK_SIZE      = 2^LOG2_CHUNK_SIZE
 
 function _index_to_endian_bits(ix::Int, qubit_count::Int)
@@ -203,7 +205,7 @@ basis rotation instructions if running with non-zero shots. Return the `Program`
 parsing and the qubit count of the circuit.
 """
 function _prepare_program(circuit_ir::OpenQasmProgram, inputs::Dict{String, <:Any}, shots::Int)
-    ir_inputs = isnothing(circuit_ir.inputs) ? Dict{String, Float64}() : circuit_ir.inputs 
+    ir_inputs     = isnothing(circuit_ir.inputs) ? Dict{String, Float64}() : circuit_ir.inputs
     merged_inputs = merge(ir_inputs, inputs)
     src           = circuit_ir.source::String
     circuit       = to_circuit(src, merged_inputs)
@@ -704,11 +706,14 @@ include("dm_simulator.jl")
         bit[3] b;
         qubit[3] q;
         rx(0.1) q[0];
+        rx(1) q[0];
         prx(0.1, 0.2) q[0];
         x q[0];
         ry(0.1) q[0];
+        ry(1) q[0];
         y q[0];
         rz(0.1) q[0];
+        rz(1) q[0];
         z q[0];
         h q[0];
         i q[0];
@@ -758,6 +763,11 @@ include("dm_simulator.jl")
         #pragma braket result density_matrix
         #pragma braket result probability
         #pragma braket result expectation x(q[0])
+        #pragma braket result expectation x(q[0]) @ x(q[1])
+        #pragma braket result expectation z(q[0]) @ z(q[1])
+        #pragma braket result expectation y(q[0]) @ y(q[1])
+        #pragma braket result expectation h(q[0]) @ h(q[1])
+        #pragma braket result expectation i(q[0]) @ i(q[1])
         #pragma braket result variance x(q[0]) @ y(q[1])
         """
     dm_exact_results_qasm = """
@@ -772,16 +782,22 @@ include("dm_simulator.jl")
         """
     shots_results_qasm = """
         OPENQASM 3.0;
-        qubit[2] q;
+        qubit[10] q;
         h q;
         #pragma braket result probability
         #pragma braket result expectation x(q[0])
         #pragma braket result variance x(q[0]) @ y(q[1])
         #pragma braket result sample x(q[0]) @ y(q[1])
+        #pragma braket result expectation z(q[2]) @ z(q[3])
+        #pragma braket result expectation x(q[4]) @ x(q[5])
+        #pragma braket result expectation y(q[6]) @ y(q[7])
+        #pragma braket result expectation h(q[8]) @ h(q[9])
         """
     @compile_workload begin
         using BraketSimulator, Quasar
         Quasar.builtin_gates[] = BraketSimulator.builtin_gates
+        Quasar.parse_pragma[]  = BraketSimulator.parse_pragma
+        Quasar.visit_pragma[]  = BraketSimulator.visit_pragma
         simulator = StateVectorSimulator(5, 0)
         oq3_program = OpenQasmProgram(braketSchemaHeader("braket.ir.openqasm.program", "1"), custom_qasm, nothing)
         simulate(simulator, oq3_program, 100)

src/circuit.jl

@@ -169,11 +169,12 @@ function basis_rotation_instructions!(c::Circuit)
         c.basis_rotation_instructions = reduce(vcat, _observable_to_instruction(all_qubit_observable, target) for target in qubits(c))
         return c
     end
-    unsorted = collect(Set(values(c.qubit_observable_target_mapping)))
-    target_lists = sort(unsorted)
+    mapping_vals = collect(values(c.qubit_observable_target_mapping))
+    target_lists = unique(mapping_vals)
     for target_list in target_lists
-        observable = c.qubit_observable_mapping[first(target_list)]
-        append!(basis_rotation_instructions, _observable_to_instruction(observable, target_list))
+        observable    = c.qubit_observable_mapping[first(target_list)]
+        observable_ix = _observable_to_instruction(observable, target_list)
+        append!(basis_rotation_instructions, observable_ix)
     end
     c.basis_rotation_instructions = basis_rotation_instructions
     return c

src/gate_kernels.jl

@@ -96,13 +96,14 @@ function matrix_rep(g::PRx)
 end
 
 for G in (:Rx, :Ry, :Rz, :PhaseShift)
-    @eval function matrix_rep(g::$G)
+    @eval function matrix_rep(g::$G)::SMatrix{2,2,ComplexF64}
         n = g.pow_exponent::Float64
         θ = @inbounds g.angle[1]
-        iszero(n)    && return matrix_rep_raw(I())::SMatrix{2,2,ComplexF64}
-        isone(n)     && return matrix_rep_raw(g, θ)::SMatrix{2,2,ComplexF64}
-        isinteger(n) && return matrix_rep_raw(g, θ*n)::SMatrix{2,2,ComplexF64}
-        return SMatrix{2,2,ComplexF64}(matrix_rep_raw(g, θ) ^ n)
+        one_mat = matrix_rep_raw(g, θ)
+        iszero(n)    && return matrix_rep_raw(I())
+        isone(n)     && return one_mat
+        isinteger(n) && return matrix_rep_raw(g, θ*n)
+        return SMatrix{2,2,ComplexF64}(one_mat ^ n)
     end
 end
 
@@ -211,9 +212,9 @@ matrix_rep_raw(g::PRx) = SMatrix{2,2}(
         -im*exp(im*g.angle[2])*sin(g.angle[1]/2.0) cos(g.angle[1] / 2.0)
     ],
 )
-matrix_rep_raw(g::Rz, ϕ) = (θ = ϕ/2.0; return SMatrix{2,2}(exp(-im*θ), 0.0, 0.0, exp(im*θ)))
-matrix_rep_raw(g::Rx, ϕ) = ((sθ, cθ) = sincos(ϕ/2.0); return SMatrix{2,2}(cθ, -im*sθ, -im*sθ, cθ))
-matrix_rep_raw(g::Ry, ϕ) = ((sθ, cθ) = sincos(ϕ/2.0); return SMatrix{2,2}(complex(cθ), complex(sθ), -complex(sθ), complex(cθ)))
+matrix_rep_raw(g::Rz, ϕ)::SMatrix{2,2,ComplexF64} = ((sθ, cθ) = sincos(ϕ/2.0); return SMatrix{2,2}(cθ - im*sθ, 0.0, 0.0, cθ + im*sθ))
+matrix_rep_raw(g::Rx, ϕ)::SMatrix{2,2,ComplexF64} = ((sθ, cθ) = sincos(ϕ/2.0); return SMatrix{2,2}(cθ, -im*sθ, -im*sθ, cθ))
+matrix_rep_raw(g::Ry, ϕ)::SMatrix{2,2,ComplexF64} = ((sθ, cθ) = sincos(ϕ/2.0); return SMatrix{2,2}(complex(cθ), complex(sθ), -complex(sθ), complex(cθ)))
 matrix_rep_raw(g::GPi) =
     SMatrix{2,2}(complex([0 exp(-im * g.angle[1]); exp(im * g.angle[1]) 0]))
 
@@ -311,14 +312,11 @@ function apply_gate!(
     g_00, g_10, g_01, g_11 = g_matrix
     Threads.@threads for chunk_index = 0:n_chunks-1
         # my_amps is the group of amplitude generators which this `Task` will process
-        my_amps = if n_chunks > 1
-            chunk_index*CHUNK_SIZE:((chunk_index+1)*CHUNK_SIZE-1)
-        else
-            0:n_tasks-1
-        end
-        lower_ix  = pad_bit(my_amps[1], endian_qubit) + 1
+        first_amp = n_chunks > 1 ? chunk_index*CHUNK_SIZE : 0
+        amp_block = n_chunks > 1 ? CHUNK_SIZE : n_tasks
+        lower_ix  = pad_bit(first_amp, endian_qubit) + 1
         higher_ix = lower_ix + flipper
-        for task_amp = 0:length(my_amps)-1
+        for task_amp = 0:amp_block-1
             if is_small_target && div(task_amp, flipper) > 0 && mod(task_amp, flipper) == 0
                 lower_ix  = higher_ix
                 higher_ix = lower_ix + flipper

src/gates.jl

@@ -112,7 +112,7 @@ mutable struct Unitary <: Gate
     Unitary(matrix::Matrix{<:Number}, pow_exponent=1.0) = new(ComplexF64.(matrix), Float64(pow_exponent))
 end
 Base.:(==)(u1::Unitary, u2::Unitary) = u1.matrix == u2.matrix && u1.pow_exponent == u2.pow_exponent
-qubit_count(g::Unitary) = convert(Int, log2(size(g.matrix, 1)))
+qubit_count(g::Unitary) = qubit_count(g.matrix)
 StructTypes.constructfrom(::Type{Unitary}, nt::Quasar.CircuitInstruction) = Unitary(only(nt.arguments), nt.exponent)
 
 Parametrizable(g::AngledGate) = Parametrized()

src/observables.jl

@@ -71,7 +71,7 @@ HermitianObservable(v::Vector{Vector{Vector{T}}}) where {T<:Number} = HermitianO
 Base.copy(o::HermitianObservable) = HermitianObservable(copy(o.matrix))
 StructTypes.lower(x::HermitianObservable) = Union{String, Vector{Vector{Vector{Float64}}}}[complex_matrix_to_ir(ComplexF64.(x.matrix))]
 Base.:(==)(h1::HermitianObservable, h2::HermitianObservable) = (size(h1.matrix) == size(h2.matrix) && h1.matrix ≈ h2.matrix)
-qubit_count(o::HermitianObservable) = convert(Int, log2(size(o.matrix, 1)))
+qubit_count(o::HermitianObservable) = qubit_count(o.matrix) 
 LinearAlgebra.eigvals(o::HermitianObservable) = eigvals(Hermitian(o.matrix))
 unscaled(o::HermitianObservable) = o
 Base.:(*)(o::HermitianObservable, n::Real) = HermitianObservable(Float64(n) .* o.matrix)

src/operators.jl

@@ -21,6 +21,8 @@ StructTypes.StructType(::Type{QuantumOperator}) = StructTypes.AbstractType()
 StructTypes.subtypes(::Type{QuantumOperator}) = (h=H, i=I, x=X, y=Y, z=Z, s=S, si=Si, t=T, ti=Ti, v=V, vi=Vi, cnot=CNot, swap=Swap, iswap=ISwap, cv=CV, cy=CY, cz=CZ, ecr=ECR, ccnot=CCNot, cswap=CSwap, unitary=Unitary, rx=Rx, ry=Ry, rz=Rz, phaseshift=PhaseShift, pswap=PSwap, xy=XY, cphaseshift=CPhaseShift, cphaseshift00=CPhaseShift00, cphaseshift01=CPhaseShift01, cphaseshift10=CPhaseShift10, xx=XX, yy=YY, zz=ZZ, gpi=GPi, gpi2=GPi2, ms=MS, prx=PRx, u=U, gphase=GPhase, kraus=Kraus, bit_flip=BitFlip, phase_flip=PhaseFlip, pauli_channel=PauliChannel, amplitude_damping=AmplitudeDamping, phase_damping=PhaseDamping, depolarizing=Depolarizing, two_qubit_dephasing=TwoQubitDephasing, two_qubit_depolarizing=TwoQubitDepolarizing, generalized_amplitude_damping=GeneralizedAmplitudeDamping, multi_qubit_pauli_channel=MultiQubitPauliChannel, measure=Measure, reset=Reset, barrier=Barrier, delay=Delay)
 parameters(::QuantumOperator) = FreeParameter[]
 
+qubit_count(o::Matrix) = Int(log2(size(o, 1))) 
+
 struct PauliEigenvalues{N}
     coeff::Float64
     PauliEigenvalues{N}(coeff::Float64=1.0) where {N} = new(coeff)
@@ -99,4 +101,4 @@ for T in (:Barrier, :Reset, :Delay, :Measure)
         qubit_count(o::$T) = qubit_count($T)
         Parametrizable(::$T) = NonParametrized()
     end
-end
+end

src/pragmas.jl

@@ -1,5 +1,4 @@
-Quasar.qubit_count(o::String) = length(o)
-Quasar.qubit_count(o::Matrix) = Int(log2(size(o, 1))) 
+qubit_count(o::String) = length(o)
 
 function _observable_targets_error(observable::Matrix{ComplexF64}, targets)
     mat = Vector{Vector{Vector{Float64}}}(undef, size(observable, 1))
@@ -14,7 +13,7 @@ end
 _observable_targets_error(::String, targets) = throw(Quasar.QasmVisitorError("Standard observable target must be exactly 1 qubit.", "ValueError"))
 
 function _check_observable_targets(observable::Union{Matrix{ComplexF64}, String}, targets)
-    qc = Quasar.qubit_count(observable)
+    qc = qubit_count(observable)
     qc == 1 && (isempty(targets) || length(targets) == 1) && return
     qc == length(targets) && return
     _observable_targets_error(observable, targets)
@@ -42,7 +41,7 @@ function visit_observable(v, expr)
     end
 end
 
-function Quasar.visit_pragma(v, program_expr)
+function visit_pragma(v, program_expr)
     pragma_type::Symbol = program_expr.args[1]
     if pragma_type == :result
         result_type = program_expr.args[2]
@@ -95,12 +94,12 @@ function Quasar.visit_pragma(v, program_expr)
 end
 
 function parse_matrix(tokens::Vector{Tuple{Int64, Int32, Quasar.Token}}, stack, start, qasm)
-    inner = Quasar.extract_braced_block(tokens, stack, start, qasm)
+    inner = Quasar.extract_expression(tokens, Quasar.lbracket, Quasar.rbracket, stack, start, qasm)
     n_rows = count(triplet->triplet[end] == Quasar.lbracket, inner)
     matrix = Matrix{Quasar.QasmExpression}(undef, n_rows, n_rows)
     row = 1
     while !isempty(inner)
-        row_tokens = Quasar.extract_braced_block(inner, stack, start, qasm)
+        row_tokens = Quasar.extract_expression(inner, Quasar.lbracket, Quasar.rbracket, stack, start, qasm)
         push!(row_tokens, (-1, Int32(-1), Quasar.semicolon))
         col = 1
         while !isempty(row_tokens)
@@ -126,7 +125,7 @@ function parse_pragma_observables(tokens::Vector{Tuple{Int64, Int32, Quasar.Toke
         observable_token = popfirst!(tokens)
         observable_id    = Quasar.parse_identifier(observable_token, qasm)
         if observable_id.args[1] == "hermitian"
-            matrix_tokens = Quasar.extract_parensed(tokens, stack, start, qasm)
+            matrix_tokens = Quasar.extract_expression(tokens, Quasar.lparen, Quasar.rparen, stack, start, qasm)
             # next token is targets
             h_mat = parse_matrix(matrix_tokens, stack, start, qasm)
             # next token is targets
@@ -146,7 +145,7 @@ function parse_pragma_observables(tokens::Vector{Tuple{Int64, Int32, Quasar.Toke
             break
         else
             if !isempty(tokens) && first(tokens)[end] == Quasar.lparen
-                arg_tokens = Quasar.extract_parensed(tokens, stack, start, qasm)
+                arg_tokens = Quasar.extract_expression(tokens, Quasar.lparen, Quasar.rparen, stack, start, qasm)
                 push!(arg_tokens, (-1, Int32(-1), Quasar.semicolon))
                 target_expr = Quasar.parse_expression(arg_tokens, stack, start, qasm)
                 push!(obs_targets, target_expr)
@@ -175,7 +174,7 @@ function parse_pragma_targets(tokens::Vector{Tuple{Int64, Int32, Quasar.Token}},
 end
 
 
-function Quasar.parse_pragma(tokens, stack, start, qasm)
+function parse_pragma(tokens, stack, start, qasm)
     prefix    = popfirst!(tokens)
     prefix_id = Quasar.parse_identifier(prefix, qasm)
     prefix_id.args[1] == "braket" || throw(Quasar.QasmParseError("pragma expression must begin with `#pragma braket`", stack, start, qasm))
@@ -204,7 +203,7 @@ function Quasar.parse_pragma(tokens, stack, start, qasm)
         end
     elseif pragma_type == "unitary"
         push!(expr, :unitary)
-        matrix_tokens  = Quasar.extract_parensed(tokens, stack, start, qasm)
+        matrix_tokens  = Quasar.extract_expression(tokens, Quasar.lparen, Quasar.rparen, stack, start, qasm)
         unitary_matrix = parse_matrix(matrix_tokens, stack, start, qasm)
         push!(expr, unitary_matrix)
         target_expr = parse_pragma_targets(tokens, stack, start, qasm)
@@ -213,8 +212,8 @@ function Quasar.parse_pragma(tokens, stack, start, qasm)
         push!(expr, :noise)
         noise_type = Quasar.parse_identifier(popfirst!(tokens), qasm)::Quasar.QasmExpression
         if noise_type.args[1] == "kraus"
-            matrix_tokens = Quasar.extract_parensed(tokens, stack, start, qasm)
-            all(triplet->triplet[end] == Quasar.lbracket, matrix_tokens[1:3]) && (matrix_tokens = Quasar.extract_braced_block(matrix_tokens, stack, start, qasm))
+            matrix_tokens = Quasar.extract_expression(tokens, Quasar.lparen, Quasar.rparen, stack, start, qasm)
+            all(triplet->triplet[end] == Quasar.lbracket, matrix_tokens[1:3]) && (matrix_tokens = Quasar.extract_expression(matrix_tokens, Quasar.lbracket, Quasar.rbracket, stack, start, qasm))
             mats = Matrix{Quasar.QasmExpression}[]
             while !isempty(matrix_tokens)
                 push!(mats, parse_matrix(matrix_tokens, stack, start, qasm))

src/validation.jl

@@ -164,6 +164,7 @@ function _check_observable(observable_map, observable, qubits)
     observable_map[qubits] = observable
     return observable_map
 end
+_check_observable(observable_map, observable, qubits::Int) = _check_observable(observable_map, observable, [qubits])
 
 function _combine_obs_and_targets(observable::Observables.HermitianObservable, result_targets::Vector{Int})
     obs_qc = qubit_count(observable)
@@ -175,12 +176,14 @@ _combine_obs_and_targets(observable::Observables.TensorProduct, result_targets::
 _combine_obs_and_targets(observable, result_targets::Vector{Int}) = length(result_targets) == 1 ? [(observable, result_targets)] : [(copy(observable), t) for t in result_targets]
 
 function _verify_openqasm_shots_observables(circuit::Circuit, n_qubits::Int)
-    observable_map = Dict()
+    observable_map = LittleDict{Vector{Int}, Observables.Observable}()
     for result in filter(rt->rt isa ObservableResult, circuit.result_types)
         result.observable isa Observables.I && continue
         result_targets = isempty(result.targets) ? collect(0:n_qubits-1) : collect(result.targets)
         for obs_and_target in _combine_obs_and_targets(result.observable, result_targets)
-            observable_map = _check_observable(observable_map, obs_and_target...)
+            obs  = obs_and_target[1]
+            targ = obs_and_target[2]
+            observable_map = _check_observable(observable_map, obs, targ)
         end
     end
     return

test/runtests.jl

@@ -1,6 +1,6 @@
 using Test, Aqua, Documenter, BraketSimulator
 
-Aqua.test_all(BraketSimulator, ambiguities=false, piracies=false)
+Aqua.test_all(BraketSimulator, ambiguities=false)
 Aqua.test_ambiguities(BraketSimulator)
 dir_list = filter(x-> startswith(x, "test_") && endswith(x, ".jl"), readdir(@__DIR__))
 

test/test_openqasm3.jl

@@ -762,9 +762,7 @@ get_tol(shots::Int) = return (
             qubit[3] q;
             i q;
             #pragma braket result expectation x(q[2])
-            // # noqa: E501
             #pragma braket result expectation hermitian([[-6+0im, 2+1im, -3+0im, -5+2im], [2-1im, 0im, 2-1im, -5+4im], [-3+0im, 2+1im, 0im, -4+3im], [-5-2im, -5-4im, -4-3im, -6+0im]]) q[0:1]
-            // # noqa: E501
             #pragma braket result expectation x(q[2]) @ hermitian([[-6+0im, 2+1im, -3+0im, -5+2im], [2-1im, 0im, 2-1im, -5+4im], [-3+0im, 2+1im, 0im, -4+3im], [-5-2im, -5-4im, -4-3im, -6+0im]]) q[0:1]
             """
             circuit = BraketSimulator.to_circuit(qasm) 
@@ -773,7 +771,8 @@ get_tol(shots::Int) = return (
                     2-1im 0 2-1im -5+4im;
                    -3 2+1im 0 -4+3im;
                    -5-2im -5-4im -4-3im -6]
-            h = BraketSimulator.Observables.HermitianObservable(arr)
+            h    = BraketSimulator.Observables.HermitianObservable(arr)
+            @test circuit.result_types[2].observable.matrix == arr
             bris = vcat(BraketSimulator.diagonalizing_gates(h, [0, 1]), BraketSimulator.Instruction(BraketSimulator.H(), [2]))
             for (ix, bix) in zip(circuit.basis_rotation_instructions, bris)
                 @test Matrix(BraketSimulator.matrix_rep(ix.operator)) ≈ adjoint(BraketSimulator.fix_endianness(Matrix(BraketSimulator.matrix_rep(bix.operator))))