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:
Source Node Type
Target Node Type
- 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 aImplementedResult
. WrapsCodeGenState
with more expression-specific information.This data is passed through
InlinedExpressionGenMapper
via arguments, and is also used byImplementedResult.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
.
- 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 ofImplementedResult
.
- 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 storeresult – 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.