Code Generation#

loopy Target#

pytato.generate_loopy(result: ~typing.Union[~pytato.array.Array, ~pytato.array.DictOfNamedArrays, ~typing.Dict[str, ~pytato.array.Array]], target: ~typing.Optional[~pytato.target.loopy.LoopyTarget] = None, options: ~typing.Optional[~loopy.options.Options] = None, *, function_name: str = '_pt_kernel', cl_device: ~typing.Optional[~pyopencl._cl.Device] = None, array_tag_t_to_not_propagate: ~typing.FrozenSet[~typing.Type[~pytools.tag.Tag]] = frozenset({<class 'pytato.tags.ImplStored'>, <class 'pytato.tags.Named'>, <class 'pytato.tags.PrefixNamed'>}), axis_tag_t_to_not_propagate: ~typing.FrozenSet[~typing.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.

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: Optional[Device] = 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: Any, bound_arguments: Mapping[str, Any], target: Target)[source]#

A wrapper around a loopy kernel for execution with pyopencl.

__call__(queue: CommandQueue, allocator=None, wait_for=None, out_host=None, **kwargs: Any) Any[source]#

Convenience function for launching a pyopencl computation.

copy(*, program: Optional[TranslationUnit] = None, bound_arguments: Optional[Mapping[str, Any]] = None, target: Optional[Target] = None) BoundPyOpenCLProgram[source]#
with_transformed_program(f: Callable[[TranslationUnit], TranslationUnit]) BoundPyOpenCLProgram[source]#

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

pytato.generate_jax(expr: Union[Array, Mapping[str, Array], DictOfNamedArrays], *, target: Optional[JAXPythonTarget] = 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.