Unverified Commit cec804e2 authored by Mccaskey, Alex's avatar Mccaskey, Alex Committed by GitHub
Browse files

Merge pull request #45 from tnguyen-ornl/tnguyen/pyxasm

Work on PyXASM syntax handler
parents 0083c54d c0d98372
Pipeline #123674 passed with stage
in 18 minutes and 50 seconds
......@@ -55,9 +55,8 @@ void PyXasmTokenCollector::collect(clang::Preprocessor &PP,
std::vector<std::pair<std::string, int>> lines;
std::string line = "";
auto current_line_number = sm.getSpellingLineNumber(Toks[0].getLocation());
line += PP.getSpelling(Toks[0]);
int last_col_number = 0;
for (int i = 1; i < Toks.size(); i++) {
for (int i = 0; i < Toks.size(); i++) {
// std::cout << PP.getSpelling(Toks[i]) << "\n";
auto location = Toks[i].getLocation();
auto col_number = sm.getSpellingColumnNumber(location);
......@@ -90,17 +89,29 @@ void PyXasmTokenCollector::collect(clang::Preprocessor &PP,
using namespace antlr4;
int previous_col = lines[0].second;
bool is_in_for_loop = false;
int line_counter = 0;
// Tracking the scope of for loops by their indent
std::stack<int> for_loop_indent;
for (const auto &line : lines) {
// std::cout << "processing line " << line_counter << " of " << lines.size()
// << ": " << line.first << ", " << line.second << std::boolalpha
// << ", " << is_in_for_loop << "\n";
pyxasm_visitor visitor;
// << ", " << !for_loop_indent.empty() << "\n";
pyxasm_visitor visitor(bufferNames);
// Should we close a 'for' scope after this statement
// If > 0, indicate the number of for blocks to be closed.
int close_for_scopes = 0;
// If the stack is not empty and this line changed column to an outside
// scope:
while (!for_loop_indent.empty() && line.second < for_loop_indent.top()) {
// Pop the stack and flag to close the scope afterward
for_loop_indent.pop();
close_for_scopes++;
}
// Enter a new for loop -> push to the stack
if (line.first.find("for ") != std::string::npos) {
is_in_for_loop = true;
for_loop_indent.push(line.second);
}
// is_in_for_loop = line.first.find("for ") != std::string::npos &&
......@@ -128,17 +139,21 @@ void PyXasmTokenCollector::collect(clang::Preprocessor &PP,
ss << visitor.result.first;
}
if ((is_in_for_loop && line.second < previous_col) ||
(is_in_for_loop && line_counter == lines.size() - 1)) {
// we are now not in a for loop...
is_in_for_loop = false;
if (close_for_scopes > 0) {
// std::cout << "Close " << close_for_scopes << " for scopes.\n";
// need to close out the c++ or loop
ss << "}\n";
for (int i = 0; i < close_for_scopes; ++i) {
ss << "}\n";
}
}
previous_col = line.second;
line_counter++;
}
// If there are open for scope blocks here,
// i.e. for loops at the end of the function body.
while (!for_loop_indent.empty()) {
for_loop_indent.pop();
ss << "}\n";
}
}
} // namespace qcor
\ No newline at end of file
......@@ -15,18 +15,20 @@ using pyxasm_result_type =
std::pair<std::string, std::shared_ptr<xacc::Instruction>>;
class pyxasm_visitor : public pyxasmBaseVisitor {
protected:
protected:
std::shared_ptr<xacc::IRProvider> provider;
public:
pyxasm_visitor()
: provider(xacc::getIRProvider("quantum")) {}
// List of buffers in the *context* of this XASM visitor
std::vector<std::string> bufferNames;
public:
pyxasm_visitor(const std::vector<std::string> &buffers = {})
: provider(xacc::getIRProvider("quantum")), bufferNames(buffers) {}
pyxasm_result_type result;
bool in_for_loop = false;
antlrcpp::Any visitAtom_expr(
pyxasmParser::Atom_exprContext *context) override {
antlrcpp::Any
visitAtom_expr(pyxasmParser::Atom_exprContext *context) override {
if (context->atom()->NAME() != nullptr) {
auto inst_name = context->atom()->NAME()->getText();
......@@ -67,9 +69,9 @@ class pyxasm_visitor : public pyxasmBaseVisitor {
auto found_bracket = bit_expr_str.find_first_of("[");
if (found_bracket != std::string::npos) {
auto buffer_name = bit_expr_str.substr(0, found_bracket);
auto bit_idx_expr = bit_expr_str.substr(
found_bracket + 1,
bit_expr_str.length() - found_bracket - 2);
auto bit_idx_expr = bit_expr_str.substr(found_bracket + 1,
bit_expr_str.length() -
found_bracket - 2);
buffer_names.push_back(buffer_name);
inst->setBitExpression(i, bit_idx_expr);
} else {
......@@ -81,14 +83,68 @@ class pyxasm_visitor : public pyxasmBaseVisitor {
// Get the parameter expressions
int counter = 0;
for (int i = required_bits; i < atom_n_args; i++) {
inst->setParameter(
counter,
context->trailer()[0]->arglist()->argument()[i]->getText());
inst->setParameter(counter,
replacePythonConstants(context->trailer()[0]
->arglist()
->argument()[i]
->getText()));
counter++;
}
}
result.second = inst;
} else {
// Composite instructions, e.g. exp_i_theta
if (inst_name == "exp_i_theta") {
// Expected 3 params:
if (context->trailer()[0]->arglist()->argument().size() != 3) {
xacc::error(
"Invalid number of arguments for the 'exp_i_theta' "
"instruction. Expected 3, got " +
std::to_string(
context->trailer()[0]->arglist()->argument().size()) +
". Please check your input.");
}
std::stringstream ss;
// Delegate to the QRT call directly.
ss << "quantum::exp("
<< context->trailer()[0]->arglist()->argument(0)->getText()
<< ", "
<< context->trailer()[0]->arglist()->argument(1)->getText()
<< ", "
<< context->trailer()[0]->arglist()->argument(2)->getText()
<< ");\n";
result.first = ss.str();
} else {
xacc::error("Composite instruction '" + inst_name +
"' is not currently supported.");
}
}
} else {
// This kernel *callable* is not an intrinsic instruction, just
// reassemble the call:
// Check that the *first* argument is a *qreg* in the current context of
// *this* kernel.
if (!context->trailer().empty() &&
!context->trailer()[0]->arglist()->argument().empty() &&
xacc::container::contains(
bufferNames,
context->trailer()[0]->arglist()->argument(0)->getText())) {
std::stringstream ss;
// Use the kernel call with a parent kernel arg.
ss << inst_name << "(parent_kernel, ";
// TODO: We potentially need to handle *inline* expressions in the
// function call.
const auto &argList = context->trailer()[0]->arglist()->argument();
for (size_t i = 0; i < argList.size(); ++i) {
ss << argList[i]->getText();
if (i != argList.size() - 1) {
ss << ", ";
}
}
ss << ");\n";
result.first = ss.str();
}
result.second = inst;
}
}
return 0;
......@@ -96,28 +152,47 @@ class pyxasm_visitor : public pyxasmBaseVisitor {
antlrcpp::Any visitFor_stmt(pyxasmParser::For_stmtContext *context) override {
auto counter_expr = context->exprlist()->expr()[0];
auto iter_container = context->testlist()->test()[0]->getText();
// Rewrite:
// Python: "for <var> in <expr>:"
// C++: for (auto& var: <expr>) {}
// Note: we add range(int) as a C++ function to support this common pattern.
std::stringstream ss;
ss << "for (auto &" << counter_expr->getText() << " : " << iter_container
<< ") {\n";
result.first = ss.str();
in_for_loop = true;
return 0;
}
if (context->testlist()->test()[0]->getText().find("range") !=
std::string::npos) {
auto range_str = context->testlist()->test()[0]->getText();
auto found_paren = range_str.find_first_of("(");
auto range_contents = range_str.substr(
found_paren + 1, range_str.length() - found_paren - 2);
antlrcpp::Any visitExpr_stmt(pyxasmParser::Expr_stmtContext *ctx) override {
if (ctx->ASSIGN().size() == 1 && ctx->testlist_star_expr().size() == 2) {
// Handle simple assignment: a = expr
std::stringstream ss;
ss << "for (int " << counter_expr->getText() << " = 0; "
<< counter_expr->getText() << " < " << range_contents << "; ++"
<< counter_expr->getText() << " ) {\n";
const std::string lhs = ctx->testlist_star_expr(0)->getText();
const std::string rhs = ctx->testlist_star_expr(1)->getText();
ss << "auto " << lhs << " = " << rhs << "; \n";
result.first = ss.str();
in_for_loop = true;
return 0;
} else {
xacc::error(
"QCOR PyXasm can only handle 'for VAR in range(QREG.size())' at the "
"moment.");
return visitChildren(ctx);
}
}
return 0;
private:
// Replaces common Python constants, e.g. 'math.pi' or 'numpy.pi'.
// Note: the library names have been resolved to their original names.
std::string replacePythonConstants(const std::string &in_pyExpr) const {
// List of all keywords to be replaced
const std::map<std::string, std::string> REPLACE_MAP{{"math.pi", "M_PI"},
{"numpy.pi", "M_PI"}};
std::string newSrc = in_pyExpr;
for (const auto &[key, value] : REPLACE_MAP) {
const auto pos = newSrc.find(key);
if (pos != std::string::npos) {
newSrc.replace(pos, key.length(), value);
}
}
return newSrc;
}
};
\ No newline at end of file
......@@ -31,7 +31,7 @@ TEST(PyXASMTokenCollectorTester, checkSimple) {
EXPECT_EQ(R"#(quantum::h(qb[0]);
quantum::cnot(qb[0], qb[1]);
for (int i = 0; i < qb.size(); ++i ) {
for (auto &i : range(qb.size())) {
quantum::x(qb[i]);
quantum::x(qb[i]);
quantum::mz(qb[i]);
......
......@@ -154,9 +154,9 @@ TEST(TokenCollectorTester, checkPyXasm) {
auto results =
qcor::run_token_collector(*PP, cached, {"qb"});
std::cout << results << "\n";
EXPECT_EQ(R"#(quantum::h(qb[0]);
EXPECT_EQ(R"#(quantum::h(qb[0]);
quantum::cnot(qb[0], qb[1]);
for (int i = 0; i < qb.size(); ++i ) {
for (auto &i : range(qb.size())) {
quantum::x(qb[i]);
quantum::x(qb[i]);
quantum::mz(qb[i]);
......
# Run this from the command line like this
#
# python3 exp_i_theta.py -shots 100
from qcor import qjit, qalloc, qreg
# To create QCOR quantum kernels in Python one
# simply creates a Python function, writes Pythonic,
# XASM-like quantum code, and annotates the kernel
# to indicate it is meant for QCOR just in time compilation
# NOTE Programmers must type annotate their function arguments
# Define a XASM kernel
@qjit
def exp_circuit(q : qreg, t0: float, t1: float):
exponent_op1 = X(0) * Y(1) - Y(0) * X(1)
exponent_op2 = X(0) * Z(1) * Y(2) - X(2) * Z(1) * Y(0)
X(q[0])
exp_i_theta(q, t0, exponent_op1)
exp_i_theta(q, t1, exponent_op2)
for i in range(q.size()):
Measure(q[i])
# Allocate 3 qubits
q = qalloc(3)
# Run the experiment with some random angles
theta1 = 1.234
theta2 = 2.345
# Examine the circuit QASM
comp = exp_circuit.extract_composite(q, theta1, theta2)
print(comp.toString())
# Execute
exp_circuit(q, theta1, theta2)
# Print the results
q.print()
\ No newline at end of file
# Run this from the command line like this
#
# python3 multiple_kernels.py -shots 100
from qcor import qjit, qalloc, qreg
# To create QCOR quantum kernels in Python one
# simply creates a Python function, writes Pythonic,
# XASM-like quantum code, and annotates the kernel
# to indicate it is meant for QCOR just in time compilation
# NOTE Programmers must type annotate their function arguments
@qjit
def measure_all_qubits(q : qreg):
for i in range(q.size()):
Measure(q[i])
# Define a Bell kernel
@qjit
def bell_test(q : qreg):
H(q[0])
CX(q[0], q[1])
# Call other kernels
measure_all_qubits(q)
# Allocate 2 qubits
q = qalloc(2)
# Inspect the IR
comp = bell_test.extract_composite(q)
print(comp.toString())
# Run the bell experiment
bell_test(q)
# Print the results
q.print()
\ No newline at end of file
# Run this from the command line like this
#
# python3 qaoa_circuit.py
from qcor import *
import numpy as np
from types import MethodType
# Define a QAOA kernel with variational parameters (theta and beta angles)
@qjit
def qaoa_circ(q: qreg, cost_ham: PauliOperator, nbSteps: int, theta: List[float], beta: List[float]):
# Start off in the uniform superposition
for i in range(q.size()):
H(q[i])
terms = cost_ham.getNonIdentitySubTerms()
for step in range(nbSteps):
for term in terms:
exp_i_theta(q, theta[step], term)
# Reference Hamiltonian:
for i in range(len(q)):
ref_ham_term = X(i)
exp_i_theta(q, beta[step], ref_ham_term)
# Allocate 4 qubits
q = qalloc(4)
n_steps = 3
# Hamiltonion:
H = -5.0 - 0.5 * (Z(0) - Z(3) - Z(1) * Z(2)) - Z(2) + 2 * Z(0) * Z(2) + 2.5 * Z(2) * Z(3)
# Custom arg_translator in a Pythonic way
def qaoa_translate(self, q: qreg, x: List[float]):
ret_dict = {}
ret_dict["q"] = q
ret_dict["cost_ham"] = H
ret_dict["nbSteps"] = n_steps
ret_dict["theta"] = x[:n_steps]
ret_dict["beta"] = x[n_steps:]
return ret_dict
# Rebind arg translate:
qaoa_circ.translate = MethodType(qaoa_translate, qjit)
# Use the standard parameterization scheme:
# one theta + one beta per step
n_params = 2 * n_steps
obj = createObjectiveFunction(qaoa_circ, H, n_params)
# Run optimization
optimizer = createOptimizer('nlopt', {'initial-parameters': np.random.rand(n_params)})
results = optimizer.optimize(obj)
......@@ -54,7 +54,7 @@ namespace {
// Here we enumerate them as a Variant
using AllowedKernelArgTypes =
xacc::Variant<bool, int, double, std::string, xacc::internal_compiler::qreg,
std::vector<double>>;
std::vector<double>, qcor::PauliOperator>;
// We will take as input a mapping of arg variable names to the argument itself.
using KernelArgDict = std::map<std::string, AllowedKernelArgTypes>;
......@@ -76,8 +76,9 @@ class KernelArgDictToHeterogeneousMap {
};
// Add type name to this list to support receiving from Python.
using PyHeterogeneousMapTypes = xacc::Variant<bool, int, double, std::string,
std::shared_ptr<qcor::Optimizer>>;
using PyHeterogeneousMapTypes =
xacc::Variant<bool, int, double, std::string,
std::shared_ptr<qcor::Optimizer>, std::vector<double>>;
using PyHeterogeneousMap = std::map<std::string, PyHeterogeneousMapTypes>;
// Helper to convert a Python *dict* (as a map of variants) into a native
......@@ -299,9 +300,10 @@ PYBIND11_MODULE(_pyqcor, m) {
.def("jit_compile", &qcor::QJIT::jit_compile, "")
.def(
"internal_python_jit_compile",
[](qcor::QJIT &qjit, const std::string src) {
[](qcor::QJIT &qjit, const std::string src,
const std::vector<std::string> &dependency = {}) {
bool turn_on_hetmap_kernel_ctor = true;
qjit.jit_compile(src, turn_on_hetmap_kernel_ctor);
qjit.jit_compile(src, turn_on_hetmap_kernel_ctor, dependency);
},
"")
.def("run_syntax_handler", &qcor::QJIT::run_syntax_handler, "")
......
......@@ -4,9 +4,11 @@ import sys
import inspect
from typing import List
import typing
import re
from collections import defaultdict
List = typing.List
PauliOperator = xacc.quantum.PauliOperator
def X(idx):
return xacc.quantum.PauliOperator({idx: 'X'}, 1.0)
......@@ -19,6 +21,70 @@ def Y(idx):
def Z(idx):
return xacc.quantum.PauliOperator({idx: 'Z'}, 1.0)
# Simple graph class to help resolve kernel dependency (via topological sort)
class KernelGraph(object):
def __init__(self):
self.graph = defaultdict(list)
self.V = 0
self.kernel_idx_dep_map = {}
self.kernel_name_list = []
def addKernelDependency(self, kernelName, depList):
self.kernel_name_list.append(kernelName)
self.kernel_idx_dep_map[self.V] = []
for dep_ker_name in depList:
self.kernel_idx_dep_map[self.V].append(self.kernel_name_list.index(dep_ker_name))
self.V += 1
def addEdge(self, u, v):
self.graph[u].append(v)
# Topological Sort.
def topologicalSort(self):
self.graph = defaultdict(list)
for sub_ker_idx in self.kernel_idx_dep_map:
for dep_sub_idx in self.kernel_idx_dep_map[sub_ker_idx]:
self.addEdge(dep_sub_idx, sub_ker_idx)
in_degree = [0]*(self.V)
for i in self.graph:
for j in self.graph[i]:
in_degree[j] += 1
queue = []
for i in range(self.V):
if in_degree[i] == 0:
queue.append(i)
cnt = 0
top_order = []
while queue:
u = queue.pop(0)
top_order.append(u)
for i in self.graph[u]:
in_degree[i] -= 1
if in_degree[i] == 0:
queue.append(i)
cnt += 1
sortedDep = []
for sorted_dep_idx in top_order:
sortedDep.append(self.kernel_name_list[sorted_dep_idx])
return sortedDep
def getSortedDependency(self, kernelName):
kernel_idx = self.kernel_name_list.index(kernelName)
# No dependency
if len(self.kernel_idx_dep_map[kernel_idx]) == 0:
return []
sorted_dep = self.topologicalSort()
result_dep = []
for dep_name in sorted_dep:
if dep_name == kernelName:
return result_dep
else:
result_dep.append(dep_name)
class qjit(object):
"""
......@@ -60,7 +126,9 @@ class qjit(object):
self.kwargs = kwargs
self.function = function
self.allowed_type_cpp_map = {'<class \'_pyqcor.qreg\'>': 'qreg',
'<class \'float\'>': 'double', 'typing.List[float]': 'std::vector<double>'}
'<class \'float\'>': 'double', 'typing.List[float]': 'std::vector<double>',
'<class \'int\'>': 'int',
'<class \'_pyxacc.quantum.PauliOperator\'>': 'qcor::PauliOperator'}
self.__dict__.update(kwargs)
# Create the qcor just in time engine
......@@ -88,16 +156,64 @@ class qjit(object):
self.allowed_type_cpp_map[str(_type)] + ' ' + arg
cpp_arg_str = cpp_arg_str[1:]
globalVarDecl = []
# Get all globals currently defined at this stack frame
globalsInStack = inspect.stack()[1][0].f_globals
globalVars = globalsInStack.copy()
importedModules = {}
for key in globalVars:
descStr = str(globalVars[key])
# Cache module import and its potential alias
# e.g. import abc as abc_alias
if descStr.startswith("<module "):
moduleName = descStr.split()[1].replace("'", "")
importedModules[key] = moduleName
else:
# Import global variables:
# Only support float atm
if (isinstance(globalVars[key], float)):
globalVarDecl.append(key + " = " + str(globalVars[key]))
# Inject these global declarations into the function body.
separator = "\n"
globalDeclStr = separator.join(globalVarDecl)
# Handle common modules like numpy or math
# e.g. if seeing `import numpy as np`, we'll have <'np' -> 'numpy'> in the importedModules dict.
# We'll replace any module alias by its original name,
# i.e. 'np.pi' -> 'numpy.pi', etc.
for moduleAlias in importedModules:
if moduleAlias != importedModules[moduleAlias]:
aliasModuleStr = moduleAlias + '.'
originalModuleStr = importedModules[moduleAlias] + '.'
fbody_src = fbody_src.replace(aliasModuleStr, originalModuleStr)
# Create the qcor quantum kernel function src for QJIT and the Clang syntax handler
self.src = '__qpu__ void '+self.function.__name__ + \
'('+cpp_arg_str+') {\nusing qcor::pyxasm;\n'+fbody_src+"}\n"
'('+cpp_arg_str+') {\nusing qcor::pyxasm;\n' + globalDeclStr + '\n' + fbody_src +"}\n"
# Handle nested kernels:
dependency = []
for kernelName in self.__compiled__kernels:
kernelCall = kernelName + '('
# Check that this kernel *calls* a previously-compiled kernel:
# pattern: "<white space> kernel("
if re.search(r"\b" + re.escape(kernelCall), self.src):
dependency.append(kernelName)
self.__kernels__graph.addKernelDependency(self.function.__name__, dependency)
sorted_kernel_dep = self.__kernels__graph.getSortedDependency(self.function.__name__)
# Run the QJIT compile step to store function pointers internally
self._qjit.internal_python_jit_compile(self.src)
self._qjit.internal_python_jit_compile(self.src, sorted_kernel_dep)