Branch-based 'if' is needed to handle break/continue (a08f5355) · Commits · ORNL Quantum Computing Institute / qcor

mlir/parsers/qasm3/qasm3_visitor.hpp

+11 −0

Original line number	Diff line number	Diff line
		@@ -267,6 +267,17 @@ class qasm3_visitor : public qasm3::qasm3BaseVisitor {

		return allocation;
		}

		template <class NodeType>
		bool hasChildNodeOfType(antlr4::tree::ParseTree &in_node) {
		for (auto &child_node : in_node.children) {
		if (dynamic_cast<NodeType *>(child_node) \|\|
		hasChildNodeOfType<NodeType>(*child_node)) {
		return true;
		}
		}
		return false;
		}
		};

		} // namespace qcor

mlir/parsers/qasm3/visitor_handlers/conditional_handler.cpp

+124 −38

Original line number	Diff line number	Diff line

		#include "expression_handler.hpp"
		#include "qasm3_visitor.hpp"
		#include "mlir/Dialect/SCF/SCF.h"
		#include "qasm3_visitor.hpp"
		namespace {
		// ATM, we don't try to convert everything to the
		// special Quantum If-Then-Else Op.
		@@ -10,7 +10,8 @@ namespace {
		// i.e., this serve mainly as a stop-gap before fully-FTQC runtimes become
		// available.

		// FIXME: Define a Target Capability setting and make the compiler aware of that.
		// FIXME: Define a Target Capability setting and make the compiler aware of
		// that.

		// Capture binary comparison conditional.
		// Note: currently, only bit-level is modeled.
		@@ -201,6 +202,90 @@ antlrcpp::Any qasm3_visitor::visitBranchingStatement(
		}
		}

		// Manually write the conditional block:
		// If there is 'break', 'continue' (ControlDirective) in the body.
		// The reason being: these break/continue will be translated to BranchOp
		// which are overlapping with the BranchOp implicitly added at the end of SCF::IfOp.
		// e.g.,
		// br ^bb1 (e.g., out of the outer loop) <-- added by ControlDirectiveContext handler
		// br ^bb2 (e.g., to the end of the if statement) <-- added by the implicit yield op
		// The verify step (MLIR -> LLVM) will complain this....
		if (hasChildNodeOfType<qasm3Parser::ControlDirectiveContext>(*context)) {
		qasm3_expression_generator exp_generator(builder, symbol_table, file_name);
		exp_generator.visit(conditional_expr);
		auto expr_value = exp_generator.current_value;

		// build up the program block
		auto currRegion = builder.getBlock()->getParent();
		auto savept = builder.saveInsertionPoint();
		auto thenBlock = builder.createBlock(currRegion, currRegion->end());
		auto elseBlock = builder.createBlock(currRegion, currRegion->end());
		mlir::Block *exitBlock = nullptr;
		// If we have an else block from programBlock,
		// then create a stand alone exit block that both
		// then and else can fall to
		if (context->programBlock().size() == 2) {
		exitBlock = builder.createBlock(currRegion, currRegion->end());
		} else {
		exitBlock = elseBlock;
		}

		// Build up the THEN Block
		builder.setInsertionPointToStart(thenBlock);
		symbol_table.enter_new_scope();
		// Get the conditional code and visit the nodes
		auto conditional_code = context->programBlock(0);
		visitChildren(conditional_code);

		// Need to check if we have a branch out of
		// this thenBlock, if so do not add a branch
		// to the exitblock
		mlir::Operation &last_op = thenBlock->back();
		auto branchOps = thenBlock->getOps<mlir::BranchOp>();
		auto branch_to_exit = true;
		for (auto b : branchOps) {
		if (b.dest() == current_loop_exit_block \|\|
		b.dest() == current_loop_header_block \|\|
		b.dest() == current_loop_incrementor_block) {
		branch_to_exit = false;
		break;
		}
		}
		if (branch_to_exit) {
		builder.create<mlir::BranchOp>(location, exitBlock);
		}
		symbol_table.exit_scope();

		// If we have a second program block then we have an else stmt
		builder.setInsertionPointToStart(elseBlock);
		if (context->programBlock().size() == 2) {
		symbol_table.enter_new_scope();
		visitChildren(context->programBlock(1));
		branch_to_exit = true;
		for (auto b : branchOps) {
		if (b.dest() == current_loop_exit_block \|\|
		b.dest() == current_loop_header_block \|\|
		b.dest() == current_loop_incrementor_block) {
		branch_to_exit = false;
		break;
		}
		}
		if (branch_to_exit) {
		builder.create<mlir::BranchOp>(location, exitBlock);
		}

		symbol_table.exit_scope();
		}

		// Restore the insertion point and create the conditional statement
		builder.restoreInsertionPoint(savept);
		builder.create<mlir::CondBranchOp>(location, expr_value, thenBlock,
		elseBlock);
		builder.setInsertionPointToStart(exitBlock);

		symbol_table.set_last_created_block(exitBlock);
		} else {
		// Using SCF::IfOp
		// Map it to a Value
		qasm3_expression_generator exp_generator(builder, symbol_table, file_name);
		exp_generator.visit(conditional_expr);
		@@ -237,6 +322,7 @@ antlrcpp::Any qasm3_visitor::visitBranchingStatement(
		}
		// Restore builder
		builder = cached_builder;
		}
		return 0;
		}
		} // namespace qcor
		No newline at end of file

mlir/parsers/qasm3/visitor_handlers/loop_stmt_handler.cpp

+6 −10

Original line number	Diff line number	Diff line
		@@ -307,17 +307,13 @@ antlrcpp::Any qasm3_visitor::visitLoopStatement(
		const std::string program_block_str = program_block->getText();
		// std::cout << "HOWDY:\n" << program_block_str << "\n";

		// HACK: Currently, we don't handle 'break', 'continue'
		// HACK: Currently, we don't handle 'break', 'continue' or nested loop
		// in the Affine for loop yet.
		if (program_block_str.find("break") == std::string::npos &&
		program_block_str.find("continue") == std::string::npos &&
		// This is equivalent to an "if"
		program_block_str.find("QCOR_EXPECT_TRUE") == std::string::npos &&
		// We can only handle nested for loops if the inner one is also an
		// affine one For now, don't do that since we're not sure.
		program_block_str.find("for") == std::string::npos &&
		// While loop is not converted to affine yet.
		program_block_str.find("while") == std::string::npos) {
		if (!hasChildNodeOfType<qasm3Parser::ControlDirectiveContext>(
		*program_block) &&
		!hasChildNodeOfType<qasm3Parser::LoopStatementContext>(
		*program_block) &&
		(program_block_str.find("QCOR_EXPECT_TRUE") == std::string::npos)) {
		// Can use Affine for loop....
		affineLoopBuilder(
		a_value, b_value, c,