Internals

Generic Code Generation Functionality

class pytato.codegen.CodeGenPreprocessor(target)[source]

A mapper that preprocesses graphs to simplify code generation.

The following node simplifications are performed:

class pytato.codegen.PreprocessResult(outputs, compute_order, bound_arguments)[source]
pytato.codegen.preprocess(outputs, target)[source]

Preprocess a computation for code generation.

Return type

PreprocessResult

pytato.codegen.normalize_outputs(result)[source]

Convert outputs of a computation to the canonical form.

Performs a conversion to DictOfNamedArrays if necessary.

Parameters

result (Union[Array, DictOfNamedArrays, Dict[str, Array]]) – Outputs of the computation.

Return type

DictOfNamedArrays

loopy Code Generation

class pytato.target.loopy.codegen.PersistentExpressionContext(state, _depends_on=<factory>)[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)[source]
Return type

None

class pytato.target.loopy.codegen.LocalExpressionContext(num_indices, local_namespace, reduction_bounds)[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)[source]
Return type

Array

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, expr_context)[source]

Return a loopy expression for this result.

Parameters
  • indices (Tuple[Union[int, Expression], …]) – symbolic expressions for the indices of the array

  • expr_context (PersistentExpressionContext) –

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

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

Return type

Union[Number, bool_, bool, Expression]

class pytato.target.loopy.codegen.StoredResult(name, num_indices, depends_on)[source]

An array expression generated as a loopy array.

See also: pytato.array.ImplStored.

class pytato.target.loopy.codegen.InlinedResult(expr, num_indices, depends_on)[source]

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

See also: pytato.array.ImplInlined.

class pytato.target.loopy.codegen.CodeGenState(_program, results)[source]

A container for data kept by CodeGenMapper.

_program

The partial loopy.LoopKernel or loopy.TranslationUnit being built.

results

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

var_name_gen
insn_id_gen
update_kernel(kernel)[source]
Return type

None

class pytato.target.loopy.codegen.CodeGenMapper[source]

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

class pytato.target.loopy.codegen.InlinedExpressionGenMapper(codegen_mapper)[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, shape, reductions)[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 (Tuple[str, …]) – A tuple of strings, the names of the axes. These become set dimensions in the returned domain.

  • shape (Tuple[Union[Number, bool_, bool, Expression], …]) – 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 (Dict[str, Tuple[Union[Number, bool_, bool, Expression], Union[Number, bool_, bool, Expression]]]) – 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.

Return type

BasicSet

pytato.target.loopy.codegen.get_loopy_temporary(name, expr, cgen_mapper, state)[source]
Return type

TemporaryVariable

pytato.target.loopy.codegen.add_store(name, expr, result, state, cgen_mapper, output_to_temporary=False)[source]

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

Parameters
Return type

str

Returns

the id of the generated instruction

pytato.target.loopy.codegen.normalize_outputs(result)[source]

Convert outputs of a computation to the canonical form.

Performs a conversion to DictOfNamedArrays if necessary.

Parameters

result (Union[Array, DictOfNamedArrays, Dict[str, Array]]) – Outputs of the computation.

Return type

DictOfNamedArrays

pytato.target.loopy.codegen.get_initial_codegen_state(target, options)[source]
Return type

CodeGenState