Add ispc stream harness

examples/python/ispc-harness.py

+import loopy as lp
+import numpy as np
+import numpy.linalg as la
+import ctypes
+import os
+from time import time
+
+
+# {{{ temporary directory
+
+class TemporaryDirectory(object):
+    """Create and return a temporary directory.  This has the same
+    behavior as mkdtemp but can be used as a context manager.  For
+    example:
+
+        with TemporaryDirectory() as tmpdir:
+            ...
+
+    Upon exiting the context, the directory and everything contained
+    in it are removed.
+    """
+
+    # Yanked from
+    # https://hg.python.org/cpython/file/3.3/Lib/tempfile.py
+
+    # Handle mkdtemp raising an exception
+    name = None
+    _closed = False
+
+    def __init__(self, suffix="", prefix="tmp", dir=None):
+        from tempfile import mkdtemp
+        self.name = mkdtemp(suffix, prefix, dir)
+
+    def __repr__(self):
+        return "<{} {!r}>".format(self.__class__.__name__, self.name)
+
+    def __enter__(self):
+        return self.name
+
+    def cleanup(self, _warn=False):
+        import warnings
+        if self.name and not self._closed:
+            from shutil import rmtree
+            try:
+                rmtree(self.name)
+            except (TypeError, AttributeError) as ex:
+                if "None" not in '%s' % (ex,):
+                    raise
+                self._rmtree(self.name)
+            self._closed = True
+            if _warn and warnings.warn:
+                warnings.warn("Implicitly cleaning up {!r}".format(self))
+
+    def __exit__(self, exc, value, tb):
+        self.cleanup()
+
+    def __del__(self):
+        # Issue a ResourceWarning if implicit cleanup needed
+        self.cleanup(_warn=True)
+
+# }}}
+
+
+# {{{ build_ispc_shared_lib
+
+def build_ispc_shared_lib(
+        cwd, ispc_sources, cxx_sources,
+        ispc_options=[], cxx_options=[]):
+    from os.path import join
+
+    ispc_source_names = []
+    for name, contents in ispc_sources:
+        ispc_source_names.append(name)
+
+        with open(join(cwd, name), "w") as srcf:
+            srcf.write(contents)
+
+    cxx_source_names = []
+    for name, contents in cxx_sources:
+        cxx_source_names.append(name)
+
+        with open(join(cwd, name), "w") as srcf:
+            srcf.write(contents)
+
+    from subprocess import check_call
+
+    check_call(
+            ["ispc",
+                "--pic",
+                "--opt=force-aligned-memory",
+                "--target=avx2-i32x8",
+                "-o", "ispc.o"]
+            + ispc_options
+            + list(ispc_source_names),
+            cwd=cwd)
+
+    check_call(
+            [
+                "g++",
+                "-shared", "-fopenmp", "-Wl,--export-dynamic",
+                "-fPIC",
+                "-oshared.so",
+                "ispc.o",
+                ]
+            + cxx_options
+            + list(cxx_source_names),
+            cwd=cwd)
+
+# }}}
+
+
+def cptr_from_numpy(obj):
+    ary_intf = getattr(obj, "__array_interface__", None)
+    if ary_intf is None:
+        raise RuntimeError("no array interface")
+
+    buf_base, is_read_only = ary_intf["data"]
+    return ctypes.c_void_p(buf_base + ary_intf.get("offset", 0))
+
+
+def main():
+    with open("tasksys.cpp", "r") as ts_file:
+        tasksys_source = ts_file.read()
+
+    from loopy.target.ispc import ISPCTarget
+    stream_knl = lp.make_kernel(
+            "{[i]: 0<=i<n}",
+            "z[i] = a*x[i] + y[i]",
+            target=ISPCTarget())
+
+    stream_dtype = np.float64
+    stream_ctype = ctypes.c_double
+
+    stream_knl = lp.add_and_infer_dtypes(stream_knl, {
+        "a": stream_dtype,
+        "x": stream_dtype,
+        "y": stream_dtype
+        })
+
+    stream_knl = lp.assume(stream_knl, "n>0")
+    stream_knl = lp.split_iname(stream_knl, "i", 8, inner_tag="l.0")
+    stream_knl = lp.split_iname(stream_knl,
+            "i_outer", 2**22, outer_tag="g.0")
+    stream_knl = lp.preprocess_kernel(stream_knl)
+    stream_knl = lp.get_one_scheduled_kernel(stream_knl)
+    stream_knl = lp.set_argument_order(stream_knl, "n,a,x,y,z")
+    ispc_code, arg_info = lp.generate_code(stream_knl)
+
+    with TemporaryDirectory() as tmpdir:
+        build_ispc_shared_lib(
+                tmpdir,
+                [("stream.ispc", ispc_code)],
+                [("tasksys.cpp", tasksys_source)])
+
+        print(ispc_code)
+        knl_lib = ctypes.cdll.LoadLibrary(os.path.join(tmpdir, "shared.so"))
+
+        n = 2**28
+        a = 5
+        x = np.empty(n, dtype=stream_dtype)
+        y = np.empty(n, dtype=stream_dtype)
+        z = np.empty(n, dtype=stream_dtype)
+
+        nruns = 30
+        start_time = time()
+        for irun in range(nruns):
+            knl_lib.loopy_kernel(
+                    ctypes.c_int(n), stream_ctype(a),
+                    cptr_from_numpy(x),
+                    cptr_from_numpy(y),
+                    cptr_from_numpy(z))
+        elapsed = time() - start_time
+
+        print(1e-9*3*x.nbytes*nruns/elapsed, "GB/s")
+
+        print(la.norm(z-a*x+y))
+
+
+
+if __name__ == "__main__":
+    main()
+
+# vim: foldmethod=marker