Commit 50b2fa57 authored by Nguyen, Thien Minh's avatar Nguyen, Thien Minh
Browse files

Fixed a bug w/ nested for loops 



Need to handle closing multiple for loops based on code indentation.

Signed-off-by: default avatarThien Nguyen <nguyentm@ornl.gov>
parent 5b3c25ab

handlers/token_collector/pyxasm/pyxasm_token_collector.cpp

+10 −6
Original line number Diff line number Diff line
@@ -95,17 +95,18 @@ void PyXasmTokenCollector::collect(clang::Preprocessor &PP, 
  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";

    //           << ", " << !for_loop_indent.empty() << "\n";



    pyxasm_visitor visitor(bufferNames);

    // Should we close a 'for' scope after this statement

    bool close_for_scope = false;

    // 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:

    if (!for_loop_indent.empty() && line.second < for_loop_indent.top()) {

    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_scope = true;

      close_for_scopes++;

    }



    // Enter a new for loop -> push to the stack
@@ -138,10 +139,13 @@ void PyXasmTokenCollector::collect(clang::Preprocessor &PP, 
      ss << visitor.result.first;

    }



    if (close_for_scope) {

    if (close_for_scopes > 0) {

      // std::cout << "Close " << close_for_scopes << " for scopes.\n";

      // need to close out the c++ or loop

      for (int i = 0; i < close_for_scopes; ++i) {

        ss << "}\n";

      }

    }

    previous_col = line.second;

    line_counter++;

  }

python/tests/test_kernel_jit.py

+26 −5
Original line number Diff line number Diff line
@@ -8,7 +8,7 @@ MY_PI = 3.1416 
class TestSimpleKernelJIT(unittest.TestCase):

    def test_simple_bell(self):



        set_qpu('qpp', {'shots':8192})

        set_qpu('qpp', {'shots':1024})



        @qjit

        def bell(q : qreg):
@@ -33,7 +33,7 @@ class TestSimpleKernelJIT(unittest.TestCase): 


    def test_assignment(self):



        set_qpu('qpp', {'shots':8192})

        set_qpu('qpp', {'shots':1024})



        @qjit

        def varAssignKernel(q : qreg):
@@ -61,7 +61,7 @@ class TestSimpleKernelJIT(unittest.TestCase): 
    

    def test_globalVar(self):



        set_qpu('qpp', {'shots':8192})

        set_qpu('qpp', {'shots':1024})



        @qjit

        def kernelUseGlobals(q : qreg):
@@ -80,11 +80,11 @@ class TestSimpleKernelJIT(unittest.TestCase): 
        q.print()

        counts = q.counts()

        # Pi pulse -> X gate

        self.assertTrue(counts['11'] > 8000)

        self.assertTrue(counts['11'] > 1000)

    

    def test_ModuleConstants(self):



        set_qpu('qpp', {'shots':8192})

        set_qpu('qpp', {'shots':1024})



        @qjit

        def kernelUseConstants(q : qreg):
@@ -198,5 +198,26 @@ class TestSimpleKernelJIT(unittest.TestCase): 
        for i in range(5, 10):

            self.assertEqual(comp.getInstruction(i).name(), "Measure") 



    def test_for_loop(self):

        @qjit

        def kernels_w_loops(q : qreg, thetas : List[float], betas : List[float]):

            for i in range(len(q)):

                for theta in thetas:

                    for beta in betas:

                        angle = theta + beta

                        Rx(q[i], angle)

            for i in range(q.size()):

                Measure(q[i])



        list1 = [1.0, 2.0, 3.0]

        list2 = [4.0, 5.0, 6.0]

        q = qalloc(2)

        comp = kernels_w_loops.extract_composite(q, list1, list2)

        self.assertEqual(comp.nInstructions(), q.size() * len(list1) * len(list2) + q.size())

        for i in range(0, q.size() * len(list1) * len(list2)):

            self.assertEqual(comp.getInstruction(i).name(), "Rx") 

        for i in range(q.size() * len(list1) * len(list2), comp.nInstructions()):

            self.assertEqual(comp.getInstruction(i).name(), "Measure")    



if __name__ == '__main__':

  unittest.main()
 
 No newline at end of file