TriangularSolve

#include <poplin/TriangularSolve.hpp>

Solving linear equations using triangular matrices.

namespace poplin

Linear algebra functions.

Functions

poplar::Tensor createTriangularSolveInputLHS(poplar::Graph &graph, const poplar::Type &inputType, const poplar::Type &outputType, const std::vector<std::size_t> &aShape, const std::vector<std::size_t> &bShape, bool leftSide, const poplar::DebugContext &debugContext, const poplar::OptionFlags &options = {}, PlanningCache *cache = nullptr)

Create a tensor that is used as the left operand of triangular solve.

This will create a 2D/3D tensor in the graph. The ordering and tile mapping of the tensor will be set to make a triangular solver with this tensor as the left argument efficient.

Supported options:

  • blockSize

    A hint for the size of block to be used. See triangularSolve() for more information.

See matMul() for additional options.

Parameters

  • graph – The Poplar graph.

  • inputType – The input data type.

  • outputType – The data type of the returned tensor.

  • aShape – The shape of the left operand.

  • bShape – The shape of the right operand.

  • leftSide – Solve AX = B if true, XA = B otherwise.

  • debugContext – Debug information.

  • options – The implementation options of the triangular solver.

  • cache – Optional pointer to a planning cache to use.

Returns

A matrix of type type and shape aShape. The tensor will have been mapped to tiles.

poplar::Tensor createTriangularSolveInputRHS(poplar::Graph &graph, const poplar::Type &inputType, const poplar::Type &outputType, const std::vector<std::size_t> &aShape, const std::vector<std::size_t> &bShape, bool leftSide, const poplar::DebugContext &debugContext, const poplar::OptionFlags &options = {}, PlanningCache *cache = nullptr)

Create a tensor that is used as the right operand of triangular solve.

This will create a 2D/3D tensor in the graph. The ordering and tile mapping of the tensor will be set to make a triangular solver with this tensor as the left argument efficient.

Supported options:

  • blockSize

    A hint for the size of block to be used. See triangularSolve() for more information.

See matMul() for additional options.

Parameters

  • graph – The Poplar graph.

  • inputType – The input data type.

  • outputType – The data type of the returned tensor.

  • aShape – The shape of the left operand.

  • bShape – The shape of the right operand.

  • leftSide – Solve AX = B if true, XA = B otherwise.

  • debugContext – Debug information.

  • options – The implementation options of the triangular solver.

  • cache – Optional pointer to a planning cache to use.

Returns

A matrix of type type and shape bShape. The tensor will have been mapped to tiles.

poplar::Tensor triangularMask(poplar::Graph &graph, const poplar::Tensor &a, bool lower, bool unitDiagonal, poplar::program::Sequence &prog, const poplar::DebugContext &debugContext = {})

Masks the unused components of the input tensor with zeroes, optionally allowing for a unit diagonal.

Parameters

  • graph – The Poplar graph.

  • a – Tensor of floating-point type with shape […, N,N].

  • lower – Whether to use the upper or lower triangle of a.

  • unitDiagonal – If true, the diagonal elements of a are assumed to be 1 and not accessed.

  • prog – A reference to a program sequence which the code to perform the arrangement will be appended to.

  • debugContext – Optional debug information.

Returns

A tensor with the same shape as a with all unused values masked.

poplar::Tensor triangularSolve(poplar::Graph &graph, const poplar::Tensor &a, const poplar::Tensor &b, bool leftSide, bool lower, bool unitDiagonal, poplar::program::Sequence &prog, const poplar::DebugContext &debugContext = {}, const poplar::OptionFlags &options = {}, PlanningCache *cache = nullptr)

Solves systems of linear equations with lower or upper triangular coefficients.

Supported options:

  • blockSize

    A hint for the size of block to be used. The triangularSolve() function uses a divide-and-conquer approach, where the input matrices are split into four quadrants that can be solved independently, and in parallel. This process is repeated until the quadrants are less than the block size. Although smaller blocks can enable more parallelism there is an overhead for combining the results from each step.

    The default value is 64.

See matMul() for additional options.

Parameters

  • graph – The Poplar graph.

  • a – Tensor of floating-point type with shape […, N,N].

  • b – Tensor of the same type with shape […, N, K] if left_side is true, […,K, N] otherwise.

  • leftSide – Solve AX = B if true, XA = B otherwise.

  • lower – Use the upper or lower triangle of a.

  • unitDiagonal – If true, the diagonal elements of a are assumed to be 1 and not accessed.

  • prog – A reference to a program sequence which the code to perform the arrangement will be appended to.

  • debugContext – Optional debug information.

  • options – The implementation options of the triangular solver.

  • cache – Optional pointer to a planning cache to use.

Returns

Tensor with shape of b with linear system solution.

std::vector<std::pair<MatMulParams, poplar::OptionFlags>> getTriangularSolveMatMulPrePlanParameters(const poplar::Type &inputType, const poplar::Type &outputType, const std::vector<std::size_t> &aShape, const std::vector<std::size_t> &bShape, bool leftSide, bool lower, const poplar::OptionFlags &options)

Plan matrix multiplication for given triangular solver.

Parameters

  • inputType – The data type of the lhs tensor.

  • outputType – The data type of the rhs tensor.

  • aShape – The shape of the left operand.

  • bShape – The shape of the right operand.

  • leftSide – Solve AX = B if true, XA = B otherwise.

  • options – The implementation options of the triangular solver.

Returns

Matmul preplan parameters.