From 46025bb148064f7b9ba6ff12b23928e67436b8b9 Mon Sep 17 00:00:00 2001 From: Ali Asadi <10773383+maliasadi@users.noreply.github.com> Date: Tue, 30 Jun 2026 18:16:17 -0400 Subject: [PATCH 1/3] Add abstract hamiltonian mlir def --- mlir/include/Quantum/IR/QuantumOps.td | 70 +++++++++++++++++++++++++ mlir/include/Quantum/IR/QuantumTypes.td | 10 ++++ 2 files changed, 80 insertions(+) diff --git a/mlir/include/Quantum/IR/QuantumOps.td b/mlir/include/Quantum/IR/QuantumOps.td index 5b7db9bebf..3f0fb5a82b 100644 --- a/mlir/include/Quantum/IR/QuantumOps.td +++ b/mlir/include/Quantum/IR/QuantumOps.td @@ -1148,6 +1148,76 @@ def HamiltonianOp : Observable_Op<"hamiltonian"> { // ----- +def AbstractHamiltonianOp : Quantum_Op<"abs_hamiltonian"> { + let summary = "Construct a representation-agnostic Hamiltonian handle."; + let description = [{ + The `quantum.abs_hamiltonian` operation builds an opaque, representation-agnostic handle + to a Hamiltonian (`!quantum.abs_hamiltonian`). The numeric data of the Hamiltonian (e.g. + Pauli coefficients, or the core/leaf factor tensors of a compressed double-factorized + Hamiltonian) flows in as a variadic list of ranked tensor or memref values, while the + representation and its static structure ride along as attributes: + + * `kind` -- the representation discriminator: `"pauli"`, `"lcu"` or `"cdf"`. + * `structure` -- representation-specific static metadata (Pauli words, factor ranks, ...). + + Example: + + ```mlir + func.func @foo(%coeffs: tensor<2xf64>) { + %h = quantum.abs_hamiltonian(%coeffs : tensor<2xf64>) + {kind = "pauli", structure = {words = ["XX", "Z"]}} : !quantum.abs_hamiltonian + func.return + } + ``` + }]; + + let arguments = (ins + Variadic>:$data, + StrAttr:$kind, + DefaultValuedAttr:$structure + ); + + let results = (outs + AbstractHamiltonianType:$hamiltonian + ); + + let assemblyFormat = [{ + `(` $data `:` type($data) `)` attr-dict `:` type(results) + }]; +} + +def CustomHOp : Quantum_Op<"custom_h"> { + let summary = "Apply a custom Hamiltonian-based operation."; + let description = [{ + The `quantum.custom_h` operation consumes a representation-agnostic Hamiltonian handle + (`!quantum.abs_hamiltonian`) together with an evolution `time` value and a pair of static + integer attributes `n` and `m`. It is the lowering target of the frontend `CustomHOp` + operation. + + Example: + + ```mlir + func.func @foo(%h: !quantum.abs_hamiltonian, %time: tensor) { + quantum.custom_h %h, %time {n = 50 : i64, m = 1 : i64} : tensor + func.return + } + ``` + }]; + + let arguments = (ins + AbstractHamiltonianType:$hamiltonian, + AnyTypeOf<[F64, AnyRankedTensor, AnyMemRef]>:$time, + I64Attr:$n, + I64Attr:$m + ); + + let assemblyFormat = [{ + $hamiltonian `,` $time attr-dict `:` type($time) + }]; +} + +// ----- + class Measurement_Op traits = []> : Quantum_Op; diff --git a/mlir/include/Quantum/IR/QuantumTypes.td b/mlir/include/Quantum/IR/QuantumTypes.td index 77a2c02923..89bea21c40 100644 --- a/mlir/include/Quantum/IR/QuantumTypes.td +++ b/mlir/include/Quantum/IR/QuantumTypes.td @@ -37,6 +37,16 @@ def QuregType : Quantum_Type<"Qureg", "reg"> { def ObservableType : Quantum_Type<"Observable", "obs"> { let summary = "A quantum observable for use in measurements."; } +def AbstractHamiltonianType : Quantum_Type<"AbstractHamiltonian", "abs_hamiltonian"> { + let summary = "A representation-agnostic handle to a Hamiltonian."; + let description = [{ + An opaque, compile-time handle to a Hamiltonian, independent of its underlying + representation (Pauli-word, linear-combination-of-unitaries, compressed + double-factorized, ...). It carries no payload itself: the numeric data flows alongside + as ordinary tensor/memref values, and the structure rides as attributes on the producing + operation. This mirrors how `!quantum.obs` is used as an opaque observable token. + }]; +} def ResultType : Quantum_Type<"Result", "res"> { let summary = "A quantum measurement result."; } From 79a0c6dc3fbc19db5f0757f377c0703d7f006eb6 Mon Sep 17 00:00:00 2001 From: Ali Asadi <10773383+maliasadi@users.noreply.github.com> Date: Tue, 7 Jul 2026 17:09:14 -0400 Subject: [PATCH 2/3] Add mlir tests --- .../test/Quantum/AbstractHamiltonianTest.mlir | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 mlir/test/Quantum/AbstractHamiltonianTest.mlir diff --git a/mlir/test/Quantum/AbstractHamiltonianTest.mlir b/mlir/test/Quantum/AbstractHamiltonianTest.mlir new file mode 100644 index 0000000000..5533f9cccf --- /dev/null +++ b/mlir/test/Quantum/AbstractHamiltonianTest.mlir @@ -0,0 +1,44 @@ +// Copyright 2026 Xanadu Quantum Technologies Inc. + +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at + +// http://www.apache.org/licenses/LICENSE-2.0 + +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// RUN: quantum-opt --split-input-file %s | FileCheck %s + +// CHECK-LABEL: func.func @pauli +func.func @pauli(%coeffs: tensor<2xf64>, %time: tensor) { + // CHECK: [[h:%.+]] = quantum.abs_hamiltonian(%arg0 : tensor<2xf64>) {kind = "pauli", structure = {words = ["XX", "Z"]}} : !quantum.abs_hamiltonian + %h = quantum.abs_hamiltonian(%coeffs : tensor<2xf64>) {kind = "pauli", structure = {words = ["XX", "Z"]}} : !quantum.abs_hamiltonian + // CHECK: quantum.custom_h [[h]], %arg1 {m = 1 : i64, n = 50 : i64} : tensor + quantum.custom_h %h, %time {n = 50 : i64, m = 1 : i64} : tensor + func.return +} + +// ----- + +// CHECK-LABEL: func.func @lcu +func.func @lcu(%coeffs: tensor<2xf64>, %time: tensor) { + // CHECK: quantum.abs_hamiltonian(%arg0 : tensor<2xf64>) {kind = "lcu" + %h = quantum.abs_hamiltonian(%coeffs : tensor<2xf64>) {kind = "lcu", structure = {terms = ["X(0)", "Z(0)"]}} : !quantum.abs_hamiltonian + quantum.custom_h %h, %time {n = 50 : i64, m = 1 : i64} : tensor + func.return +} + +// ----- + +// CHECK-LABEL: func.func @cdf +func.func @cdf(%core: tensor<2x2x2x3x3xf64>, %leaf: tensor<2x2x3x3xf64>, %nuc: tensor, %time: tensor) { + // CHECK: quantum.abs_hamiltonian(%arg0, %arg1, %arg2 : tensor<2x2x2x3x3xf64>, tensor<2x2x3x3xf64>, tensor) {kind = "cdf", structure = {L = 2 : i64, M = 2 : i64, N = 3 : i64}} : !quantum.abs_hamiltonian + %h = quantum.abs_hamiltonian(%core, %leaf, %nuc : tensor<2x2x2x3x3xf64>, tensor<2x2x3x3xf64>, tensor) {kind = "cdf", structure = {L = 2 : i64, M = 2 : i64, N = 3 : i64}} : !quantum.abs_hamiltonian + quantum.custom_h %h, %time {n = 50 : i64, m = 1 : i64} : tensor + func.return +} From 010018169ed1f3d34157e6ebc8b4357b9e1393d5 Mon Sep 17 00:00:00 2001 From: Ali Asadi <10773383+maliasadi@users.noreply.github.com> Date: Tue, 7 Jul 2026 18:03:38 -0400 Subject: [PATCH 3/3] Add the python interface for abs_ham mlir definition --- frontend/catalyst/abstract_hamiltonian.py | 190 ++++++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 frontend/catalyst/abstract_hamiltonian.py diff --git a/frontend/catalyst/abstract_hamiltonian.py b/frontend/catalyst/abstract_hamiltonian.py new file mode 100644 index 0000000000..640b159df0 --- /dev/null +++ b/frontend/catalyst/abstract_hamiltonian.py @@ -0,0 +1,190 @@ +# Copyright 2026 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +A representation-agnostic Hamiltonian abstraction for Catalyst. + +This module is the frontend (PennyLane) half of the abstract-Hamiltonian feature. It defines +a single canonical container, :class:`AbstractHamiltonian`, that can hold any of the supported +Hamiltonian representations, and :func:`convert_h`, which normalizes the various user-facing +inputs into that container. + +The design follows the same "opaque handle + numeric leaves + static structure" split that +Catalyst already uses for observables (see ``AbstractObs`` / ``HamiltonianOp`` in +``catalyst.jax_primitives``): + +* the *numeric leaves* (coefficients, factor tensors, ...) flow through tracing / lowering as + ordinary array values, and +* the *structure* (which representation, Pauli words, tensor ranks, ...) rides along as static, + hashable metadata. + +Because :class:`AbstractHamiltonian` is registered as a JAX pytree whose children are exactly +the numeric leaves and whose auxiliary data is the (hashable) structure, an instance can be +passed as a runtime argument to a ``qjit`` / captured ``QNode``: the leaves are traced while the +structure is preserved as pytree metadata. + +Supported representations (the ``kind`` discriminator): + +* ``"pauli"`` -- a mapping of Pauli words to coefficients, e.g. ``{"XX": 50, "Z": 51}``. +* ``"lcu"`` -- a linear combination of operators, e.g. ``qp.dot(coeffs, ops)``. +* ``"cdf"`` -- a compressed double-factorized Hamiltonian given by ``core_tensors``, + ``leaf_tensors`` and a scalar ``nuc_constant``. +""" + +import numpy as np +import pennylane as qp +from jax.tree_util import register_pytree_node_class +from pennylane.operation import Operator + +__all__ = ("AbstractHamiltonian", "convert_h", "CustomHOp") + + +@register_pytree_node_class +class AbstractHamiltonian: + """Canonical, representation-agnostic container for a Hamiltonian. + + Args: + kind (str): the representation discriminator, one of ``"pauli"``, ``"lcu"`` or ``"cdf"``. + leaves (Sequence): the numeric leaves of the Hamiltonian (arrays). These are the values + that are traced / lowered. + structure (tuple): static, hashable metadata describing the structure. Stored as a tuple + of ``(key, value)`` pairs so that it can be used as JAX pytree auxiliary data. + + The class is a JAX pytree: ``leaves`` are the children and ``(kind, structure)`` is the + auxiliary data, so an instance can be passed as a runtime argument to a captured ``QNode``. + """ + + def __init__(self, kind, leaves, structure): + self.kind = kind + self.leaves = tuple(leaves) + self.structure = tuple(structure) + + @property + def structure_dict(self): + """Return the structure as a plain dictionary for convenient inspection.""" + return dict(self.structure) + + def tree_flatten(self): + """Children are the numeric leaves; aux data is the (static) kind and structure.""" + return (self.leaves, (self.kind, self.structure)) + + @classmethod + def tree_unflatten(cls, aux_data, children): + """Rebuild from traced leaves and static aux data.""" + kind, structure = aux_data + return cls(kind, children, structure) + + def __repr__(self): + return f"AbstractHamiltonian(kind={self.kind!r}, structure={self.structure_dict!r})" + + +def _convert_pauli(mapping): + """Normalize a ``{pauli_word: coefficient}`` mapping into an AbstractHamiltonian.""" + words = tuple(mapping.keys()) + coeffs = np.array([mapping[w] for w in words], dtype=float) + return AbstractHamiltonian("pauli", (coeffs,), (("words", words),)) + + +def _convert_cdf(mapping): + """Normalize a compressed double-factorized dict into an AbstractHamiltonian. + + Expects keys ``"core_tensors"`` (shape ``(L, M, M, N, N)``), ``"leaf_tensors"`` (shape + ``(L, M, N, N)``) and a scalar ``"nuc_constant"``. + """ + core = np.asarray(mapping["core_tensors"], dtype=float) + leaf = np.asarray(mapping["leaf_tensors"], dtype=float) + nuc = np.asarray(mapping["nuc_constant"], dtype=float) + L, M, _, N, _ = core.shape + structure = (("L", int(L)), ("M", int(M)), ("N", int(N))) + return AbstractHamiltonian("cdf", (core, leaf, nuc), structure) + + +def _convert_lcu(op): + """Normalize a PennyLane operator (e.g. ``qp.dot(coeffs, ops)``) into an AbstractHamiltonian. + + The coefficients are the numeric leaf; the per-term operator descriptions are stored as + static structure. + """ + coeffs, ops = op.terms() + coeffs = np.array(coeffs, dtype=float) + terms = tuple(str(o) for o in ops) + return AbstractHamiltonian("lcu", (coeffs,), (("terms", terms),)) + + +def convert_h(H): + """Normalize a user-supplied Hamiltonian into the canonical :class:`AbstractHamiltonian`. + + Supported inputs: + + * ``dict`` with ``"core_tensors"`` -> a compressed double-factorized (``"cdf"``) Hamiltonian. + * any other ``dict`` -> a Pauli-word (``"pauli"``) Hamiltonian, mapping words to coefficients. + * a PennyLane operator exposing ``.terms()`` (e.g. ``qp.dot(coeffs, ops)``) -> an ``"lcu"`` + Hamiltonian. + + Args: + H: the Hamiltonian in one of the supported representations. + + Returns: + AbstractHamiltonian: the canonical container. + + Raises: + TypeError: if ``H`` is not a supported representation. + """ + if isinstance(H, AbstractHamiltonian): + return H + if isinstance(H, dict): + if "core_tensors" in H: + return _convert_cdf(H) + return _convert_pauli(H) + if isinstance(H, Operator) or hasattr(H, "terms"): + return _convert_lcu(H) + raise TypeError(f"Unsupported Hamiltonian representation: {type(H)!r}") + + +class CustomHOp(Operator): + """A representation-agnostic Hamiltonian-evolution operation. + + This operation consumes an :class:`AbstractHamiltonian` together with an evolution ``time`` + and a pair of integer hyperparameters. The Hamiltonian's numeric leaves become the operation's + trainable/traced data, while its ``kind`` and ``structure`` ride along as static + hyperparameters. This mirrors how Catalyst's ``HamiltonianOp`` carries coefficients as SSA + tensor values and structure as attributes. + + Args: + H (AbstractHamiltonian): the Hamiltonian to evolve under. + time: the evolution time (traced). + n (int): a static integer hyperparameter. + m (int): a static integer hyperparameter. + wires: the wires the operation acts on. + """ + + def __init__(self, H, time, n, m, wires): + H = convert_h(H) + self._hyperparameters = { + "kind": H.kind, + "structure": H.structure, + "n": n, + "m": m, + } + super().__init__(*H.leaves, time, wires=wires) + + @property + def kind(self): + """The Hamiltonian representation discriminator.""" + return self.hyperparameters["kind"] + + @property + def structure(self): + """The static structure metadata as a dictionary.""" + return dict(self.hyperparameters["structure"])