|
23 | 23 | THE SOFTWARE. |
24 | 24 | """ |
25 | 25 |
|
| 26 | +from typing import TYPE_CHECKING, Any |
| 27 | + |
26 | 28 | from boxtree.array_context import PyOpenCLArrayContext as PyOpenCLArrayContextBase |
| 29 | +from typing_extensions import override |
27 | 30 |
|
| 31 | +import loopy as lp |
28 | 32 | from arraycontext.pytest import ( |
29 | 33 | _PytestPyOpenCLArrayContextFactoryWithClass, |
30 | 34 | register_pytest_array_context_factory, |
31 | 35 | ) |
32 | 36 |
|
33 | 37 |
|
| 38 | +if TYPE_CHECKING: |
| 39 | + from collections.abc import Iterator |
| 40 | + |
| 41 | + from numpy.typing import DTypeLike |
| 42 | + |
| 43 | + from arraycontext import ArrayContext |
| 44 | + from loopy import TranslationUnit |
| 45 | + from loopy.codegen import PreambleInfo |
| 46 | + from pytools.tag import ToTagSetConvertible |
| 47 | + |
| 48 | + |
34 | 49 | __doc__ = """ |
35 | 50 | Array Context |
36 | 51 | ------------- |
37 | 52 |
|
| 53 | +.. autofunction:: make_loopy_program |
38 | 54 | .. autoclass:: PyOpenCLArrayContext |
39 | 55 | """ |
40 | 56 |
|
41 | 57 |
|
42 | 58 | # {{{ PyOpenCLArrayContext |
43 | 59 |
|
| 60 | +def make_loopy_program( |
| 61 | + domains, statements, |
| 62 | + kernel_data: list[Any] | None = None, *, |
| 63 | + name: str = "sumpy_loopy_kernel", |
| 64 | + silenced_warnings: list[str] | str | None = None, |
| 65 | + assumptions: str = "", |
| 66 | + fixed_parameters: dict[str, Any] | None = None, |
| 67 | + index_dtype: DTypeLike | None = None, |
| 68 | + tags: ToTagSetConvertible = None): |
| 69 | + """Return a :class:`loopy.LoopKernel` suitable for use with |
| 70 | + :meth:`arraycontext.ArrayContext.call_loopy`. |
| 71 | + """ |
| 72 | + if kernel_data is None: |
| 73 | + kernel_data = [...] |
| 74 | + |
| 75 | + if silenced_warnings is None: |
| 76 | + silenced_warnings = [] |
| 77 | + |
| 78 | + import loopy as lp |
| 79 | + from arraycontext.loopy import _DEFAULT_LOOPY_OPTIONS |
| 80 | + |
| 81 | + return lp.make_kernel( |
| 82 | + domains, |
| 83 | + statements, |
| 84 | + kernel_data=kernel_data, |
| 85 | + options=_DEFAULT_LOOPY_OPTIONS, |
| 86 | + default_offset=lp.auto, |
| 87 | + name=name, |
| 88 | + lang_version=lp.MOST_RECENT_LANGUAGE_VERSION, |
| 89 | + assumptions=assumptions, |
| 90 | + fixed_parameters=fixed_parameters, |
| 91 | + silenced_warnings=silenced_warnings, |
| 92 | + index_dtype=index_dtype, |
| 93 | + tags=tags) |
| 94 | + |
| 95 | + |
| 96 | +def _fp_contract_fast_preamble( |
| 97 | + preamble_info: PreambleInfo |
| 98 | + ) -> Iterator[tuple[str, str]]: |
| 99 | + yield ("fp_contract_fast_pocl", "#pragma clang fp contract(fast)") |
| 100 | + |
| 101 | + |
44 | 102 | class PyOpenCLArrayContext(PyOpenCLArrayContextBase): |
45 | | - def transform_loopy_program(self, t_unit): |
46 | | - default_ep = t_unit.default_entrypoint |
47 | | - options = default_ep.options |
| 103 | + @override |
| 104 | + def transform_loopy_program(self, t_unit: TranslationUnit): |
| 105 | + import pyopencl as cl |
| 106 | + device = self.queue.device |
| 107 | + if (device.platform.name == "Portable Computing Language" |
| 108 | + and (device.type & cl.device_type.GPU)): |
| 109 | + t_unit = lp.register_preamble_generators( |
| 110 | + t_unit, |
| 111 | + [_fp_contract_fast_preamble]) |
| 112 | + |
| 113 | + return t_unit |
48 | 114 |
|
49 | | - if not (options.return_dict and options.no_numpy): |
50 | | - raise ValueError("Loopy kernel passed to call_loopy must " |
51 | | - "have return_dict and no_numpy options set. " |
52 | | - "Did you use arraycontext.make_loopy_program " |
53 | | - "to create this kernel?") |
54 | 115 |
|
55 | | - return super().transform_loopy_program(t_unit) |
| 116 | +def is_cl_cpu(actx: ArrayContext) -> bool: |
| 117 | + if not isinstance(actx, PyOpenCLArrayContext): |
| 118 | + return False |
| 119 | + |
| 120 | + import pyopencl as cl |
| 121 | + return all(dev.type & cl.device_type.CPU for dev in actx.context.devices) |
56 | 122 |
|
57 | 123 | # }}} |
58 | 124 |
|
|
0 commit comments