Callbacks with parameters.

doc/changelog.md

@@ -4,21 +4,28 @@
 
 * Support for callbacks in Catalyst.
   [(#540)](https://github.com/PennyLaneAI/catalyst/pull/540)
+  [(#596)](https://github.com/PennyLaneAI/catalyst/pull/596)
 
-  Catalyst now supports callbacks without parameters nor return values.
+  Catalyst now supports callbacks with parameters but no return values.
   This is the very first step in supporting callbacks.
-  The following is now possible.
+  The following is now possible:
 
   ```py
   @callback
-  def foo():
-    print("Hello world")
+  def foo(val):
+    print("Hello world", val)
 
   @qjit
   def circuit(*args, **kwargs):
     ...
-    foo()
+    foo(123)
     ...
+
+  ```
+
+  ```pycon
+  >>> circuit()
+  Hello world 123
   ```
 
 * The python layer of the OQC-Catalyst device is now available.

doc/conf.py

@@ -84,6 +84,7 @@ def __getattr__(cls, name):
 MOCK_MODULES = [
     "mlir_quantum",
     "mlir_quantum.runtime",
+    "mlir_quantum.runtime.np_to_memref",
     "mlir_quantum.dialects",
     "mlir_quantum.dialects.arith",
     "mlir_quantum.dialects.tensor",

frontend/catalyst/compiled_functions.py

@@ -23,7 +23,6 @@
 from jax.tree_util import PyTreeDef, tree_flatten, tree_unflatten
 from mlir_quantum.runtime import (
     as_ctype,
-    get_ranked_memref_descriptor,
     make_nd_memref_descriptor,
     make_zero_d_memref_descriptor,
 )
@@ -38,6 +37,7 @@
 from catalyst.utils import wrapper  # pylint: disable=no-name-in-module
 from catalyst.utils.c_template import get_template, mlir_type_to_numpy_type
 from catalyst.utils.filesystem import Directory
+from catalyst.utils.jnp_to_memref import get_ranked_memref_descriptor
 
 
 class SharedObjectManager:

frontend/catalyst/jax_primitives.py

@@ -242,7 +242,7 @@ def _python_callback_lowering(jax_ctx: mlir.LoweringRuleContext, *args, callback
     ctx = jax_ctx.module_context.context
     i64_type = ir.IntegerType.get_signless(64, ctx)
     identifier = ir.IntegerAttr.get(i64_type, callback_id)
-    return PythonCallOp(identifier).results
+    return PythonCallOp(args, identifier).results
 
 
 #

frontend/catalyst/pennylane_extensions.py

@@ -19,6 +19,7 @@
 # pylint: disable=too-many-lines
 
 import copy
+import ctypes
 import inspect
 import numbers
 from collections.abc import Sequence, Sized
@@ -95,6 +96,10 @@
     JaxTracingContext,
 )
 from catalyst.utils.exceptions import DifferentiableCompileError
+from catalyst.utils.jnp_to_memref import (
+    get_ranked_memref_descriptor,
+    ranked_memref_to_numpy,
+)
 from catalyst.utils.runtime import (
     BackendInfo,
     device_get_toml_config,
@@ -2555,6 +2560,42 @@ def _get_batch_size(args_flat, axes_flat, axis_size):
     return batch_size
 
 
+class CallbackClosure:
+    """This is just a class containing data that is important for the callback."""
+
+    def __init__(self, *absargs, **abskwargs):
+        self.absargs = absargs
+        self.abskwargs = abskwargs
+
+    @property
+    def tree_flatten(self):
+        """Flatten args and kwargs."""
+        return tree_flatten((self.absargs, self.abskwargs))
+
+    @property
+    def low_level_sig(self):
+        """Get the memref descriptor types"""
+        flat_params, _ = self.tree_flatten
+        low_level_flat_params = []
+        for param in flat_params:
+            empty_memref_descriptor = get_ranked_memref_descriptor(param)
+            memref_type = type(empty_memref_descriptor)
+            ptr_ty = ctypes.POINTER(memref_type)
+            low_level_flat_params.append(ptr_ty)
+        return low_level_flat_params
+
+    def getArgsAsJAXArrays(self, flat_args):
+        """Get arguments as JAX arrays. Since our integration is mostly compatible with JAX,
+        it is best for the user if we continue with that idea and forward JAX arrays."""
+        jnpargs = []
+        for void_ptr, ty in zip(flat_args, self.low_level_sig):
+            memref_ty = ctypes.cast(void_ptr, ty)
+            nparray = ranked_memref_to_numpy(memref_ty)
+            jnparray = jnp.asarray(nparray)
+            jnpargs.append(jnparray)
+        return jnpargs
+
+
 def callback(func):
     """Decorator that will correctly pass the signature as arguments to the callback
     implementation.
@@ -2587,15 +2628,20 @@ def callback_implementation(
     """
 
     flat_args, in_tree = tree_flatten((args, kwargs))
+    metadata = CallbackClosure(args, kwargs)
+
+    def _flat_callback(flat_args):
+        """Each flat_arg is a pointer.
+
+        It is a pointer to a memref object.
+        To find out which element type it has, we use the signature obtained previously.
+        """
+        jnpargs = metadata.getArgsAsJAXArrays(flat_args)
 
-    def _flat_callback(*flat_args):
-        """This function packages flat arguments back into the shapes expected by the function."""
-        _args, _kwargs = tree_unflatten(in_tree, flat_args)
-        assert not _args, "Args are not yet expected here."
-        assert not _kwargs, "Kwargs are not yet supported here."
-        return tree_leaves(cb())
+        args, kwargs = tree_unflatten(in_tree, jnpargs)
+        return tree_leaves(cb(*args, **kwargs))
 
-    # TODO(@erick-xanadu): Change back once we support arguments.
+    # TODO(@erick-xanadu): Change back once we support return values.
     # I am leaving this as a to-do because otherwise the coverage will complain.
     # results_aval = tree_map(lambda x: jax.core.ShapedArray(x.shape, x.dtype), result_shape_dtypes)
     results_aval = []

frontend/catalyst/utils/jnp_to_memref.py

@@ -0,0 +1,74 @@
+#  This file is essentially a re-implementation of np_to_memref.py
+#  But we make sure that all the functions also work with JAX's numpy.
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#  This file contains functions to convert between Memrefs and NumPy arrays and vice-versa.
+
+"""
+This file is a wrapper around MLIR's np_to_memref to allow for abstract types and JAX arrays
+to be converted to and from memrefs.
+"""
+
+import numpy as np
+from mlir_quantum.runtime import as_ctype
+from mlir_quantum.runtime import (
+    get_ranked_memref_descriptor as mlir_get_ranked_memref_descriptor,
+)
+from mlir_quantum.runtime import ranked_memref_to_numpy as mlir_ranked_memref_to_numpy
+from mlir_quantum.runtime.np_to_memref import (
+    make_nd_memref_descriptor,
+    make_zero_d_memref_descriptor,
+    move_aligned_ptr_by_offset,
+    to_numpy,
+)
+
+from catalyst.jax_extras import DynamicJaxprTracer, ShapedArray
+
+
+def get_ranked_memref_descriptor_from_shaped_array(array: ShapedArray):
+    """Get a ranked memref descriptor from a shaped array.
+
+    Unlike MLIR's implementation, all values are left uninitialized.
+    This is because the values are not yet known. We only have a description
+    of the type.
+    """
+
+    ctp = as_ctype(array.dtype)
+    if array.ndim == 0:
+        return make_zero_d_memref_descriptor(ctp)()
+
+    return make_nd_memref_descriptor(array.ndim, ctp)()
+
+
+def get_ranked_memref_descriptor(array):
+    """Wrapper around MLIR's get_ranked_memref_descriptor."""
+
+    if isinstance(array, DynamicJaxprTracer):
+        array = array.aval
+
+    if isinstance(array, (int, float, bool, complex)):
+        # This is necessary for keyword arguments
+        array = np.array(array)
+
+    if isinstance(array, ShapedArray):
+        # If input is ShapedArray
+        return get_ranked_memref_descriptor_from_shaped_array(array)
+
+    # Use default implementation from MLIR's library.
+    return mlir_get_ranked_memref_descriptor(array)
+
+
+def ranked_memref_to_numpy(ranked_memref):
+    """Wrapper around MLIR's ranked_memref_to_numpy.
+
+    This wrapper succeeds when the ranked_memref is a scalar tensor.
+    """
+    try:
+        return mlir_ranked_memref_to_numpy(ranked_memref)
+    except AttributeError:
+        # zero dimensional tensor...
+        content_ptr = move_aligned_ptr_by_offset(ranked_memref[0].aligned, ranked_memref[0].offset)
+        np_arr = np.ctypeslib.as_array(content_ptr, shape=[])
+        return to_numpy(np_arr)

frontend/test/pytest/test_callback.py

@@ -21,6 +21,9 @@
 
 @pytest.mark.parametrize("arg", [1, 2, 3])
 def test_callback_no_tracing(arg):
+    """Test that when there's no tracing the behaviour of identity
+    stays the same."""
+
     @callback
     def identity(x):
         return x
@@ -78,3 +81,42 @@ def cir2():
     cir2()
     captured = capsys.readouterr()
     assert captured.out.strip() == "Hello erick"
+
+
+def test_callback_send_param(capsys):
+    """Test callback with parameters no returns"""
+
+    @callback
+    def my_callback(n) -> None:
+        print(n)
+
+    @qml.qjit
+    def cir(n):
+        my_callback(n)
+        return None
+
+    cir(0)
+    captured = capsys.readouterr()
+    assert captured.out.strip() == "0"
+
+
+def test_kwargs(capsys):
+    """Test kwargs returns"""
+
+    @callback
+    def my_callback(**kwargs) -> None:
+        for k, v in kwargs.items():
+            print(k, v)
+
+    @qml.qjit
+    def cir(a, b, c):
+        my_callback(a=a, b=b, c=c, d=3, e=4)
+        return None
+
+    captured = capsys.readouterr()
+    assert captured.out.strip() == ""
+
+    cir(0, 1, 2)
+    captured = capsys.readouterr()
+    for string in ["a 0", "b 1", "c 2", "d 3", "e 4"]:
+        assert string in captured.out

frontend/test/pytest/test_np_to_memref.py

@@ -0,0 +1,61 @@
+# Copyright 2024 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.
+"""Tests memref descriptor functions with JAX's ShapedArrays."""
+
+import ctypes
+
+import jax.numpy as jnp
+import pytest
+from jax.core import ShapedArray
+from mlir_quantum.runtime import (
+    C64,
+    C128,
+    as_ctype,
+    make_nd_memref_descriptor,
+    make_zero_d_memref_descriptor,
+)
+
+from catalyst.utils.jnp_to_memref import get_ranked_memref_descriptor
+
+
+@pytest.mark.parametrize(
+    "inp, exp",
+    [
+        (jnp.dtype(jnp.float64), ctypes.c_double),
+        (jnp.dtype(jnp.float32), ctypes.c_float),
+        (jnp.dtype(jnp.int64), ctypes.c_long),
+        (jnp.dtype(jnp.bool_), ctypes.c_bool),
+        (jnp.dtype(jnp.complex128), C128),
+        (jnp.dtype(jnp.complex64), C64),
+    ],
+)
+def test_as_ctype(inp, exp):
+    """Tests that JAX's dtypes behave the same as numpy's dtypes"""
+    obs = as_ctype(inp)
+    assert exp == obs
+
+
+@pytest.mark.parametrize(
+    "inp, exp",
+    [
+        (1, make_zero_d_memref_descriptor(ctypes.c_long)),
+        (ShapedArray([], float), make_zero_d_memref_descriptor(ctypes.c_double)),
+        (ShapedArray([1], float), make_nd_memref_descriptor(1, ctypes.c_double)),
+        (ShapedArray([2, 2], float), make_nd_memref_descriptor(2, ctypes.c_double)),
+    ],
+)
+def test_get_ranked_memref_descriptor(inp, exp):
+    """Tests that the structure has the expected fields."""
+    obs = get_ranked_memref_descriptor(inp)
+    assert exp._fields_ == obs._fields_  # pylint: disable=protected-access

mlir/include/Catalyst/IR/CatalystOps.td

@@ -114,11 +114,12 @@ def PythonCallOp: Catalyst_Op<"pycallback",
   }];
 
   let arguments = (ins
+        Variadic<AnyTypeOf<[AnyRankedTensor, MemRefOf<[AnyType]>]>>:$inputs,
         I64Attr: $identifier
   );
 
   let assemblyFormat = [{
-    attr-dict
+    `(` $inputs `)` attr-dict `:` functional-type(operands, results)
   }];
 }
 

mlir/lib/Catalyst/Transforms/BufferizationPatterns.cpp

@@ -82,13 +82,26 @@ struct BufferizeCustomCallOp : public OpConversionPattern<CustomCallOp> {
     }
 };
 
+struct BufferizePythonCallOp : public OpConversionPattern<PythonCallOp> {
+    using OpConversionPattern::OpConversionPattern;
+
+    LogicalResult matchAndRewrite(PythonCallOp op, OpAdaptor adaptor,
+                                  ConversionPatternRewriter &rewriter) const override
+    {
+        rewriter.create<PythonCallOp>(op.getLoc(), adaptor.getOperands(), adaptor.getIdentifier());
+        rewriter.eraseOp(op);
+        return success();
+    }
+};
+
 } // namespace
 
 namespace catalyst {
 
 void populateBufferizationPatterns(TypeConverter &typeConverter, RewritePatternSet &patterns)
 {
     patterns.add<BufferizeCustomCallOp>(typeConverter, patterns.getContext());
+    patterns.add<BufferizePythonCallOp>(typeConverter, patterns.getContext());
     patterns.add<BufferizePrintOp>(typeConverter, patterns.getContext());
 }
 

mlir/lib/Catalyst/Transforms/catalyst_bufferize.cpp

@@ -50,6 +50,8 @@ struct CatalystBufferizationPass : impl::CatalystBufferizationPassBase<CatalystB
             [&](PrintOp op) { return typeConverter.isLegal(op); });
         target.addDynamicallyLegalOp<CustomCallOp>(
             [&](CustomCallOp op) { return typeConverter.isLegal(op); });
+        target.addDynamicallyLegalOp<PythonCallOp>(
+            [&](PythonCallOp op) { return typeConverter.isLegal(op); });
 
         if (failed(applyPartialConversion(getOperation(), target, std::move(patterns)))) {
             signalPassFailure();