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 generated loopy.TranslationUnit.

If result is a dict or a pytato.DictOfNamedArrays and options is not supplied, then the Loopy option return_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 corresponding Axis’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 corresponding loopy.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 corresponding loopy.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 with pytato.tags.InlineCallTag can be lowered to loopy 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 program

  • bound_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 the loopy.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 with pyopencl.

__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 with pyopencl. In contrast to BoundPyOpenCLProgram, this object is specific to a given pyopencl.Context, allowing it to store a loopy.ExecutorBase instead of a loopy.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 a loopy.SubstitutionRule invocation.

Stuff that’s only here because the documentation tool wants it¶

class pyopencl._cl.MemoryObject¶

See pyopencl.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.

__call__(*args: Any, **kwargs: Any) Any[source]¶

Call self as a function.

copy(**kwargs: Any) BoundPythonProgram[source]¶
with_transformed_program(*args: Any, **kwargs: Any) BoundPythonProgram[source]¶
class pytato.target.python.JAXPythonTarget[source]¶

A target that generates code for a python program by offloading array operations to jax.numpy.

class pytato.target.python.BoundJAXPythonProgram(program: Any, bound_arguments: Mapping[str, Any], target: Target, expected_arguments: frozenset[str], entrypoint: str)[source]¶

Generic target support¶

Code Generation Targets¶

class pytato.target.Target[source]¶

An abstract code generation target.

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.

__call__()[source]¶

It is expected that every concrete subclass of this class have a __call__ method to run the generated code, however interfaces may vary based on the specific subclass.