Repository: PennyLaneAI/pennylane
Pull request #4673 (2023-10-13 15:20:50)

`apply_operation` kernel for `MultiControlledX` on many wires
**Context:**
Currently, `default.qubit` applies the operation `MultiControlledX` by building a dense matrix.
This is rather expensive for large numbers of wires.

**Description of the Change:**
This PR adds a dispatch registration for `MultiControlledX` that falls back to `einsum` for less than eight wires in the operation.
For eight wires or more, it instead uses repeated `math.roll`, `math.transpose` and a single concatenation call to achieve the effect of the gate.

**Benefits:**
- Performance improvement on eight or more wires (see plots below).
- Makes more than 13 wires accessible for this operation.

![Screenshot from 2023-10-13 16-56-16](https://github.com/PennyLaneAI/pennylane/assets/20772922/618ec9c0-c691-4eb3-9c44-1e796ed95e2b)

![Screenshot from 2023-10-13 16-54-24](https://github.com/PennyLaneAI/pennylane/assets/20772922/7a7dd53c-e10a-44c2-a988-77d3f930b7cd)


**Possible Drawbacks:**
N/A

**Related GitHub Issues:**


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4673?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified lines are covered by tests :white_check_mark:
> Comparison is base [(`7af90b4`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/7af90b44ab7a81a5fcfc748c1b1d876cda09eed2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63% compared to head [(`5fa68fb`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4673?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.65%.
> Report is 2 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4673      +/-   ##
==========================================
+ Coverage   99.63%   99.65%   +0.02%     
==========================================
  Files         376      376              
  Lines       33678    33693      +15     
==========================================
+ Hits        33554    33576      +22     
+ Misses        124      117       -7     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4673?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/qubit/apply\_operation.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4673?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvYXBwbHlfb3BlcmF0aW9uLnB5) | `99.30% <100.00%> (+0.08%)` | :arrow_up: |

... and [3 files with indirect coverage changes](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4673/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4673?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

dwierichs: Thanks for the reviews @timmysilv  and @mudit2812 .
I fixed the Tensorflow support.
I also timed the other interfaces now, and of course the picture is much more complicated then...
![image](https://github.com/PennyLaneAI/pennylane/assets/20772922/b3d13ceb-eae7-4149-bf17-61b0c14ddab4)

I also am having a benchmarking discussion with @mlxd , and there are even more factors to take into account (like the number of threads).

Therefore my suggestion for this PR would be to only activate the dedicated kernel for a number of operation wires for which the dense matrix becomes too large. Of course this depends on the used hardware as well, but maybe it is sane to limit the matrix size to 2GB/14 wires ? For 15 wires it would be 8GB.
What do you think @mudit2812 , @mlxd, @timmysilv ?
timmysilv: honestly i'm finding it quite tough to read this graph 😅 what is the main point here? it seems like there is no discussion to be had beyond N=12 though
dwierichs: >honestly i'm finding it quite tough to read this graph 😅 what is the main point here?

Sorry, it is quite dense. The blue lines act with the `MultiControlledX` on all `N` wires. We see that this gains us a speedup (numbers below 1) in all interfaces, except for a single data point in the dotted line.
When the gate acts on a subset of more qubits (orange for first half, green for second half), though, it seems that there is some serious slow-down for numpy (dashed) and autograd (long dash-dot).
mudit2812: Happy to approve once the hiccups are resolved :)

dwierichs: Okay, so here is the updated timing:
![image](https://github.com/PennyLaneAI/pennylane/assets/20772922/a944ec78-939d-421e-8862-e05bc6ee432e)
I think I had accidentally included the fall-back in the previous timings, making them invalid.

I suggest the following branching now:
```python
if len(op.wires) < 9:
    # Fall back to einsum/tensordot. based on the standard criteria for that
return dedicated_kernel(op, state, is_state_batched)
```

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5d4b56cfb..8cc6d5f68 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -9,6 +9,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* `default.qubit` no longer uses a dense matrix for `MultiControlledX` for more than 8 operation wires.
+  [(#4673)](https://github.com/PennyLaneAI/pennylane/pull/4673)
+
 * `AmplitudeEmbedding` now also supports batching when used with Tensorflow.
   [(#4818)](https://github.com/PennyLaneAI/pennylane/pull/4818)
 
diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index a0a4e0be7..51cdaa0e7 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -263,6 +263,58 @@ def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False, debugger=Non
     return math.stack([state[sl_0], state_x], axis=control_axes)
 
 
+@apply_operation.register
+def apply_multicontrolledx(
+    op: qml.MultiControlledX, state, is_state_batched: bool = False, debugger=None
+):
+    r"""Apply MultiControlledX to a state with the default einsum/tensordot choice
+    for 8 operation wires or less. Otherwise, apply a custom kernel based on
+    composing transpositions, rolling of control axes and the CNOT logic above."""
+    if len(op.wires) < 9:
+        return _apply_operation_default(op, state, is_state_batched, debugger)
+    ctrl_wires = [w + is_state_batched for w in op.control_wires]
+    # apply x on all control wires with control value 0
+    roll_axes = [
+        w for val, w in zip(op.hyperparameters["control_values"], ctrl_wires) if val == "0"
+    ]
+    for ax in roll_axes:
+        state = math.roll(state, 1, ax)
+
+    orig_shape = math.shape(state)
+    # Move the axes into the order [(batch), other, target, controls]
+    transpose_axes = (
+        np.array(
+            [
+                w - is_state_batched
+                for w in range(len(orig_shape))
+                if w - is_state_batched not in op.wires
+            ]
+            + [op.wires[-1]]
+            + op.wires[:-1].tolist()
+        )
+        + is_state_batched
+    )
+    state = math.transpose(state, transpose_axes)
+
+    # Reshape the state into 3-dimensional array with axes [batch+other, target, controls]
+    state = math.reshape(state, (-1, 2, 2 ** (len(op.wires) - 1)))
+
+    # The part of the state to which we want to apply PauliX is now in the last entry along the
+    # third axis. Extract it, apply the PauliX along the target axis (1), and append a dummy axis
+    state_x = math.roll(state[:, :, -1], 1, 1)[:, :, np.newaxis]
+
+    # Stack the transformed part of the state with the unmodified rest of the state
+    state = math.concatenate([state[:, :, :-1], state_x], axis=2)
+
+    # Reshape into original shape and undo the transposition
+    state = math.transpose(math.reshape(state, orig_shape), np.argsort(transpose_axes))
+
+    # revert x on all "wrong" controls
+    for ax in roll_axes:
+        state = math.roll(state, 1, ax)
+    return state
+
+
 @apply_operation.register
 def apply_snapshot(op: qml.Snapshot, state, is_state_batched: bool = False, debugger=None):
     """Take a snapshot of the state"""


Test patch:
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index 4dde54b40..7b44f57c2 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -800,14 +800,20 @@ class TestLargerOperations:
         ]
     )
 
-    def test_multicontrolledx(self, method):
+    @pytest.mark.parametrize("control_values", ["111", "010", None, "100"])
+    def test_multicontrolledx(self, method, control_values):
         """Tests a four qubit multi-controlled x gate."""
 
-        new_state = method(qml.MultiControlledX(wires=(0, 1, 2, 3)), self.state)
+        op = qml.MultiControlledX(wires=(0, 1, 2, 3), control_values=control_values)
+        new_state = method(op, self.state)
 
         expected_state = np.copy(self.state)
-        expected_state[1, 1, 1, 1] = self.state[1, 1, 1, 0]
-        expected_state[1, 1, 1, 0] = self.state[1, 1, 1, 1]
+        if control_values is None:
+            values = (1, 1, 1)
+        else:
+            values = tuple(map(int, control_values))
+        expected_state[values + (1,)] = self.state[values + (0,)]
+        expected_state[values + (0,)] = self.state[values + (1,)]
 
         assert qml.math.allclose(new_state, expected_state)
 
@@ -825,6 +831,114 @@ class TestLargerOperations:
         assert qml.math.allclose(state_v1, state_v2)
 
 
+class TestMultiControlledXKernel:
+    """Test the specialized kernel for MultiControlledX and its dispatching."""
+
+    # pylint: disable=too-many-arguments
+    @pytest.mark.parametrize(
+        "num_op_wires, num_state_wires, einsum_called, tdot_called",
+        [
+            # state small and matrix huge -> not possible because num_op_wires<=num_state_wires
+            # matrix large -> kernel
+            (9, 9, 0, 0),
+            # matrix large, state huge -> still kernel, not tensordot
+            (9, 9, 0, 0),
+            # matrix tiny, state not huge -> einsum
+            (2, 12, 1, 0),
+            # matrix small, state not huge -> tensordot
+            (5, 12, 0, 1),
+            # matrix tiny, state huge -> tensordot
+            (2, 13, 0, 1),
+            # matrix small, state huge -> tensordot
+            (5, 13, 0, 1),
+        ],
+    )
+    def test_multicontrolledx_dispatching(
+        self, num_op_wires, num_state_wires, einsum_called, tdot_called, mocker
+    ):
+        """Test that apply_multicontrolledx dispatches to the right method and is correct."""
+        np.random.seed(2751)
+        op = qml.MultiControlledX(wires=list(range(num_op_wires)))
+        state = np.random.random([2] * num_state_wires).astype(complex)
+        spies = [mocker.spy(qml.math, "einsum"), mocker.spy(qml.math, "tensordot")]
+        out = apply_operation(op, state, is_state_batched=False, debugger=None)
+        # Compute expected output
+        exp_out = state.copy()
+        idx = (1,) * (num_op_wires - 1)
+        exp_out[idx] = np.roll(exp_out[idx], 1, 0)
+        assert spies[0].call_count == einsum_called
+        assert spies[1].call_count == tdot_called
+        assert np.allclose(out, exp_out)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("batch_dim", [None, 1, 3])
+    def test_with_jax(self, batch_dim):
+        """Test that the custom kernel works with JAX."""
+        from jax import numpy as jnp
+
+        np.random.seed(2751)
+        op = qml.MultiControlledX(wires=[0, 4, 3, 1])
+        state_shape = ([batch_dim] if batch_dim is not None else []) + [2] * 5
+        state = np.random.random(state_shape).astype(complex)
+        jax_state = jnp.array(state)
+        out = apply_operation(op, jax_state, is_state_batched=batch_dim is not None, debugger=None)
+        # Compute expected output
+        exp_out = state.copy()
+        exp_out[..., 1, :, :, 1, 1] = np.roll(exp_out[..., 1, :, :, 1, 1], 1, -2)
+        assert qml.math.allclose(out, exp_out)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("batch_dim", [None, 1, 3])
+    def test_with_tf(self, batch_dim):
+        """Test that the custom kernel works with Tensorflow."""
+        import tensorflow as tf
+
+        np.random.seed(2751)
+        op = qml.MultiControlledX(wires=[0, 4, 3, 1])
+        state_shape = ([batch_dim] if batch_dim is not None else []) + [2] * 5
+        state = np.random.random(state_shape).astype(complex)
+        tf_state = tf.Variable(state)
+        out = apply_operation(op, tf_state, is_state_batched=batch_dim is not None, debugger=None)
+        # Compute expected output
+        exp_out = state.copy()
+        exp_out[..., 1, :, :, 1, 1] = np.roll(exp_out[..., 1, :, :, 1, 1], 1, -2)
+        assert qml.math.allclose(out, exp_out)
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("batch_dim", [None, 1, 3])
+    def test_with_autograd(self, batch_dim):
+        """Test that the custom kernel works with Autograd."""
+        np.random.seed(2751)
+        op = qml.MultiControlledX(wires=[0, 4, 3, 1])
+        state_shape = ([batch_dim] if batch_dim is not None else []) + [2] * 5
+        state = np.random.random(state_shape).astype(complex)
+        ag_state = qml.numpy.array(state)
+        out = apply_operation(op, ag_state, is_state_batched=batch_dim is not None, debugger=None)
+        # Compute expected output
+        exp_out = state.copy()
+        exp_out[..., 1, :, :, 1, 1] = np.roll(exp_out[..., 1, :, :, 1, 1], 1, -2)
+        assert qml.math.allclose(out, exp_out)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("batch_dim", [None, 1, 3])
+    def test_with_torch(self, batch_dim):
+        """Test that the custom kernel works with Torch."""
+        import torch
+
+        np.random.seed(2751)
+        op = qml.MultiControlledX(wires=[0, 4, 3, 1])
+        state_shape = ([batch_dim] if batch_dim is not None else []) + [2] * 5
+        state = np.random.random(state_shape).astype(complex)
+        torch_state = torch.tensor(state, requires_grad=True)
+        out = apply_operation(
+            op, torch_state, is_state_batched=batch_dim is not None, debugger=None
+        )
+        # Compute expected output
+        exp_out = state.copy()
+        exp_out[..., 1, :, :, 1, 1] = np.roll(exp_out[..., 1, :, :, 1, 1], 1, -2)
+        assert qml.math.allclose(out, exp_out)
+
+
 @pytest.mark.tf
 @pytest.mark.parametrize("op", (qml.PauliZ(8), qml.CNOT((5, 6))))
 def test_tf_large_state(op):