from __future__ import annotations

__copyright__ = "Copyright (C) 2022 Alexandru Fikl"

__license__ = """

Permission is hereby granted, free of charge, to any person obtaining a copy

of this software and associated documentation files (the "Software"), to deal

in the Software without restriction, including without limitation the rights

to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

copies of the Software, and to permit persons to whom the Software is

furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in

all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN

THE SOFTWARE.

"""

from typing import TYPE_CHECKING, Any

from boxtree.array_context import PyOpenCLArrayContext as PyOpenCLArrayContextBase

from typing_extensions import override

import loopy as lp

from arraycontext.pytest import (

_PytestPyOpenCLArrayContextFactoryWithClass,

register_pytest_array_context_factory,

)

if TYPE_CHECKING:

from collections.abc import Iterator, Sequence

import namedisl as nisl

from numpy.typing import DTypeLike

from arraycontext import ArrayContext

from loopy import TranslationUnit

from loopy.codegen import PreambleInfo

from loopy.kernel.instruction import InstructionBase

from pytools.tag import ToTagSetConvertible

__doc__ = """

Array Context

-------------

.. autofunction:: make_loopy_program

.. autoclass:: PyOpenCLArrayContext

"""

# {{{ PyOpenCLArrayContext

def make_loopy_program(

domains: str | Sequence[str | nisl.Set],

statements: Sequence[InstructionBase | str] | str,

kernel_data: list[Any] | None = None, *,

name: str = "sumpy_loopy_kernel",

silenced_warnings: list[str] | str | None = None,

assumptions: str = "",

fixed_parameters: dict[str, Any] | None = None,

index_dtype: DTypeLike | None = None,

tags: ToTagSetConvertible = None):

"""Return a :class:`loopy.LoopKernel` suitable for use with

:meth:`arraycontext.ArrayContext.call_loopy`.

"""

if kernel_data is None:

kernel_data = [...]

if silenced_warnings is None:

silenced_warnings = []

import loopy as lp

from arraycontext.loopy import _DEFAULT_LOOPY_OPTIONS

return lp.make_kernel(

domains,

statements,

kernel_data=kernel_data,

options=_DEFAULT_LOOPY_OPTIONS,

default_offset=lp.auto,

name=name,

lang_version=lp.MOST_RECENT_LANGUAGE_VERSION,

assumptions=assumptions,

fixed_parameters=fixed_parameters,

silenced_warnings=silenced_warnings,

index_dtype=index_dtype,

tags=tags)

def _fp_contract_fast_preamble(

preamble_info: PreambleInfo

) -> Iterator[tuple[str, str]]:

yield ("fp_contract_fast_pocl", "#pragma clang fp contract(fast)")

class PyOpenCLArrayContext(PyOpenCLArrayContextBase):

@override

def transform_loopy_program(self, t_unit: TranslationUnit):

import pyopencl as cl

device = self.queue.device

if (device.platform.name == "Portable Computing Language"

and (device.type & cl.device_type.GPU)):

t_unit = lp.register_preamble_generators(

t_unit,

[_fp_contract_fast_preamble])

return t_unit

def is_cl_cpu(actx: ArrayContext) -> bool:

if not isinstance(actx, PyOpenCLArrayContext):

return False

import pyopencl as cl

return all(dev.type & cl.device_type.CPU for dev in actx.context.devices)

# }}}

# {{{ pytest

def _acf() -> ArrayContext:

import pyopencl as cl

ctx = cl.create_some_context()

queue = cl.CommandQueue(ctx)

return PyOpenCLArrayContext(queue)

class PytestPyOpenCLArrayContextFactory(

_PytestPyOpenCLArrayContextFactoryWithClass):

@property

@override

def actx_class(self) -> type[ArrayContext]:

return PyOpenCLArrayContext

@override

def __call__(self) -> ArrayContext:

# NOTE: prevent any cache explosions during testing!

from sympy.core.cache import clear_cache

clear_cache()

return super().__call__()

register_pytest_array_context_factory(

"sumpy.pyopencl",

PytestPyOpenCLArrayContextFactory)

# }}}