Add kernel math function support

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,56 @@
+/*******************************************************************************
+ * 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 | FileCheck %s
+
+#include <cmath>
+#include <cudaq.h>
+
+struct math_functions {
+  void operator()(double theta) __qpu__ {
+    cudaq::qubit q;
+    double angle = std::sin(theta) + std::cos(theta) + std::tan(theta) +
+                   std::asin(theta) + std::acos(theta) + std::atan(theta) +
+                   std::sqrt(theta) + std::exp(theta) + std::log(theta);
+    rx(angle, q);
+  }
+};
+
+struct math_functions_float {
+  void operator()(float theta) __qpu__ {
+    cudaq::qubit q;
+    float angle = sinf(theta) + cosf(theta) + tanf(theta) + asinf(theta) +
+                  acosf(theta) + atanf(theta) + sqrtf(theta) + expf(theta) +
+                  logf(theta);
+    rx(angle, q);
+  }
+};
+
+// CHECK-LABEL:   func.func @__nvqpp__mlirgen__math_functions
+// 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__math_functions_float
+// 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

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', 'sqrt', 'ceil',
+            'exp', 'log'
+        ]
 
     def __isSupportedVectorFunction(self, id):
         return id in ['front', 'back', 'append']
@@ -3504,6 +3507,38 @@ def bodyBuilder(iterVar):
                             return
                         self.pushValue(math.SqrtOp(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 == 'asin':
+                        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 == 'acos':
+                        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 == 'atan':
+                        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 == 'exp':
                         if ComplexType.isinstance(value.type):
                             # Note: using `complex.ExpOp` results in a "can't
@@ -3527,6 +3562,14 @@ def bodyBuilder(iterVar):
                             return
                         self.pushValue(math.ExpOp(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
                     if node.func.attr == 'ceil':
                         if ComplexType.isinstance(value.type):
                             self.emitFatalError(

python/tests/kernel/test_kernel_float.py

@@ -135,6 +135,38 @@ def float_np_use() -> np.float64:
     t = np.cos(np.float64(np.pi / 2 + 1))
     assert is_close(t, float_np_use())
 
+    # Use a float inside np in a kernel (tan)
+    @cudaq.kernel
+    def float_np_use() -> np.float64:
+        return np.tan(np.float64(0.25))
+
+    t = np.tan(np.float64(0.25))
+    assert is_close(t, float_np_use())
+
+    # Use a float inside np in a kernel (asin)
+    @cudaq.kernel
+    def float_np_use() -> np.float64:
+        return np.asin(np.float64(0.25))
+
+    t = np.asin(np.float64(0.25))
+    assert is_close(t, float_np_use())
+
+    # Use a float inside np in a kernel (acos)
+    @cudaq.kernel
+    def float_np_use() -> np.float64:
+        return np.acos(np.float64(0.25))
+
+    t = np.acos(np.float64(0.25))
+    assert is_close(t, float_np_use())
+
+    # Use a float inside np in a kernel (atan)
+    @cudaq.kernel
+    def float_np_use() -> np.float64:
+        return np.atan(np.float64(0.25))
+
+    t = np.atan(np.float64(0.25))
+    assert is_close(t, float_np_use())
+
     # Use a float inside np in a kernel (sqrt)
     @cudaq.kernel
     def float_np_use() -> np.float64:
@@ -159,6 +191,14 @@ def float_np_use() -> np.float64:
     t = np.exp(np.float64(np.pi / 2 + 1))
     assert is_close(t, float_np_use())
 
+    # Use a float inside np in a kernel (log)
+    @cudaq.kernel
+    def float_np_use() -> np.float64:
+        return np.log(np.float64(np.pi / 2 + 1))
+
+    t = np.log(np.float64(np.pi / 2 + 1))
+    assert is_close(t, float_np_use())
+
 
 # np.float32
 
@@ -224,6 +264,38 @@ def float_np_use() -> np.float32:
     t = np.cos(np.float32(np.pi / 2 + 1))
     assert is_close(t, float_np_use())
 
+    # Use a float inside np in a kernel (tan)
+    @cudaq.kernel
+    def float_np_use() -> np.float32:
+        return np.tan(np.float32(0.25))
+
+    t = np.tan(np.float32(0.25))
+    assert is_close(t, float_np_use())
+
+    # Use a float inside np in a kernel (asin)
+    @cudaq.kernel
+    def float_np_use() -> np.float32:
+        return np.asin(np.float32(0.25))
+
+    t = np.asin(np.float32(0.25))
+    assert is_close(t, float_np_use())
+
+    # Use a float inside np in a kernel (acos)
+    @cudaq.kernel
+    def float_np_use() -> np.float32:
+        return np.acos(np.float32(0.25))
+
+    t = np.acos(np.float32(0.25))
+    assert is_close(t, float_np_use())
+
+    # Use a float inside np in a kernel (atan)
+    @cudaq.kernel
+    def float_np_use() -> np.float32:
+        return np.atan(np.float32(0.25))
+
+    t = np.atan(np.float32(0.25))
+    assert is_close(t, float_np_use())
+
     # Use a float inside np in a kernel (sqrt)
     @cudaq.kernel
     def float_np_use() -> np.float32:
@@ -248,6 +320,14 @@ def float_np_use() -> np.float32:
     t = np.exp(np.float32(np.pi / 2 + 1))
     assert is_close(t, float_np_use())
 
+    # Use a float inside np in a kernel (log)
+    @cudaq.kernel
+    def float_np_use() -> np.float32:
+        return np.log(np.float32(np.pi / 2 + 1))
+
+    t = np.log(np.float32(np.pi / 2 + 1))
+    assert is_close(t, float_np_use())
+
 
 def test_float_list_parameter_promotion():