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 = {}, matmul::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.

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 overwise.

  • debugContext – Debug information.

  • options – The implementation options of the triangular solver. Supported options: ‘blockSize’ - blockSize hint See matMul() for additional options.

  • 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 = {}, matmul::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.

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 overwise.

  • debugContext – Debug information.

  • options – The implementation options of the triangular solver. Supported options: ‘blockSize’ - blockSize hint See matMul() for additional options.

  • 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 = {}, matmul::PlanningCache *cache = nullptr)

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

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 overwise.

  • 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. Supported options: ‘blockSize’ - blockSize hint See matMul() for additional options.

  • 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 overwise.

  • options – The implementation options of the triangular solver.

Returns

Mat mul preplan parameters.