diff --git a/cudaq/lib/Frontend/nvqpp/ConvertExpr.cpp b/cudaq/lib/Frontend/nvqpp/ConvertExpr.cpp index 0f910d218a7..88a028ec747 100644 --- a/cudaq/lib/Frontend/nvqpp/ConvertExpr.cpp +++ b/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")) diff --git a/cudaq/test/AST-Quake/math_functions.cpp b/cudaq/test/AST-Quake/math_functions.cpp new file mode 100644 index 00000000000..5fd33685ef2 --- /dev/null +++ b/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 +#include + +// 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 diff --git a/docs/sphinx/specification/cudaq/kernels.rst b/docs/sphinx/specification/cudaq/kernels.rst index 5d686e02d91..4b9f0095c26 100644 --- a/docs/sphinx/specification/cudaq/kernels.rst +++ b/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). diff --git a/python/cudaq/kernel/ast_bridge.py b/python/cudaq/kernel/ast_bridge.py index 5cdbf342701..6baf313c0f1 100644 --- a/python/cudaq/kernel/ast_bridge.py +++ b/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) diff --git a/python/tests/kernel/test_kernel_float.py b/python/tests/kernel/test_kernel_float.py index db0c2a166a0..5142c0647ea 100644 --- a/python/tests/kernel/test_kernel_float.py +++ b/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