Added 'out' parameter to unary math functions.

inducer · Mar 5, 2014 · a1cc472 · a1cc472
1 parent dcbcc84
commit a1cc472
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 20 deletions.
diff --git a/doc/source/array.rst b/doc/source/array.rst
@@ -338,34 +338,34 @@ workalikes for the functions contained in :mod:`math`.
 Rounding and Absolute Value
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-.. function:: fabs(array, stream=None)
-.. function:: ceil(array, stream=None)
-.. function:: floor(array, stream=None)
+.. function:: fabs(array, out=None, stream=None)
+.. function:: ceil(array, out=None, stream=None)
+.. function:: floor(array, out=None, stream=None)
 
 Exponentials, Logarithms and Roots
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-.. function:: exp(array, stream=None)
-.. function:: log(array, stream=None)
-.. function:: log10(array, stream=None)
-.. function:: sqrt(array, stream=None)
+.. function:: exp(array, out=None, stream=None)
+.. function:: log(array, out=None, stream=None)
+.. function:: log10(array, out=None, stream=None)
+.. function:: sqrt(array, out=None, stream=None)
 
 Trigonometric Functions
 ^^^^^^^^^^^^^^^^^^^^^^^
 
-.. function:: sin(array, stream=None)
-.. function:: cos(array, stream=None)
-.. function:: tan(array, stream=None)
-.. function:: asin(array, stream=None)
-.. function:: acos(array, stream=None)
-.. function:: atan(array, stream=None)
+.. function:: sin(array, out=None, stream=None)
+.. function:: cos(array, out=None, stream=None)
+.. function:: tan(array, out=None, stream=None)
+.. function:: asin(array, out=None, stream=None)
+.. function:: acos(array, out=None, stream=None)
+.. function:: atan(array, out=None, stream=None)
 
 Hyperbolic Functions
 ^^^^^^^^^^^^^^^^^^^^
 
-.. function:: sinh(array, stream=None)
-.. function:: cosh(array, stream=None)
-.. function:: tanh(array, stream=None)
+.. function:: sinh(array, out=None, stream=None)
+.. function:: cosh(array, out=None, stream=None)
+.. function:: tanh(array, out=None, stream=None)
 
 Floating Point Decomposition and Assembly
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

diff --git a/pycuda/cumath.py b/pycuda/cumath.py
@@ -3,8 +3,7 @@
 import numpy as np
 
 def _make_unary_array_func(name):
-    def f(array, stream=None):
-        result = array._new_like_me()
+    def f(array, out=None, stream=None):
 
         if array.dtype == np.float32:
             func_name = name + "f"
@@ -15,11 +14,18 @@ def f(array, stream=None):
             raise RuntimeError("only contiguous arrays may "
                     "be used as arguments to this operation")
 
+        if out is None:
+            out = array._new_like_me()
+        else:
+            assert out.dtype == array.dtype
+            assert out.strides == array.strides
+            assert out.shape == array.shape
+
         func = elementwise.get_unary_func_kernel(func_name, array.dtype)
         func.prepared_async_call(array._grid, array._block, stream,
-                array.gpudata, result.gpudata, array.mem_size)
+                array.gpudata, out.gpudata, array.mem_size)
 
-        return result
+        return out
     return f
 
 fabs = _make_unary_array_func("fabs")

diff --git a/test/test_cumath.py b/test/test_cumath.py
@@ -54,6 +54,14 @@ def test():
                 assert (max_err <= threshold).all(), \
                         (max_err, name, dtype)
 
+                gpu_results2 = gpuarray.empty_like(args)
+                gr2 = gpu_func(args, out=gpu_results2)
+                assert gpu_results2 is gr2
+                gr2 = gr2.get()
+                max_err = np.max(np.abs(cpu_results - gr2))
+                assert (max_err <= threshold).all(), \
+                        (max_err, name, dtype)
+
     return mark_cuda_test(test)