Internals

Generic Code Generation Functionality

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

A mapper that preprocesses graphs to simplify code generation.

The following node simplifications are performed:

Source Node Type	Target Node Type
DataWrapper	Placeholder
Roll	IndexLambda
AxisPermutation	IndexLambda
IndexBase	IndexLambda
Reshape	IndexLambda
Concatenate	IndexLambda
Einsum	IndexLambda

class pytato.codegen.PreprocessResult(outputs: 'DictOfNamedArrays', compute_order: 'Tuple[str, ...]', bound_arguments: 'Dict')

pytato.codegen.preprocess(outputs: DictOfNamedArrays, target: Target) → PreprocessResult

Preprocess a computation for code generation.

pytato.codegen.normalize_outputs(result: Array | DictOfNamedArrays | Dict[str, Array]) → DictOfNamedArrays

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)

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

class pytato.target.loopy.codegen.LocalExpressionContext(num_indices: int, local_namespace: Mapping[str, ImplementedResult], reduction_bounds: ReductionBounds, var_to_reduction_descr: Mapping[str, ReductionDescriptor])

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

class pytato.target.loopy.codegen.ImplementedResult

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

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])

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])

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])

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])

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

class pytato.target.loopy.codegen.CodeGenMapper(array_tag_t_to_not_propagate: FrozenSet[Type[Tag]], axis_tag_t_to_not_propagate: FrozenSet[Type[Tag]])

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]])

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

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

pytato.target.loopy.codegen.add_store(name: str, expr: Array, result: ImplementedResult, state: CodeGenState, cgen_mapper: CodeGenMapper, output_to_temporary: bool = False) → str

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

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

class pytato.target.loopy.codegen.ReductionBounds

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