Add angle-producing math functions to kernels with docs and tests

cudaq/lib/Frontend/nvqpp/ConvertExpr.cpp

@@ -1460,6 +1460,21 @@ bool QuakeBridgeVisitor::visitMathLibFunc(clang::CallExpr *x,
       (funcName == "tan" || funcName == "tanf"))
     return floatOperator(math::TanOp{}, "tan");
 
+  // Handle std::asin
+  if ((isInNamespace(func, "std") || isNotInANamespace(func)) &&
+      (funcName == "asin" || funcName == "asinf"))
+    return floatOperator(math::AsinOp{}, "asin");
+
+  // Handle std::acos
+  if ((isInNamespace(func, "std") || isNotInANamespace(func)) &&
+      (funcName == "acos" || funcName == "acosf"))
+    return floatOperator(math::AcosOp{}, "acos");
+
+  // Handle std::atan
+  if ((isInNamespace(func, "std") || isNotInANamespace(func)) &&
+      (funcName == "atan" || funcName == "atanf"))
+    return floatOperator(math::AtanOp{}, "atan");
+
   // Handle std::exp
   if ((isInNamespace(func, "std") || isNotInANamespace(func)) &&
       (funcName == "exp" || funcName == "expf"))

cudaq/test/AST-Quake/math_functions.cpp

@@ -0,0 +1,70 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2026 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+// RUN: cudaq-quake %s | cudaq-opt | FileCheck %s
+
+#include <cudaq.h>
+#include <cmath>
+
+// Standard math functions are usable inside kernels and lower to the matching
+// MLIR math dialect ops. The angle is a runtime parameter so the calls are not
+// constant-folded away.
+
+struct MathFunctions {
+  void operator()(double x) __qpu__ {
+    cudaq::qubit q;
+    rx(std::sin(x), q);
+    rx(std::cos(x), q);
+    rx(std::tan(x), q);
+    rx(std::asin(x), q);
+    rx(std::acos(x), q);
+    rx(std::atan(x), q);
+    rx(std::sqrt(x), q);
+    rx(std::exp(x), q);
+    rx(std::log(x), q);
+  }
+};
+
+struct MathFunctionsFloat {
+  void operator()(float x) __qpu__ {
+    cudaq::qubit q;
+    rx(sinf(x), q);
+    rx(cosf(x), q);
+    rx(tanf(x), q);
+    rx(asinf(x), q);
+    rx(acosf(x), q);
+    rx(atanf(x), q);
+    rx(sqrtf(x), q);
+    rx(expf(x), q);
+    rx(logf(x), q);
+  }
+};
+
+// CHECK-LABEL: func.func @__nvqpp__mlirgen__MathFunctions
+// CHECK-DAG: math.sin
+// CHECK-DAG: math.cos
+// CHECK-DAG: math.tan
+// CHECK-DAG: math.asin
+// CHECK-DAG: math.acos
+// CHECK-DAG: math.atan
+// CHECK-DAG: math.sqrt
+// CHECK-DAG: math.exp
+// CHECK-DAG: math.log
+// CHECK: quake.rx
+
+// CHECK-LABEL: func.func @__nvqpp__mlirgen__MathFunctionsFloat
+// CHECK-DAG: math.sin
+// CHECK-DAG: math.cos
+// CHECK-DAG: math.tan
+// CHECK-DAG: math.asin
+// CHECK-DAG: math.acos
+// CHECK-DAG: math.atan
+// CHECK-DAG: math.sqrt
+// CHECK-DAG: math.exp
+// CHECK-DAG: math.log
+// CHECK: quake.rx

docs/sphinx/specification/cudaq/kernels.rst

@@ -59,6 +59,12 @@ Kernels can be composed of the following:
   * Classical control flow constructs from the classical language (:code:`if`, :code:`for`, :code:`while`, etc.)
   * Stack variable declarations for supported types. 
   * Arithmetic operations on integer and floating point stack variables
+  * Calls to a defined set of standard mathematical functions on floating-point
+    operands: the real-valued, angle-producing functions :code:`sin`,
+    :code:`cos`, :code:`tan`, :code:`asin`, :code:`acos`, :code:`atan`,
+    :code:`sqrt`, :code:`exp`, and :code:`log` (in C++, the corresponding
+    :code:`std::` functions; in Python, the corresponding :code:`numpy`
+    functions)
   * Coherent conditional execution - :code:`if ( boolExprFromQubitMeasurement ) { x (another_qubit); }` 
   * Syntax for common quantum programming patterns (e.g. compute-action-uncompute).
 

python/cudaq/kernel/ast_bridge.py

@@ -815,7 +815,10 @@ def isArithmeticType(self, type):
             type) or F32Type.isinstance(type) or ComplexType.isinstance(type)
 
     def __isSupportedNumpyFunction(self, id):
-        return id in ['sin', 'cos', 'sqrt', 'ceil', 'exp']
+        return id in [
+            'sin', 'cos', 'tan', 'asin', 'acos', 'atan', 'arcsin', 'arccos',
+            'arctan', 'sqrt', 'ceil', 'exp', 'log'
+        ]
 
     def __isSupportedVectorFunction(self, id):
         return id in ['front', 'back', 'append']
@@ -3535,6 +3538,46 @@ def bodyBuilder(iterVar):
                             return
                         self.pushValue(math.CeilOp(value).result)
                         return
