#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"

Macros
#define	DEBUG_TYPE "vector-contract-lowering"

Functions
static std::optional< int64_t >	getResultIndex (AffineMap map, int64_t index)
static SmallVector< Attribute >	adjustIter (ArrayAttr iteratorTypes, int64_t index)
static AffineMap	adjustMap (AffineMap map, int64_t index, PatternRewriter &rewriter)
static Value	reshapeLoad (Location loc, Value val, int64_t index, int64_t pos, PatternRewriter &rewriter)
	Returns val with the dimension at position index dropped by indexing that dimension with pos.
static Value	reshapeStore (Location loc, Value val, Value result, int64_t index, int64_t pos, PatternRewriter &rewriter)
	Inserts val into result at position pos along dimension index.
static std::optional< Value >	createContractArithOp (Location loc, Value x, Value y, Value acc, vector::CombiningKind kind, PatternRewriter &rewriter, bool isInt, Value mask=Value(), arith::FastMathFlagsAttr fmf={})
	Helper to create arithmetic operation associated with a kind of contraction.
static SmallVector< int64_t >	getReductionIndex (AffineMap map, ArrayAttr iteratorTypes)
	Return the positions of the reductions in the given map.
static std::optional< unsigned >	getDimPosition (AffineMap map, unsigned dim)
	Look for a given dimension in an affine map and return its position.
static Value	createAdd (Location loc, Value x, Value y, bool isInt, PatternRewriter &rewriter, arith::FastMathFlagsAttr fmf={})
	Creates an AddIOp if isInt is true otherwise create an arith::AddFOp using operands x and y.
FailureOr< Value >	outerProd (Value lhs, Value rhs, Value res, VectorType lhsType, int reductionSize, std::optional< Value > maybeMask=std::nullopt)
std::optional< int64_t >	getReductionSize (VectorType vecType, int64_t reductionDim)
	Helper function for matmat, matvec, tmatvec. Returns the size of dimension reductionDim. If the dimension is a scalable dimension, returns "nullopt".
FailureOr< Value >	matmat ()
	Two outer parallel, one inner reduction (matmat flavor).
FailureOr< Value >	matvec ()
FailureOr< Value >	tmatvec ()

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE "vector-contract-lowering"

Definition at line 26 of file LowerVectorContract.cpp.

Function Documentation

◆ adjustIter()

SmallVector< Attribute > adjustIter	(	ArrayAttr	iteratorTypes,
		int64_t	index )

static

Definition at line 45 of file LowerVectorContract.cpp.

References ArrayAttr().

◆ adjustMap()

AffineMap adjustMap	(	AffineMap	map,
		int64_t	index,
		PatternRewriter &	rewriter )

static

Definition at line 58 of file LowerVectorContract.cpp.

References mlir::AffineMap::get(), mlir::getAffineDimExpr(), mlir::Builder::getContext(), mlir::AffineMap::getDimPosition(), mlir::AffineMap::getNumDims(), and mlir::AffineMap::getNumResults().

◆ createAdd()

Value createAdd	(	Location	loc,
		Value	x,
		Value	y,
		bool	isInt,
		PatternRewriter &	rewriter,
		arith::FastMathFlagsAttr	fmf = {} )

static

Creates an AddIOp if isInt is true otherwise create an arith::AddFOp using operands x and y.

Definition at line 204 of file LowerVectorContract.cpp.

◆ createContractArithOp()

std::optional< Value > createContractArithOp	(	Location	loc,
		Value	x,
		Value	y,
		Value	acc,
		vector::CombiningKind	kind,
		PatternRewriter &	rewriter,
		bool	isInt,
		Value	mask = Value(),
		arith::FastMathFlagsAttr	fmf = {} )

static

Helper to create arithmetic operation associated with a kind of contraction.

Definition at line 142 of file LowerVectorContract.cpp.

◆ getDimPosition()

std::optional< unsigned > getDimPosition	(	AffineMap	map,
		unsigned	dim )

static

Look for a given dimension in an affine map and return its position.

Return std::nullopt if the dimension is not in the map results.

Definition at line 194 of file LowerVectorContract.cpp.

References mlir::AffineMap::getDimPosition(), and mlir::AffineMap::getNumResults().

◆ getReductionIndex()

SmallVector< int64_t > getReductionIndex	(	AffineMap	map,
		ArrayAttr	iteratorTypes )

static

Return the positions of the reductions in the given map.

Definition at line 182 of file LowerVectorContract.cpp.

References ArrayAttr(), mlir::AffineMap::getDimPosition(), mlir::AffineMap::getNumResults(), and mlir::vector::isReductionIterator().

