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

Handle Op-Assign (<<=, +=, etc.) codegen in PyXASM 



also, add the test for bit reverse in qjit kernel

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

handlers/token_collector/pyxasm/pyxasm_visitor.hpp

+36 −1
Original line number Diff line number Diff line
@@ -457,7 +457,17 @@ class pyxasm_visitor : public pyxasmBaseVisitor { 
        return 0;

      }

    } else {

      return visitChildren(ctx);

      // Visit child node:

      auto child_result = visitChildren(ctx);

      const auto translated_src = sub_node_translation.str();

      sub_node_translation.str(std::string());

      // If no child nodes, perform the codegen (result.first is not set)

      // but just appending the incremental translation collector;

      // return the collected C++ statement.

      if (result.first.empty() && !translated_src.empty()) {

        result.first = translated_src + ";\n";

      }

      return child_result;

    }

  }
@@ -509,6 +519,31 @@ class pyxasm_visitor : public pyxasmBaseVisitor { 
    return 0;

  }



  virtual antlrcpp::Any visitTestlist_star_expr(

      pyxasmParser::Testlist_star_exprContext *context) override {

    // std::cout << "Testlist_star_exprContext:" << context->getText() << "\n";

    const auto var_name = context->getText();

    if (xacc::container::contains(declared_var_names, var_name)) {

      sub_node_translation << var_name << " ";

      return 0;

    }

    return visitChildren(context);

  }



  virtual antlrcpp::Any

  visitAugassign(pyxasmParser::AugassignContext *context) override {

    // std::cout << "Augassign:" << context->getText() << "\n";

    sub_node_translation << context->getText() << " ";

    return 0;

  }



  virtual antlrcpp::Any

  visitTestlist(pyxasmParser::TestlistContext *context) override {

    // std::cout << "visitTestlist:" << context->getText() << "\n";

    sub_node_translation << context->getText() << " ";

    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.

python/tests/test_jit_simple.py

+29 −0
Original line number Diff line number Diff line
@@ -232,5 +232,34 @@ class TestKernelJIT(unittest.TestCase): 
        self.assertEqual(q.counts()["0"], r.counts()["0"])

        self.assertEqual(q.counts()["1"], r.counts()["1"])

    

    def test_while_loop(self):

        @qjit 

        def while_loop_kernel(q: qubit, x: int, exp_inv: int):

            rev = 0

            while x:

                rev <<= 1

                rev += x & 1

                x >>= 1



            if rev == exp_inv:

                print("Success")

                X(q)



        # Reference: pure python to check

        def reverse_bit(num):

            result = 0

            while num:

                result = (result << 1) + (num & 1)

                num >>= 1

            return result

        

        q = qalloc(1)

        test_val = 123

        expected = reverse_bit(test_val)

        comp0 = while_loop_kernel.extract_composite(q[0], test_val, expected)    

        print("Comp:", comp0)

        # Has X applied.

        self.assertEqual(comp0.nInstructions(), 1) 



if __name__ == '__main__':

  unittest.main()
 
 No newline at end of file