+                    if node.func.attr == 'tan':
+                        if ComplexType.isinstance(value.type):
+                            self.emitFatalError(
+                                f"numpy call ({node.func.attr}) is not "
+                                f"supported for complex numbers", node)
+                            return
+                        self.pushValue(math.TanOp(value).result)
+                        return
+                    if node.func.attr in ('asin', 'arcsin'):
+                        if ComplexType.isinstance(value.type):
+                            self.emitFatalError(
+                                f"numpy call ({node.func.attr}) is not "
+                                f"supported for complex numbers", node)
+                            return
+                        self.pushValue(math.AsinOp(value).result)
+                        return
+                    if node.func.attr in ('acos', 'arccos'):
+                        if ComplexType.isinstance(value.type):
+                            self.emitFatalError(
+                                f"numpy call ({node.func.attr}) is not "
+                                f"supported for complex numbers", node)
+                            return
+                        self.pushValue(math.AcosOp(value).result)
+                        return
+                    if node.func.attr in ('atan', 'arctan'):
+                        if ComplexType.isinstance(value.type):
+                            self.emitFatalError(
+                                f"numpy call ({node.func.attr}) is not "
+                                f"supported for complex numbers", node)
+                            return
+                        self.pushValue(math.AtanOp(value).result)
+                        return
+                    if node.func.attr == 'log':
+                        if ComplexType.isinstance(value.type):
+                            self.emitFatalError(
+                                f"numpy call ({node.func.attr}) is not "
+                                f"supported for complex numbers", node)
+                            return
+                        self.pushValue(math.LogOp(value).result)
+                        return
 
                     self.emitFatalError(
                         f"unsupported NumPy call ({node.func.attr})", node)

python/tests/kernel/test_kernel_float.py

@@ -310,3 +310,133 @@ def check(c: any):
     check([np.float32(np.pi / 2), 0])
     check([1, 0])
     check([np.pi / 2, 0, True])
+
+
+def test_math_functions_float64():
+    """Test that math functions can be used inside kernels and match host numpy (float64)."""
+
+    @cudaq.kernel
+    def f_sin() -> np.float64:
+        return np.sin(np.float64(0.5))
+
+    assert is_close(np.sin(0.5), f_sin())
+
+    @cudaq.kernel
+    def f_cos() -> np.float64:
+        return np.cos(np.float64(0.5))
+
+    assert is_close(np.cos(0.5), f_cos())
+
+    @cudaq.kernel
+    def f_tan() -> np.float64:
+        return np.tan(np.float64(0.5))
+
+    assert is_close(np.tan(0.5), f_tan())
+
+    @cudaq.kernel
+    def f_arcsin() -> np.float64:
+        return np.arcsin(np.float64(0.5))
+
+    assert is_close(np.arcsin(0.5), f_arcsin())
+
+    @cudaq.kernel
+    def f_arccos() -> np.float64:
+        return np.arccos(np.float64(0.5))
+
+    assert is_close(np.arccos(0.5), f_arccos())
+
+    @cudaq.kernel
+    def f_arctan() -> np.float64:
+        return np.arctan(np.float64(0.5))
+
+    assert is_close(np.arctan(0.5), f_arctan())
+
+    @cudaq.kernel
+    def f_sqrt() -> np.float64:
+        return np.sqrt(np.float64(0.5))
+
+    assert is_close(np.sqrt(0.5), f_sqrt())
+
+    @cudaq.kernel
+    def f_exp() -> np.float64:
+        return np.exp(np.float64(0.5))
+
+    assert is_close(np.exp(0.5), f_exp())
+
+    @cudaq.kernel
+    def f_log() -> np.float64:
+        return np.log(np.float64(0.5))
+
+    assert is_close(np.log(0.5), f_log())
+
+
+def test_math_functions_float32():
+    """Test that math functions can be used inside kernels and match host numpy (float32)."""
+
+    @cudaq.kernel
+    def f_sin() -> np.float32:
+        return np.sin(np.float32(0.5))
+
+    assert is_close(np.sin(np.float32(0.5)), f_sin())
+
+    @cudaq.kernel
+    def f_cos() -> np.float32:
+        return np.cos(np.float32(0.5))
+
+    assert is_close(np.cos(np.float32(0.5)), f_cos())
+
+    @cudaq.kernel
+    def f_tan() -> np.float32:
+        return np.tan(np.float32(0.5))
+
+    assert is_close(np.tan(np.float32(0.5)), f_tan())
+
+    @cudaq.kernel
+    def f_arcsin() -> np.float32:
+        return np.arcsin(np.float32(0.5))
+
+    assert is_close(np.arcsin(np.float32(0.5)), f_arcsin())
+
+    @cudaq.kernel
+    def f_arccos() -> np.float32:
+        return np.arccos(np.float32(0.5))
+
+    assert is_close(np.arccos(np.float32(0.5)), f_arccos())
+
+    @cudaq.kernel
+    def f_arctan() -> np.float32:
+        return np.arctan(np.float32(0.5))
+
+    assert is_close(np.arctan(np.float32(0.5)), f_arctan())
+
+    @cudaq.kernel
+    def f_sqrt() -> np.float32:
+        return np.sqrt(np.float32(0.5))
+
+    assert is_close(np.sqrt(np.float32(0.5)), f_sqrt())
+
+    @cudaq.kernel
+    def f_exp() -> np.float32:
+        return np.exp(np.float32(0.5))
+
+    assert is_close(np.exp(np.float32(0.5)), f_exp())
+
+    @cudaq.kernel
+    def f_log() -> np.float32:
+        return np.log(np.float32(0.5))
+
+    assert is_close(np.log(np.float32(0.5)), f_log())
+
+
+def test_math_function_as_gate_parameter():
+    """Test that a math function result can be used as a gate parameter."""
+
+    @cudaq.kernel
+    def prep(p: float):
+        q = cudaq.qubit()
+        ry(2.0 * np.asin(np.sqrt(p)), q)
+
+    cudaq.set_random_seed(13)
+    p = 0.3
+    counts = cudaq.sample(prep, p, shots_count=20000)
+    assert abs(counts.probability('1') - p) < 0.02