Code Generation

loopy Target

pytato.generate_loopy(result: Array | DictOfNamedArrays | dict[str, Array], target: LoopyTarget | None = None, options: lp.Options | None = None, *, function_name: str = '_pt_kernel', cl_device: pyopencl.Device | None = None, array_tag_t_to_not_propagate: frozenset[type[Tag]] = frozenset({<class 'pytato.tags.ImplStored'>, <class 'pytato.tags.Named'>, <class 'pytato.tags.PrefixNamed'>}), axis_tag_t_to_not_propagate: frozenset[type[Tag]] = frozenset({})) → BoundProgram

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

An loopy target.

get_loopy_target() → TargetBase

Return the corresponding loopy target.

bind_program(program: loopy.TranslationUnit, bound_arguments: Mapping[str, Any]) → BoundProgram

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: pyopencl.Device | None = None)

A pyopencl code generation target.

device

The pyopencl device used to construct the loopy.PyOpenCLTarget, or None.

class pytato.target.loopy.BoundPyOpenCLProgram(program: loopy.TranslationUnit, bound_arguments: Mapping[str, Any], target: Target, _processed_bound_args_cache: dict[pyopencl.Context, Mapping[str, Any]] = <factory>)

A wrapper around a loopy kernel for execution with pyopencl.

__call__(queue: pyopencl.CommandQueue, allocator=None, wait_for=None, out_host: bool | None = None, **kwargs: Any) → Any

Convenience function for launching a pyopencl computation.

copy(*, program: loopy.TranslationUnit | None = None, bound_arguments: Mapping[str, Any] | None = None, target: Target | None = None) → BoundPyOpenCLProgram

with_transformed_program(f: Callable[[loopy.TranslationUnit], loopy.TranslationUnit]) → BoundPyOpenCLProgram

Returns a copy of self with an f-transformed loopy translation unit.

bind_to_context(context: pyopencl.Context, allocator: Callable[[int], pyopencl.MemoryObject] | None = None) → BoundPyOpenCLExecutable

class pytato.target.loopy.BoundPyOpenCLExecutable(program: loopy.ExecutorBase, bound_arguments: Mapping[str, Any], target: Target, cl_context: pyopencl.Context)

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: pyopencl.CommandQueue, allocator=None, wait_for=None, **kwargs: Any) → Any

Convenience function for launching a pyopencl computation.

with_transformed_translation_unit(f: Callable[[loopy.TranslationUnit], loopy.TranslationUnit]) → BoundPyOpenCLExecutable

Returns a copy of self with an f-transformed loopy translation unit.

class pytato.target.loopy.ImplSubstitution

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

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

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)

A wrapper for executing python programs with bound arguments.

__call__(*args: Any, **kwargs: Any) → Any

Call self as a function.

copy(**kwargs: Any) → BoundPythonProgram

with_transformed_program(*args: Any, **kwargs: Any) → BoundPythonProgram

class pytato.target.python.JAXPythonTarget

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)

Generic target support

Code Generation Targets

class pytato.target.Target

An abstract code generation target.

class pytato.target.BoundProgram(program: Any, bound_arguments: Mapping[str, Any], target: Target)

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__()

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.