Merge branch 'main' into better_loop_around_nest_map

Author: inducer

loopy/target/pyopencl.py

@@ -51,7 +51,7 @@ def with_types(self, arg_id_to_dtype, callables_table):
         for id in arg_id_to_dtype:
             # since all the below functions are single arg.
             if not -1 <= id <= 0:
-                raise LoopyError("%s can only take one argument." % name)
+                raise LoopyError(f"{name} can only take one argument")
 
         if 0 not in arg_id_to_dtype or arg_id_to_dtype[0] is None:
             # the types provided aren't mature enough to specialize the
@@ -69,14 +69,23 @@ def with_types(self, arg_id_to_dtype, callables_table):
                 elif dtype.numpy_dtype == np.complex128:
                     tpname = "cdouble"
                 else:
-                    raise LoopyTypeError("unexpected complex type '%s'" % dtype)
+                    raise LoopyTypeError(f"unexpected complex type '{dtype}'")
 
                 return (
                         self.copy(name_in_target=f"{tpname}_{name}",
                             arg_id_to_dtype={0: dtype, -1: NumpyType(
                                 np.dtype(dtype.numpy_dtype.type(0).real))}),
                         callables_table)
 
+        if name in ["real", "imag"]:
+            if not dtype.is_complex():
+                tpname = dtype.numpy_dtype.type.__name__
+                return (
+                        self.copy(
+                            name_in_target=f"lpy_{name}_{tpname}",
+                            arg_id_to_dtype={0: dtype, -1: dtype}),
+                        callables_table)
+
         if name in ["sqrt", "exp", "log",
                 "sin", "cos", "tan",
                 "sinh", "cosh", "tanh",
@@ -110,6 +119,23 @@ def with_types(self, arg_id_to_dtype, callables_table):
                 self.copy(arg_id_to_dtype=arg_id_to_dtype),
                 callables_table)
 
+    def generate_preambles(self, target):
+        name = self.name_in_target
+        if name.startswith("lpy_real") or name.startswith("lpy_imag"):
+            if name.startswith("lpy_real"):
+                ret = "x"
+            else:
+                ret = "0"
+
+            dtype = self.arg_id_to_dtype[-1]
+            ctype = target.dtype_to_typename(dtype)
+
+            yield(f"40_{name}", f"""
+            static inline {ctype} {name}({ctype} x) {{
+                return {ret};
+            }}
+            """)
+
 
 def get_pyopencl_callables():
     pyopencl_ids = ["sqrt", "exp", "log", "sin", "cos", "tan", "sinh", "cosh",

loopy/target/pyopencl_execution.py

@@ -101,7 +101,7 @@ def handle_alloc(self, gen, arg, kernel_arg, strify, skip_arg_checks):
 
         if not skip_arg_checks:
             for i in range(num_axes):
-                gen("assert _lpy_ustrides_%d > 0, "
+                gen("assert _lpy_ustrides_%d >= 0, "
                         "\"'%s' has negative stride in axis %d\""
                         % (i, arg.name, i))
 

test/test_expression.py

@@ -548,6 +548,25 @@ def test_complex_support(ctx_factory, target):
                                (0.5*n*(n-1) - 0.5*n*(n-1)*1j) ** 2)
 
 
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_real_with_real_argument(ctx_factory, dtype):
+    ctx = ctx_factory()
+    queue = cl.CommandQueue(ctx)
+
+    knl = lp.make_kernel(
+            "{[i]: 0 <= i < nresult}",
+            "result[i] = real(ary[i])",
+            )
+
+    rng = np.random.default_rng()
+    ary = cl.array.to_device(queue, rng.random(128).astype(dtype))
+
+    _, (result,) = knl(queue, ary=ary)
+
+    assert result.dtype == ary.dtype
+    np.testing.assert_allclose(result.get(), np.real(ary.get()))
+
+
 def test_bool_type_context(ctx_factory):
     # Checks if a boolean type context is correctly handled in codegen phase.
     # See https://github.com/inducer/loopy/pull/258

test/test_loopy.py

@@ -3216,6 +3216,21 @@ def test_get_return_from_kernel_mapping():
     assert ret_from_knl_idx[9] == 10
 
 
+def test_zero_stride_array(ctx_factory):
+    ctx = ctx_factory()
+    cq = cl.CommandQueue(ctx)
+
+    knl = lp.make_kernel(
+        ["{[i]: 0<=i<10}",
+         "{[j]: 1=0}"],
+        """
+        y[i, j] = 1
+        """, [lp.GlobalArg("y", shape=(10, 0))])
+
+    evt, (out,) = knl(cq)
+    assert out.shape == (10, 0)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])