◆ getReductionSize()

std::optional< int64_t > getReductionSize	(	VectorType	vecType,
		int64_t	reductionDim )

Helper function for matmat, matvec, tmatvec. Returns the size of dimension reductionDim. If the dimension is a scalable dimension, returns "nullopt".

Definition at line 429 of file LowerVectorContract.cpp.

Referenced by matmat(), matvec(), and tmatvec().

◆ getResultIndex()

std::optional< int64_t > getResultIndex	(	AffineMap	map,
		int64_t	index )

static

Definition at line 35 of file LowerVectorContract.cpp.

References mlir::AffineMap::getDimPosition(), and mlir::AffineMap::getNumResults().

◆ matmat()

FailureOr< Value > matmat ( )

Two outer parallel, one inner reduction (matmat flavor).

Definition at line 441 of file LowerVectorContract.cpp.

References mlir::bindDims(), mlir::Builder::getContext(), getReductionSize(), mlir::StructuredGenerator< StructuredOpInterface, IteratorTypeT >::iters(), mlir::StructuredGenerator< StructuredOpInterface, IteratorTypeT >::layout(), and outerProd().

◆ matvec()

FailureOr< Value > matvec ( )

Definition at line 521 of file LowerVectorContract.cpp.

References mlir::bindDims(), mlir::Builder::getContext(), getReductionSize(), mlir::StructuredGenerator< StructuredOpInterface, IteratorTypeT >::iters(), mlir::StructuredGenerator< StructuredOpInterface, IteratorTypeT >::layout(), and outerProd().

◆ outerProd()

FailureOr< Value > outerProd	(	Value	lhs,
		Value	rhs,
		Value	res,
		VectorType	lhsType,
		int	reductionSize,
		std::optional< Value >	maybeMask = std::nullopt )

Definition at line 401 of file LowerVectorContract.cpp.

References mlir::Operation::getResult(), and mlir::vector::maskOperation().

Referenced by matmat(), matvec(), and tmatvec().

◆ reshapeLoad()

Value reshapeLoad	(	Location	loc,
		Value	val,
		int64_t	index,
		int64_t	pos,
		PatternRewriter &	rewriter )

static

Returns val with the dimension at position index dropped by indexing that dimension with pos.

If index == -1, returns val unchanged. If index == 0, the result is a single vector.extract val[pos].

Example (index == 0): extract the sub-vector at pos along the leading dimension. // val : vector<4x8xf32>, pos = 2 res = vector.extract val[2] : vector<8xf32> from vector<4x8xf32>

For index > 0, recursively applies the same drop to each sub-vector of the leading dimension and reassembles the result.

Definition at line 87 of file LowerVectorContract.cpp.

References mlir::VectorType::Builder::dropDim(), mlir::Value::getType(), mlir::Builder::getZeroAttr(), load, reshapeLoad(), and result.

Referenced by reshapeLoad().

◆ reshapeStore()

Value reshapeStore	(	Location	loc,
		Value	val,
		Value	result,
		int64_t	index,
		int64_t	pos,
		PatternRewriter &	rewriter )

static

Inserts val into result at position pos along dimension index.

This is the inverse of reshapeLoad. If index == -1, returns val. If index == 0, the result is a single vector.insert val, result [pos].

Example (index == 0): insert val at pos along the leading dimension. // val : vector<4xf32>, acc : vector<2x4xf32>, pos = 1 res = vector.insert val, acc [1] : vector<4xf32> into vector<2x4xf32>

For index > 0, recursively applies the same insertion to each sub-vector of the leading dimension and reassembles the result.

Definition at line 120 of file LowerVectorContract.cpp.

References reshapeStore(), and result.

Referenced by reshapeStore().

◆ tmatvec()

FailureOr< Value > tmatvec ( )

Definition at line 564 of file LowerVectorContract.cpp.

References mlir::bindDims(), mlir::Builder::getContext(), getReductionSize(), mlir::StructuredGenerator< StructuredOpInterface, IteratorTypeT >::iters(), mlir::StructuredGenerator< StructuredOpInterface, IteratorTypeT >::layout(), and outerProd().

Macros

Functions

Macro Definition Documentation

◆ DEBUG_TYPE

Function Documentation

◆ adjustIter()

◆ adjustMap()

◆ createAdd()

◆ createContractArithOp()

◆ getDimPosition()

◆ getReductionIndex()

◆ getReductionSize()

◆ getResultIndex()

◆ matmat()

◆ matvec()

◆ outerProd()

◆ reshapeLoad()

◆ reshapeStore()

◆ tmatvec()