Infinity in CUDA kernels

pygpu/tests/test_elemwise.py

@@ -1,5 +1,6 @@
 import operator
 import numpy
+from mako.template import Template
 
 from unittest import TestCase
 from pygpu import gpuarray, ndgpuarray as elemary
@@ -297,3 +298,37 @@ def broadcast(shapea, shapeb):
     rg = ag + bg
 
     check_meta_content(rg, rc)
+
+
+_inf_preamb_tpl = Template('''
+WITHIN_KERNEL ${flt}
+infinity() {return INFINITY;}
+
+WITHIN_KERNEL ${flt}
+neg_infinity() {return -INFINITY;}
+''')
+
+
+def test_infinity():
+    for dtype in ['float32', 'float64']:
+        ac, ag = gen_gpuarray((2,), dtype, ctx=context, cls=elemary)
+        out_g = ag._empty_like_me()
+        flt = 'ga_float' if dtype == 'float32' else 'ga_double'
+        out_arg = arg('out', out_g.dtype, scalar=False, read=False, write=True)
+        preamble = _inf_preamb_tpl.render(flt=flt)
+
+        # +infinity
+        ac[:] = numpy.inf
+        expr_inf = 'out = infinity()'
+        kernel = GpuElemwise(context, expr_inf, [out_arg],
+                             preamble=preamble)
+        kernel(out_g)
+        assert numpy.array_equal(ac, numpy.asarray(out_g))
+
+        # -infinity
+        ac[:] = -numpy.inf
+        expr_neginf = 'out = neg_infinity()'
+        kernel = GpuElemwise(context, expr_neginf, [out_arg],
+                             preamble=preamble)
+        kernel(out_g)
+        assert numpy.array_equal(ac, numpy.asarray(out_g))

src/gpuarray_buffer_cuda.c

@@ -318,6 +318,10 @@ static const char CUDA_PREAMBLE[] =
     "#undef NAN\n"
     "#endif\n"
     "#define NAN __int_as_float(0x7fffffff)\n"
+    "#ifdef INFINITY\n"
+    "#undef INFINITY\n"
+    "#endif\n"
+    "#define INFINITY __int_as_float(0x7f800000)\n"
     "#define LID_0 threadIdx.x\n"
     "#define LID_1 threadIdx.y\n"
     "#define LID_2 threadIdx.z\n"