Merge branch 'main' into get-max-node-depth

Author: majosm

doc/distributed.rst

@@ -1,5 +1,5 @@
-Support for Distributed-Memory/Message Passing
-==============================================
+Distributed-Memory/Message Passing
+==================================
 
 .. automodule:: pytato.distributed
 

examples/demo.py

@@ -11,17 +11,31 @@
 result = pt.DictOfNamedArrays({"a2a": a2a, "aat": aat})
 
 
-# {{{ execute
+# {{{ execute with loopy/pyopencl
 
 import pyopencl as cl
+import pyopencl.array as cla
 ctx = cl.create_some_context()
 queue = cl.CommandQueue(ctx)
-prg = pt.generate_loopy(result, cl_device=queue.device)
+
+# simple interface
+prg = pt.generate_loopy(result)
 a = np.random.randn(20, 20)
 _, out = prg(queue, a=a)
 assert np.allclose(out["a2a"], a@(2*a))
 assert np.allclose(out["aat"], a@a.T)
 
+# interface for efficient execution
+import loopy as lp
+prg = pt.generate_loopy(
+        result, options=lp.Options(no_numpy=True, return_dict=True)
+        ).bind_to_context(ctx)
+a = np.random.randn(20, 20)
+a_dev = cla.to_device(queue, a)
+_, out = prg(queue, a=a_dev)
+assert np.allclose(out["a2a"].get(), a@(2*a))
+assert np.allclose(out["aat"].get(), a@a.T)
+
 # }}}
 
 # {{{ print generated Loopy and OpenCL code