Rc sync to main (#571)

Co-authored-by: David Ittah <dime10@users.noreply.github.com> Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>
PennyLaneAI · Mar 1, 2024 · 5ed2778 · 5ed2778
1 parent 2689a35
commit 5ed2778
Show file tree

Hide file tree

Showing 24 changed files with 384 additions and 71 deletions.
diff --git a/.dep-versions b/.dep-versions
@@ -3,4 +3,7 @@ jax=0.4.23
 mhlo=4611968a5f6818e6bdfb82217b9e836e0400bba9
 llvm=cd9a641613eddf25d4b25eaa96b2c393d401d42c
 enzyme=1beb98b51442d50652eaa3ffb9574f4720d611f1
+
+# Always remove custom PL/LQ versions before release.
 pennylane=f638bc53a09724a83d9d58964bf37bfd438e6aa3
+lightning=0.35.0-dev21
diff --git a/.github/workflows/build-wheel-linux-x86_64.yaml b/.github/workflows/build-wheel-linux-x86_64.yaml
@@ -292,7 +292,7 @@ jobs:
               -DPYTHON_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -Dpybind11_DIR=$(python${{ matrix.python_version }} -c "import pybind11; print(pybind11.get_cmake_dir())") \
               -DENABLE_LIGHTNING_KOKKOS=ON \
-              -DLIGHTNING_GIT_TAG="2716864521c539429bd382b92151a918d1a076ce" \
+              -DLIGHTNING_GIT_TAG="03e465d96a665a2d1bcea2f1e6d79321e9d1c3fe" \
               -DKokkos_ENABLE_SERIAL=ON \
               -DKokkos_ENABLE_OPENMP=ON \
               -DENABLE_WARNINGS=OFF \
@@ -376,15 +376,15 @@ jobs:
       run: |
         python${{ matrix.python_version }} -m pip install pytest pytest-xdist
 
-    - name: Install Catalyst
-      run: |
-        python${{ matrix.python_version }} -m pip install $GITHUB_WORKSPACE/dist/*.whl
-
     - name: Install PennyLane Plugins
       run: |
         python${{ matrix.python_version }} -m pip install PennyLane-Lightning-Kokkos
         python${{ matrix.python_version }} -m pip install amazon-braket-pennylane-plugin "boto3==1.26"
 
+    - name: Install Catalyst
+      run: |
+        python${{ matrix.python_version }} -m pip install $GITHUB_WORKSPACE/dist/*.whl --extra-index-url https://test.pypi.org/simple
+
     - name: Run Python Pytest Tests
       run: |
         python${{ matrix.python_version }} -m pytest $GITHUB_WORKSPACE/frontend/test/pytest -n auto

diff --git a/.github/workflows/build-wheel-macos-arm64.yaml b/.github/workflows/build-wheel-macos-arm64.yaml
@@ -289,7 +289,7 @@ jobs:
               -DPYTHON_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -Dpybind11_DIR=$(python${{ matrix.python_version }} -c "import pybind11; print(pybind11.get_cmake_dir())") \
               -DENABLE_LIGHTNING_KOKKOS=ON \
-              -DLIGHTNING_GIT_TAG="2716864521c539429bd382b92151a918d1a076ce" \
+              -DLIGHTNING_GIT_TAG="03e465d96a665a2d1bcea2f1e6d79321e9d1c3fe" \
               -DKokkos_ENABLE_SERIAL=ON \
               -DKokkos_ENABLE_OPENMP=ON \
               -DKokkos_ENABLE_COMPLEX_ALIGN=OFF \
@@ -385,15 +385,15 @@ jobs:
       run: |
         python${{ matrix.python_version }} -m pip install pytest pytest-xdist
 
-    - name: Install Catalyst
-      run: |
-        python${{ matrix.python_version }} -m pip install $GITHUB_WORKSPACE/dist/*.whl
-
     - name: Install PennyLane Plugins
       run: |
         python${{ matrix.python_version }} -m pip install PennyLane-Lightning-Kokkos
         python${{ matrix.python_version }} -m pip install amazon-braket-pennylane-plugin
 
+    - name: Install Catalyst
+      run: |
+        python${{ matrix.python_version }} -m pip install $GITHUB_WORKSPACE/dist/*.whl --extra-index-url https://test.pypi.org/simple
+
     - name: Run Python Pytest Tests
       run: |
         python${{ matrix.python_version }} -m pytest -v $GITHUB_WORKSPACE/frontend/test/pytest -n auto

diff --git a/.github/workflows/build-wheel-macos-x86_64.yaml b/.github/workflows/build-wheel-macos-x86_64.yaml
@@ -261,7 +261,7 @@ jobs:
               -DCMAKE_LIBRARY_OUTPUT_DIRECTORY=$GITHUB_WORKSPACE/runtime-build/lib \
               -DPYTHON_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -Dpybind11_DIR=$(python${{ matrix.python_version }} -c "import pybind11; print(pybind11.get_cmake_dir())") \
-              -DLIGHTNING_GIT_TAG="2716864521c539429bd382b92151a918d1a076ce" \
+              -DLIGHTNING_GIT_TAG="03e465d96a665a2d1bcea2f1e6d79321e9d1c3fe" \
               -DENABLE_LIGHTNING_KOKKOS=ON \
               -DKokkos_ENABLE_SERIAL=ON \
               -DKokkos_ENABLE_OPENMP=OFF \
@@ -355,15 +355,15 @@ jobs:
       run: |
         python${{ matrix.python_version }} -m pip install pytest pytest-xdist
 
-    - name: Install Catalyst
-      run: |
-        python${{ matrix.python_version }} -m pip install $GITHUB_WORKSPACE/dist/*.whl
-
     - name: Install PennyLane Plugins
       run: |
         python${{ matrix.python_version }} -m pip install PennyLane-Lightning-Kokkos
         python${{ matrix.python_version }} -m pip install amazon-braket-pennylane-plugin
 
+    - name: Install Catalyst
+      run: |
+        python${{ matrix.python_version }} -m pip install $GITHUB_WORKSPACE/dist/*.whl --extra-index-url https://test.pypi.org/simple
+
     - name: Run Python Pytest Tests
       run: |
         python${{ matrix.python_version }} -m pytest $GITHUB_WORKSPACE/frontend/test/pytest -n auto

diff --git a/.github/workflows/check-catalyst.yaml b/.github/workflows/check-catalyst.yaml
@@ -378,14 +378,14 @@ jobs:
     - name: Install Deps
       run: |
         sudo apt-get update
-        sudo apt-get install -y python3 python3-pip libomp-dev libasan6
+        sudo apt-get install -y python3 python3-pip libomp-dev libasan6 make
         python3 --version | grep ${{ needs.constants.outputs.primary_python_version }}
         python3 -m pip install -r requirements.txt
         # cuda-quantum is added manually here.
         # It can't be in requirements.txt as that will break
         # macOS requirements.txt
         python3 -m pip install cuda-quantum
-        python3 -m pip install .
+        make frontend
 
     - name: Get Cached LLVM Build
       id: cache-llvm-build
@@ -447,13 +447,17 @@ jobs:
     - name: Checkout Catalyst repo
       uses: actions/checkout@v3
 
+    - name: Install lightning.kokkos used in Python tests
+      run: |
+        pip install PennyLane-Lightning-Kokkos
+
     - name: Install Deps
       run: |
         sudo apt-get update
-        sudo apt-get install -y python3 python3-pip libomp-dev libasan6
+        sudo apt-get install -y python3 python3-pip libomp-dev libasan6 make
         python3 --version | grep ${{ needs.constants.outputs.primary_python_version }}
         python3 -m pip install -r requirements.txt
-        python3 -m pip install .
+        make frontend
 
     - name: Get Cached LLVM Build
       id: cache-llvm-build
@@ -481,10 +485,6 @@ jobs:
         echo "RUNTIME_LIB_DIR=$(pwd)/runtime-build/lib" >> $GITHUB_ENV
         echo "MLIR_LIB_DIR=$(pwd)/llvm-build/lib" >> $GITHUB_ENV
 
-    - name: Install lightning.kokkos used in Python tests
-      run: |
-        pip install PennyLane-Lightning-Kokkos
-
     - name: Run Python Pytest Tests (backend=lightning.kokkos)
       run: |
         make pytest TEST_BACKEND="lightning.kokkos"
@@ -504,10 +504,10 @@ jobs:
     - name: Install Deps
       run: |
         sudo apt-get update
-        sudo apt-get install -y python3 python3-pip libomp-dev libasan6
+        sudo apt-get install -y python3 python3-pip libomp-dev libasan6 make
         python3 --version | grep ${{ needs.constants.outputs.primary_python_version }}
         python3 -m pip install -r requirements.txt
-        python3 -m pip install .
+        make frontend
 
     - name: Get Cached LLVM Build
       id: cache-llvm-build

diff --git a/.github/workflows/check-pl-compat.yaml b/.github/workflows/check-pl-compat.yaml
@@ -101,13 +101,7 @@ jobs:
         DIALECTS_BUILD_DIR="$(pwd)/quantum-build" \
         ENABLE_LLD=ON \
         make dialects
-        if [ ${{ inputs.pennylane }} = "stable" ]; then
-            pip install --upgrade .
-        else
-            # TODO(@erick-xanadu): Remove after release. See issue
-            # https://github.com/PennyLaneAI/catalyst/issues/494
-            pl_version="==0.35.0.dev0" pip install --upgrade .
-        fi
+        make frontend
 
     - name: Build Catalyst Runtime (latest)
       if: ${{ inputs.lightning == 'latest' }}
@@ -128,6 +122,7 @@ jobs:
         C_COMPILER=$(which gcc }}) \
         CXX_COMPILER=$(which g++ }}) \
         RT_BUILD_DIR="$(pwd)/runtime-build" \
+        LIGHTNING_GIT_TAG_VALUE=latest_release \
         ENABLE_LIGHTNING_KOKKOS=ON \
         ENABLE_OPENQASM=ON \
         make runtime

diff --git a/Makefile b/Makefile
@@ -82,7 +82,7 @@ all: runtime mlir frontend
 .PHONY: frontend
 frontend:
 	@echo "install Catalyst Frontend"
-	$(PYTHON) -m pip install -e .
+	$(PYTHON) -m pip install -e . --extra-index-url https://test.pypi.org/simple
 	rm -r frontend/PennyLane_Catalyst.egg-info
 
 .PHONY: mlir llvm mhlo enzyme dialects runtime
@@ -202,7 +202,7 @@ coverage-frontend:
 	@echo "Generating coverage report for the frontend"
 	$(ASAN_COMMAND) $(PYTHON) -m pytest frontend/test/pytest $(PARALLELIZE) --cov=catalyst --tb=native --cov-report=$(COVERAGE_REPORT)
 ifeq ($(TEST_BRAKET), NONE)
-	$(ASAN_COMMAND) $(PYTHON) -m pytest frontend/test/async_tests --tb=native --backend=$(TEST_BACKEND) --tb=native 
+	$(ASAN_COMMAND) $(PYTHON) -m pytest frontend/test/async_tests --tb=native --backend=$(TEST_BACKEND) --tb=native
 endif
 
 coverage-runtime:

diff --git a/doc/changelog.md b/doc/changelog.md
@@ -71,6 +71,7 @@ This release contains contributions from (in alphabetical order):
 * Catalyst now supports QJIT compatible `catalyst.vmap` of hybrid programs.
   `catalyst.vmap` offers the vectorization mapping backed by `catalyst.for_loop`.
   [(#497)](https://github.com/PennyLaneAI/catalyst/pull/497)
+  [(#569)](https://github.com/PennyLaneAI/catalyst/pull/569)
 
   For example,
 
@@ -102,6 +103,8 @@ This release contains contributions from (in alphabetical order):
   >>> catalyst.qjit(catalyst.grad(f))(x)
   [1. 0. 1. 0.]
   ```
+* Add support for `GlobalPhase` gate in the runtime.
+  [(#563)](https://github.com/PennyLaneAI/catalyst/pull/563)
 
 * Catalyst no longer relies on a TensorFlow installation for its AutoGraph functionality. Instead,
   the standalone `diastatic-malt` package is used and automatically installed as a dependency.

diff --git a/doc/requirements.txt b/doc/requirements.txt
@@ -25,3 +25,7 @@ m2r2
 mistune==0.8.4
 sphinxext-opengraph==0.9.0
 matplotlib==3.8.0
+
+# Pre-install development lightning wheels
+--extra-index-url https://test.pypi.org/simple/
+pennylane-lightning==0.35.0-dev21
diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py
@@ -46,6 +46,7 @@
     DeviceReleaseOp,
     ExpvalOp,
     ExtractOp,
+    GlobalPhaseOp,
     HamiltonianOp,
     HermitianOp,
     InsertOp,
@@ -174,6 +175,8 @@ def _obs_lowering(aval):
 qdealloc_p.multiple_results = True
 qextract_p = core.Primitive("qextract")
 qinsert_p = core.Primitive("qinsert")
+gphase_p = core.Primitive("gphase")
+gphase_p.multiple_results = True
 qinst_p = core.Primitive("qinst")
 qinst_p.multiple_results = True
 qunitary_p = core.Primitive("qunitary")
@@ -677,6 +680,81 @@ def _qinsert_lowering(
     return InsertOp(qreg_type, qreg_old, qubit, idx=qubit_idx).results
 
 
+#
+# gphase
+#
+@gphase_p.def_abstract_eval
+def _gphase_abstract_eval(
+    *qubits_or_params, op=None, qubits_len: int = 0, params_len: int = 0, ctrl_len: int = 0
+):
+    # The signature here is: (using * to denote zero or more)
+    # qubits*, params*, ctrl_qubits*, ctrl_values*
+    qubits = qubits_or_params[:qubits_len]
+    ctrl_qubits = qubits_or_params[-2 * ctrl_len : -ctrl_len]
+    all_qubits = qubits + ctrl_qubits
+    for idx in range(qubits_len + ctrl_len):
+        qubit = all_qubits[idx]
+        assert isinstance(qubit, AbstractQbit)
+    return (AbstractQbit(),) * (qubits_len + ctrl_len)
+
+
+@qinst_p.def_impl
+def _gphase_abstract_eval(
+    *qubits_or_params, op=None, qubits_len: int = 0, params_len: int = 0, ctrl_len: int = 0
+):
+    """Not implemented"""
+    raise NotImplementedError()
+
+
+def _gphase_lowering(
+    jax_ctx: mlir.LoweringRuleContext,
+    *qubits_or_params,
+    op=None,
+    qubits_len: int = 0,
+    params_len: int = 0,
+    ctrl_len: int = 0,
+):
+    ctx = jax_ctx.module_context.context
+    ctx.allow_unregistered_dialects = True
+
+    qubits = qubits_or_params[:qubits_len]
+    params = qubits_or_params[qubits_len : qubits_len + params_len]
+    ctrl_qubits = qubits_or_params[qubits_len + params_len : qubits_len + params_len + ctrl_len]
+    ctrl_values = qubits_or_params[qubits_len + params_len + ctrl_len :]
+
+    float_params = []
+    assert 1 == len(params), "Only one param in GlobalPhase"
+    for p in params:
+        if ir.RankedTensorType.isinstance(p.type) and ir.RankedTensorType(p.type).shape == []:
+            baseType = ir.RankedTensorType(p.type).element_type
+
+        if not ir.F64Type.isinstance(baseType):
+            baseType = ir.F64Type.get()
+            resultTensorType = ir.RankedTensorType.get((), baseType)
+            p = ConvertOp(resultTensorType, p).results
+
+        p = TensorExtractOp(baseType, p, []).result
+
+        assert ir.F64Type.isinstance(
+            p.type
+        ), "Only scalar double parameters are allowed for quantum gates!"
+
+        float_params.append(p)
+
+    ctrl_values_i1 = []
+    for v in ctrl_values:
+        p = TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result
+        ctrl_values_i1.append(p)
+
+    GlobalPhaseOp(
+        params=float_params[0],
+        out_ctrl_qubits=[qubit.type for qubit in ctrl_qubits],
+        in_ctrl_qubits=ctrl_qubits,
+        in_ctrl_values=ctrl_values_i1,
+    )
+    return qubits + ctrl_qubits
+
+
 #
 # qinst
 #
@@ -1576,6 +1654,7 @@ def _adjoint_lowering(
 mlir.register_lowering(qextract_p, _qextract_lowering)
 mlir.register_lowering(qinsert_p, _qinsert_lowering)
 mlir.register_lowering(qinst_p, _qinst_lowering)
+mlir.register_lowering(gphase_p, _gphase_lowering)
 mlir.register_lowering(qunitary_p, _qunitary_lowering)
 mlir.register_lowering(qmeasure_p, _qmeasure_lowering)
 mlir.register_lowering(compbasis_p, _compbasis_lowering)

diff --git a/frontend/catalyst/jax_tracer.py b/frontend/catalyst/jax_tracer.py
@@ -57,6 +57,7 @@
     counts_p,
     expval_p,
     func_p,
+    gphase_p,
     hamiltonian_p,
     hermitian_p,
     mlir_fn_cache,
@@ -116,7 +117,7 @@ def _eval_jaxpr(*args):
 # Take care when adding primitives to this set in order to avoid introducing a quadratic number of
 # edges to the jaxpr equation graph in ``sort_eqns()``. Each equation with a primitive in this set
 # is constrained to occur before all subsequent equations in the quantum operations trace.
-FORCED_ORDER_PRIMITIVES = {qdevice_p}
+FORCED_ORDER_PRIMITIVES = {qdevice_p, gphase_p}
 
 PAULI_NAMED_MAP = {
     "I": "Identity",
@@ -435,6 +436,17 @@ def _bind_native_controlled_op(qrp, op, controlled_wires, controlled_values):
             )
             qrp.insert(op.wires, qubits2[: len(qubits)])
             qrp.insert(controlled_wires, qubits2[len(qubits) :])
+        elif isinstance(op, qml.GlobalPhase):
+            qubits = qrp.extract(op.wires)
+            controlled_qubits = qrp.extract(controlled_wires)
+            qubits2 = gphase_p.bind(
+                *[*qubits, *op.parameters, *controlled_qubits, *controlled_values],
+                qubits_len=len(qubits),
+                params_len=len(op.parameters),
+                ctrl_len=len(controlled_qubits),
+            )
+            qrp.insert(op.wires, qubits2[: len(qubits)])
+            qrp.insert(controlled_wires, qubits2[len(qubits) :])
         else:
             qubits = qrp.extract(op.wires)
             controlled_qubits = qrp.extract(controlled_wires)

diff --git a/frontend/catalyst/pennylane_extensions.py b/frontend/catalyst/pennylane_extensions.py
@@ -2229,10 +2229,7 @@ def postcircuit(y, x, z):
         in the output. ``out_axes`` is subject to the same modes as well.
     """
 
-    # Dispatch to jax.vmap when it is called outside qjit.
-    if not EvaluationContext.is_tracing():
-        return jax.vmap(fn, in_axes, out_axes)
-
+    # Check the validity of in_axes and out_axes
     if not all(isinstance(l, int) for l in tree_leaves(in_axes)):
         raise ValueError(
             "Invalid 'in_axes'; it must be an int or a tuple of PyTrees with integer leaves, "
@@ -2248,6 +2245,10 @@ def postcircuit(y, x, z):
     def batched_fn(*args, **kwargs):
         """Vectorization wrapper around the hybrid program using catalyst.for_loop"""
 
+        # Dispatch to jax.vmap when it is called outside qjit.
+        if not EvaluationContext.is_tracing():
+            return jax.vmap(fn, in_axes, out_axes)(*args, **kwargs)
+
         args_flat, args_tree = tree_flatten(args)
         in_axes_flat, _ = tree_flatten(in_axes, is_leaf=lambda x: x is None)