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

Support common Python constants, e.g. Pi 



The Python side will make sure all module aliasing resolved.
The C++ will then map them to C++ constants.

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

handlers/token_collector/pyxasm/pyxasm_visitor.hpp

+22 −3
Original line number Diff line number Diff line
@@ -81,9 +81,11 @@ 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++;

            }

          }
@@ -134,4 +136,21 @@ class pyxasm_visitor : public pyxasmBaseVisitor { 
      return visitChildren(ctx);

    }

  }



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

python/qcor.py

+9 −1
Original line number Diff line number Diff line
@@ -110,8 +110,16 @@ class qjit(object): 
        separator = "\n"

        globalDeclStr = separator.join(globalVarDecl)



        # TODO: Handle common module like numpy or math

        # 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' + globalDeclStr + '\n' + fbody_src +"}\n"

python/tests/test_kernel_jit.py

+28 −1
Original line number Diff line number Diff line 
import unittest

from qcor import *

# Import Python math with alias

import math as myMathMod



# Some global variables for testing

MY_PI = 3.1416
@@ -80,5 +82,30 @@ class TestSimpleKernelJIT(unittest.TestCase): 
        # Pi pulse -> X gate

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

    

    def test_ModuleConstants(self):



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



        @qjit

        def kernelUseConstants(q : qreg):

            # Use math.pi constant

            Rx(q[0], myMathMod.pi/2)

            CX(q[0], q[1])

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

                Measure(q[i])



        # Allocate 2 qubits

        q = qalloc(2)



        # Run the bell experiment

        kernelUseConstants(q)



        # Print the results

        q.print()

        counts = q.counts()

        self.assertEqual(len(counts), 2)

        self.assertTrue('00' in counts)

        self.assertTrue('11' in counts)



if __name__ == '__main__':

  unittest.main()
 
 No newline at end of file