Code Generation¶
loopy
Target¶
- pytato.generate_loopy(result: ~pytato.array.Array | ~pytato.array.DictOfNamedArrays | dict[str, ~pytato.array.Array], target: ~pytato.target.loopy.LoopyTarget | None = None, options: ~loopy.options.Options | None = None, *, function_name: str = '_pt_kernel', cl_device: ~pyopencl._cl.Device | None = None, array_tag_t_to_not_propagate: frozenset[type[~pytools.tag.Tag]] = frozenset({<class 'pytato.tags.ImplStored'>, <class 'pytato.tags.Named'>, <class 'pytato.tags.PrefixNamed'>}), axis_tag_t_to_not_propagate: frozenset[type[~pytools.tag.Tag]] = frozenset({})) BoundProgram [source]¶
Code generation entry point.
- Parameters:
result – Outputs of the computation.
target – Code generation target.
options – Code generation options for the kernel.
- Returns:
A
pytato.target.BoundProgram
wrapping the generatedloopy.TranslationUnit
.
If result is a
dict
or apytato.DictOfNamedArrays
and options is not supplied, then the Loopy optionreturn_dict
will be set to True. If it is supplied,return_dict
must already be set to True.Note
pytato
metadata \(\mapsto\)loopy
metadata semantics:Inames that index over an
Array
’s axis in the allocation instruction are tagged with the correspondingAxis
’s tags. The caller may choose to not propagate axis tags of type axis_tag_t_to_not_propagate.pytato.Array.tags
of inputs/outputs in outputs would be copied over to the tags of the correspondingloopy.ArrayArg
. The caller may choose to not propagate array tags of type array_tag_t_to_not_propagate.Arrays tagged with
pytato.tags.ImplStored
would have their tags copied over to the tags of correspondingloopy.TemporaryVariable
. The caller may choose to not propagate array tags of type array_tag_t_to_not_propagate.
Warning
Currently only
Call
nodes that are tagged withpytato.tags.InlineCallTag
can be lowered toloopy
IR.
- class pytato.target.loopy.LoopyTarget[source]¶
An
loopy
target.- get_loopy_target() TargetBase [source]¶
Return the corresponding
loopy
target.
- bind_program(program: TranslationUnit, bound_arguments: Mapping[str, Any]) BoundProgram [source]¶
Create a
pytato.target.BoundProgram
for this code generation target.- Parameters:
program – the
loopy
programbound_arguments – a mapping from argument names to outputs
- class pytato.target.loopy.LoopyPyOpenCLTarget(device: Device | None = None)[source]¶
A
pyopencl
code generation target.- device¶
The
pyopencl
device used to construct theloopy.PyOpenCLTarget
, or None.
- class pytato.target.loopy.BoundPyOpenCLProgram(program: ~loopy.translation_unit.TranslationUnit, bound_arguments: ~collections.abc.Mapping[str, ~typing.Any], target: ~pytato.target.Target, _processed_bound_args_cache: dict[~pyopencl._cl.Context, ~collections.abc.Mapping[str, ~typing.Any]] = <factory>)[source]¶
A wrapper around a
loopy
kernel for execution withpyopencl
.- __call__(queue: CommandQueue, allocator=None, wait_for=None, out_host: bool | None = None, **kwargs: Any) Any [source]¶
Convenience function for launching a
pyopencl
computation.
- copy(*, program: TranslationUnit | None = None, bound_arguments: Mapping[str, Any] | None = None, target: Target | None = None) BoundPyOpenCLProgram [source]¶
- with_transformed_program(f: Callable[[TranslationUnit], TranslationUnit]) BoundPyOpenCLProgram [source]¶
Returns a copy of self with an f-transformed loopy translation unit.
- bind_to_context(context: Context, allocator: Callable[[int], MemoryObject] | None = None) BoundPyOpenCLExecutable [source]¶
- class pytato.target.loopy.BoundPyOpenCLExecutable(program: ExecutorBase, bound_arguments: Mapping[str, Any], target: Target, cl_context: Context)[source]¶
A wrapper around a
loopy
kernel for execution withpyopencl
. In contrast toBoundPyOpenCLProgram
, this object is specific to a givenpyopencl.Context
, allowing it to store aloopy.ExecutorBase
instead of aloopy.TranslationUnit
, as well as retrieving pre-transferred bound arguments without a cache lookup, permitting more efficient invocation.Create these objects using
BoundPyOpenCLProgram.bind_to_context()
.- __call__(queue: CommandQueue, allocator=None, wait_for=None, **kwargs: Any) Any [source]¶
Convenience function for launching a
pyopencl
computation.
- with_transformed_translation_unit(f: Callable[[TranslationUnit], TranslationUnit]) BoundPyOpenCLExecutable [source]¶
Returns a copy of self with an f-transformed loopy translation unit.
- class pytato.target.loopy.ImplSubstitution[source]¶
An
ImplementationStrategy
that lowers the array expression as aloopy.SubstitutionRule
invocation.
Stuff that’s only here because the documentation tool wants it¶
- class pyopencl._cl.MemoryObject¶
- pytato.generate_jax(expr: Array | Mapping[str, Array] | DictOfNamedArrays, *, target: JAXPythonTarget | None = None, jit: bool = False, function_name: str = '_pt_kernel', show_code: bool = False, colorize_show_code: bool = True) BoundJAXPythonProgram [source]¶
Returns a
pytato.target.python.BoundJAXPythonProgram
for the array expressions in expr.- Parameters:
jit – If True, the generated function is decorated with
jax.jit()
.function – Name of the entrypoint function in the generated code.
show_code – If True, the generated code is printed to
stdout
.
- class pytato.target.python.PythonTarget[source]¶
A target that generates code for a python program, typically by invoking some
numpy
-like for the array operations.
- class pytato.target.python.BoundPythonProgram(program: Any, bound_arguments: Mapping[str, Any], target: Target, expected_arguments: frozenset[str], entrypoint: str)[source]¶
A wrapper for executing python programs with bound arguments.
- copy(**kwargs: Any) BoundPythonProgram [source]¶
- with_transformed_program(*args: Any, **kwargs: Any) BoundPythonProgram [source]¶
Generic target support¶
Code Generation Targets¶
- class pytato.target.BoundProgram(program: Any, bound_arguments: Mapping[str, Any], target: Target)[source]¶
A container for the result of code generation along with data bindings for already-bound arguments.
- target¶
The code generation target.
- program¶
Description of the program as per
BoundProgram.target
.
- bound_arguments¶
A map from names to pre-bound kernel arguments.