Internals¶

Generic Code Generation Functionality¶

class pytato.codegen.CodeGenPreprocessor(target: Target, kernels_seen: dict[str, LoopKernel] | None = None)[source]¶

A mapper that preprocesses graphs to simplify code generation.

The following node simplifications are performed:

class pytato.codegen.PreprocessResult(outputs: 'DictOfNamedArrays', compute_order: 'tuple[str, ...]', bound_arguments: 'dict[str, DataInterface]')[source]¶
pytato.codegen.preprocess(outputs: DictOfNamedArrays, target: Target) PreprocessResult[source]¶

Preprocess a computation for code generation.

pytato.codegen.normalize_outputs(result: Array | DictOfNamedArrays | dict[str, Array]) DictOfNamedArrays[source]¶

Convert outputs of a computation to the canonical form.

Performs a conversion to DictOfNamedArrays if necessary.

Parameters:

result – Outputs of the computation.

loopy Code Generation¶

class pytato.target.loopy.codegen.PersistentExpressionContext(state: CodeGenState, depends_on: frozenset[str] = NOTHING)[source]¶

Mutable state used while generating loopy expressions for a ImplementedResult. Wraps CodeGenState with more expression-specific information.

This data is passed through InlinedExpressionGenMapper via arguments, and is also used by ImplementedResult.to_loopy_expression() to retrieve contextual data.

state¶

The CodeGenState.

depends_on¶

The set of statement IDs that need to be included in loopy.InstructionBase.depends_on.

update_depends_on(other: frozenset[str]) None[source]¶
class pytato.target.loopy.codegen.LocalExpressionContext(num_indices: int, local_namespace: Mapping[str, ImplementedResult], reduction_bounds: ReductionBounds, var_to_reduction_descr: Mapping[str, ReductionDescriptor])[source]¶

Records context being to be conveyed from a parent expression to its sub-expressions.

local_namespace¶

A (read-only) local name mapping used for name lookup when generating code.

num_indices¶

The number of indices of the form _0, _1, allowed in the expression.

lookup(name: str) ImplementedResult[source]¶
class pytato.target.loopy.codegen.ImplementedResult[source]¶

Generated code for a node in the computation graph (i.e., an array expression).

abstract to_loopy_expression(indices: Tuple[int | integer[Any] | Expression, ...], expr_context: PersistentExpressionContext) number[Any] | int | bool | bool | float | complex | Expression[source]¶

Return a loopy expression for this result.

Parameters:
  • indices – symbolic expressions for the indices of the array

  • expr_context –

    the associated expression context. The fields are treated as follows:

    • depends_on is populated with any dependencies needed for the generated expression.

class pytato.target.loopy.codegen.StoredResult(name: str, num_indices: int, depends_on: frozenset[str])[source]¶

An array expression generated as a loopy array.

See also: pytato.tags.ImplStored.

class pytato.target.loopy.codegen.InlinedResult(expr: number[Any] | int | bool | bool | float | complex | Expression, num_indices: int, depends_on: frozenset[str])[source]¶

An array expression generated as a loopy expression containing inlined sub-expressions.

See also: pytato.tags.ImplInlined.

class pytato.target.loopy.codegen.SubstitutionRuleResult(subst_name: str, num_args: int, depends_on: frozenset[str])[source]¶

An array expression generated as a loopy.kernel.data.SubstitutionRule.

See also: pytato.target.loopy.ImplSubstitution.

class pytato.target.loopy.codegen.CodeGenState(t_unit: TranslationUnit, results: dict[Array, ImplementedResult])[source]¶

A container for data kept by CodeGenMapper.

_t_unit¶

The partial loopy.TranslationUnit being built.

results¶

A mapping from pytato.Array instances to instances of ImplementedResult.

var_name_gen¶
insn_id_gen¶
update_kernel(kernel: LoopKernel) None[source]¶
class pytato.target.loopy.codegen.CodeGenMapper(array_tag_t_to_not_propagate: frozenset[type[Tag]], axis_tag_t_to_not_propagate: frozenset[type[Tag]])[source]¶

A mapper for generating code for nodes in the computation graph.

class pytato.target.loopy.codegen.InlinedExpressionGenMapper(axis_tag_t_to_not_propagate: frozenset[type[Tag]])[source]¶

A mapper for generating loopy expressions with inlined sub-expressions.

The inputs to this mapper are scalar expression as found in pytato.array.IndexLambda, or expressions that are compatible (e.g., shape expressions).

The outputs of this mapper are scalar expressions suitable for wrapping in InlinedResult.

pytato.target.loopy.codegen.domain_for_shape(dim_names: tuple[str, ...], shape: tuple[number[Any] | int | bool | bool | float | complex | Expression, ...], reductions: dict[str, tuple[number[Any] | int | bool | bool | float | complex | Expression, number[Any] | int | bool | bool | float | complex | Expression]]) BasicSet[source]¶

Create an islpy.BasicSet that expresses an appropriate index domain for an array of (potentially symbolic) shape shape having reduction dimensions reductions.

Parameters:
  • dim_names – A tuple of strings, the names of the axes. These become set dimensions in the returned domain.

  • shape – A tuple of constant or quasi-affine pymbolic expressions. The variables in these expressions become parameter dimensions in the returned set. Must have the same length as dim_names.

  • reductions – A map from reduction inames to (lower, upper) bounds (as half-open integer ranges). The variables in the bounds become parameter dimensions in the returned set.

pytato.target.loopy.codegen.get_loopy_temporary(name: str, expr: Array, cgen_mapper: CodeGenMapper, state: CodeGenState) TemporaryVariable[source]¶
pytato.target.loopy.codegen.add_store(name: str, expr: Array, result: ImplementedResult, state: CodeGenState, cgen_mapper: CodeGenMapper, output_to_temporary: bool = False) str[source]¶

Add an instruction that stores to a variable in the kernel.

Parameters:
  • name – name of the output array, which is created

  • expr – the Array to store

  • result – the corresponding ImplementedResult

  • state – code generation state

  • output_to_temporary – whether to generate an output argument (default) or a temporary variable

Returns:

the id of the generated instruction

pytato.target.loopy.codegen.normalize_outputs(result: Array | DictOfNamedArrays | dict[str, Array]) DictOfNamedArrays[source]¶

Convert outputs of a computation to the canonical form.

Performs a conversion to DictOfNamedArrays if necessary.

Parameters:

result – Outputs of the computation.

pytato.target.loopy.codegen.get_initial_codegen_state(target: LoopyTarget, options: Options, function_name: str) CodeGenState[source]¶
class pytato.target.loopy.codegen.ReductionBounds¶

A mapping from reduction inames to a tuple (lower_bound, upper_bound), considered half-open.