Add get_constant_from_ssa function for getting the concrete value of an xDSL variable (#8514)

There are currently two ways of defining constant values -
`arith.constant`, and `stablehlo.constant` - and we need to deal with
them in different ways when trying to extract the concrete value. This
PR adds a function called `get_constant_from_ssa`, which takes an SSA
value as input, and returns its concrete value if it is a constant, else
`None`.

Notes:
* Added a `utils.py` file to the `python_compiler` submodule
* The function can currently only extract scalar values, and explicitly
returns `None` if the value is not scalar, even if it is a constant.

[sc-100957]

---------

Co-authored-by: Andrija Paurevic <[email protected]>

Author: mudit2812

doc/releases/changelog-dev.md

@@ -11,6 +11,10 @@
 
 <h3>Improvements 🛠</h3>
 
+* A new `qml.compiler.python_compiler.utils` submodule has been added, containing general-purpose utilities for
+  working with xDSL. This includes a function that extracts the concrete value of scalar, constant SSA values.
+  [(#8514)](https://github.com/PennyLaneAI/pennylane/pull/8514)
+
 * Added a keyword argument ``recursive`` to ``qml.transforms.cancel_inverses`` that enables
   recursive cancellation of nested pairs of mutually inverse gates. This makes the transform
   more powerful, because it can cancel larger blocks of inverse gates without having to scan
@@ -214,6 +218,7 @@ Astral Cai,
 Marcus Edwards,
 Lillian Frederiksen,
 Christina Lee,
+Mudit Pandey,
 Shuli Shu,
 Jay Soni,
 David Wierichs,

pennylane/compiler/python_compiler/utils.py

@@ -0,0 +1,77 @@
+# Copyright 2025 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.
+
+"""General purpose utilities to use with xDSL."""
+
+from numbers import Number
+
+from xdsl.dialects.arith import ConstantOp as arithConstantOp
+from xdsl.dialects.builtin import ComplexType, ShapedType
+from xdsl.dialects.tensor import ExtractOp as tensorExtractOp
+from xdsl.ir import SSAValue
+
+from .dialects.stablehlo import ConstantOp as hloConstantOp
+
+
+def get_constant_from_ssa(value: SSAValue) -> Number | None:
+    """Return the concrete value corresponding to an SSA value if it is a numerical constant.
+
+    .. note::
+
+        This function currently only returns constants if they are scalar. For non-scalar
+        constants, ``None`` will be returned.
+
+    Args:
+        value (xdsl.ir.SSAValue): the SSA value to check
+
+    Returns:
+        Number or None: If the value corresponds to a constant, its concrete value will
+        be returned, else ``None``.
+    """
+
+    # If the value has a shape, we can assume that it is not scalar. We check
+    # this because constant-like operations can return container types. This includes
+    # arith.constant, which may return containers, and stablehlo.constant, which
+    # always returns a container.
+    if not isinstance(value.type, ShapedType):
+        owner = value.owner
+
+        if isinstance(owner, arithConstantOp):
+            const_attr = owner.value
+            return const_attr.value.data
+
+        # Constant-like operations can also create scalars by returning rank 0 tensors.
+        # In this case, the owner of a scalar value should be a tensor.extract, which
+        # uses the aforementioned rank 0 constant tensor as input.
+        if isinstance(owner, tensorExtractOp):
+            tensor_ = owner.tensor
+            if (
+                len(owner.indices) == 0
+                and len(tensor_.type.shape) == 0
+                and isinstance(tensor_.owner, (arithConstantOp, hloConstantOp))
+            ):
+                dense_attr = tensor_.owner.value
+                # We know that the tensor has shape (). Dense element attributes store
+                # their data as a sequence. For a scalar, this will be a sequence with
+                # a single element.
+                val = dense_attr.get_values()[0]
+                if isinstance(tensor_.type.element_type, ComplexType):
+                    # If the dtype is complex, the value will be a 2-tuple containing
+                    # the real and imaginary components of the number rather than a
+                    # Python complex number
+                    val = val[0] + 1j * val[1]
+
+                return val
+
+    return None

tests/python_compiler/test_xdsl_utils.py

@@ -0,0 +1,120 @@
+# Copyright 2025 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.
+
+"""Unit tests for xDSL utilities."""
+
+import pytest
+
+pytestmark = pytest.mark.external
+xdsl = pytest.importorskip("xdsl")
+
+# pylint: disable=wrong-import-position
+from xdsl.dialects import arith, builtin, tensor, test
+
+from pennylane.compiler.python_compiler.dialects.stablehlo import ConstantOp as hloConstantOp
+from pennylane.compiler.python_compiler.utils import get_constant_from_ssa
+
+
+class TestGetConstantFromSSA:
+    """Unit tests for ``get_constant_from_ssa``."""
+
+    def test_non_constant(self):
+        """Test that ``None`` is returned if the input is not a constant."""
+        val = test.TestOp(result_types=(builtin.Float64Type(),)).results[0]
+        assert get_constant_from_ssa(val) is None
+
+    @pytest.mark.parametrize(
+        "const, attr_type, dtype",
+        [
+            (11, builtin.IntegerAttr, builtin.IntegerType(64)),
+            (5, builtin.IntegerAttr, builtin.IndexType()),
+            (2.5, builtin.FloatAttr, builtin.Float64Type()),
+        ],
+    )
+    def test_scalar_constant_arith(self, const, attr_type, dtype):
+        """Test that constants created by ``arith.constant`` are returned correctly."""
+        const_attr = attr_type(const, dtype)
+        val = arith.ConstantOp(value=const_attr).results[0]
+
+        assert get_constant_from_ssa(val) == const
+
+    @pytest.mark.parametrize(
+        "const, elt_type",
+        [
+            (11, builtin.IntegerType(64)),
+            (9, builtin.IndexType()),
+            (2.5, builtin.Float64Type()),
+            (-1.1 + 2.3j, builtin.ComplexType(builtin.Float64Type())),
+        ],
+    )
+    @pytest.mark.parametrize("constant_op", [arith.ConstantOp, hloConstantOp])
+    def test_scalar_constant_extracted_from_rank0_tensor(self, const, elt_type, constant_op):
+        """Test that constants created by ``stablehlo.constant`` are returned correctly."""
+        data = const
+        if isinstance(const, complex):
+            # For complex numbers, the number must be split into a 2-tuple containing
+            # the real and imaginary part when initializing a dense elements attr.
+            data = (const.real, const.imag)
+
+        dense_attr = builtin.DenseIntOrFPElementsAttr.from_list(
+            type=builtin.TensorType(element_type=elt_type, shape=()),
+            data=(data,),
+        )
+        tensor_ = constant_op(value=dense_attr).results[0]
+        val = tensor.ExtractOp(tensor=tensor_, indices=[], result_type=elt_type).results[0]
+
+        assert get_constant_from_ssa(val) == const
+
+    def test_tensor_constant_arith(self):
+        """Test that ``None`` is returned if the input is a tensor created by ``arith.constant``."""
+        dense_attr = builtin.DenseIntOrFPElementsAttr.from_list(
+            type=builtin.TensorType(element_type=builtin.Float64Type(), shape=(3,)),
+            data=(1, 2, 3),
+        )
+        val = arith.ConstantOp(value=dense_attr).results[0]
+
+        assert get_constant_from_ssa(val) is None
+
+    def test_tensor_constant_stablehlo(self):
+        """Test that ``None`` is returned if the input is a tensor created by ``stablehlo.constant``."""
+        dense_attr = builtin.DenseIntOrFPElementsAttr.from_list(
+            type=builtin.TensorType(element_type=builtin.Float64Type(), shape=(3,)),
+            data=(1.0, 2.0, 3.0),
+        )
+        val = hloConstantOp(value=dense_attr).results[0]
+
+        assert get_constant_from_ssa(val) is None
+
+    def test_extract_scalar_from_constant_tensor_stablehlo(self):
+        """Test that ``None`` is returned if the input is a scalar constant, but it was extracted
+        from a non-scalar constant."""
+        # Index SSA value to be used for extracting a value from a tensor
+        dummy_index = test.TestOp(result_types=(builtin.IndexType(),)).results[0]
+
+        dense_attr = builtin.DenseIntOrFPElementsAttr.from_list(
+            type=builtin.TensorType(element_type=builtin.Float64Type(), shape=(3,)),
+            data=(1.0, 2.0, 3.0),
+        )
+        tensor_ = hloConstantOp(value=dense_attr).results[0]
+        val = tensor.ExtractOp(
+            tensor=tensor_, indices=[dummy_index], result_type=builtin.Float64Type()
+        ).results[0]
+        # val is a value that we got by indexing into a constant tensor with rank >= 1
+        assert isinstance(val.type, builtin.Float64Type)
+
+        assert get_constant_from_ssa(val) is None
+
+
+if __name__ == "__main__":
+    pytest.main(["-x", __file__])