#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineExprVisitor.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/TypeUtilities.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetOperations.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include <optional>
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.cpp.inc"

Classes
struct	HasAffineDimExprVisitor
	Visitor to check if any of the given set of positions from AffineDimExprs are used within an AffineExpr. More...

Namespaces
namespace	mlir
	Include the generated interface declarations.
namespace	mlir::linalg
namespace	mlir::linalg::detail

Enumerations
enum class	mlir::linalg::detail::MatchContractionResult { mlir::linalg::detail::Success = 0 , mlir::linalg::detail::NotLinalgOp , mlir::linalg::detail::WrongNumOperands , mlir::linalg::detail::NoReduction , mlir::linalg::detail::NotProjectedPermutations , mlir::linalg::detail::NotAddMul }
enum class	mlir::linalg::detail::MatchConvolutionResult { mlir::linalg::detail::Success = 0 , mlir::linalg::detail::NotLinalgOp , mlir::linalg::detail::WrongNumOperands , mlir::linalg::detail::WrongInputIndexingMap , mlir::linalg::detail::NotProjectedPermutations , mlir::linalg::detail::NonConvolutionLoop , mlir::linalg::detail::OutputDimsNotParallel , mlir::linalg::detail::NonOutputDimNotReduction , mlir::linalg::detail::EmptyConvolvedDims }

Functions
static std::optional< Value >	isaInlinedFillOp (GenericOp op)
	Detects if a linalg.generic operation represents a fill with an inlined constant.
static std::optional< Value >	isaExternalFillOp (GenericOp op)
	Detects if a linalg.generic operation represents an external scalar input.
static bool	isaElemwiseSingleUnaryOrBinaryOpInterface (linalg::GenericOp op, unsigned arity)
static Value	getSourceSkipUnary (Value value)
	If the value is defined by a chain of unary side effect-free, go up the use-def chain until the first value that isn't defined by such an op.
template<typename AddOpTy, typename MulOpTy, typename... Args>
static bool	isPairTemplateImpl (Operation add, Operation mul)
	Returns true if the two operations are of the kinds specified by a pair of consecutive template arguments.
template<typename... Args>
static bool	isContractionBody (Block &block)
	Returns true if the block is a body of a contraction with the kinds of operations given pairwise by template arguments.
static llvm::SmallDenseSet< int64_t >	findPermutationsIndexingOperand (AffineMap indexingMap, ArrayRef< utils::IteratorType > iterators, utils::IteratorType iter)
	Given an indexingMap and its corresponding iterators, returns the positions of the iterators of type iter that are indexed by the indexingMap as a permutation.
static FailureOr< SmallVector< utils::IteratorType > >	inferIteratorsFromOutMap (AffineMap map)
	Infer the iterator types from the init affine map.
static FailureOr< ContractionDimensions >	inferContractionDimsImpl (ArrayRef< AffineMap > indexingMaps, ArrayRef< utils::IteratorType > iterators)
	Find 2 parallel (m and n) and 1 reduction (k) dimension candidates that form a matmul subcomputation within linalgOp.
template<typename T>
static T	getAffineExprOfType (AffineExpr lhs, AffineExpr rhs)
	Of the given two expressions returns one that is of type T (lhs gets preference over rhs)
static llvm::SmallDenseSet< int64_t >	getPreservedDims (AffineMap map)
static SmallVector< int64_t, 2 >	getConstantsFromExprList (const SmallVector< AffineExpr, 2 > &exprs)
static FailureOr< ConvolutionDimensions >	inferConvolutionDimsImpl (LinalgOp linalgOp, ConvAccessExprWalker &inputExprWalker, bool allowEmptyConvolvedDims)
	Classifies dimensions in the linalgOp used by a convolution subcomputation, as captured by inputExprWalker.
static MatchFillResult	isFillInterfaceImpl (Operation *op)
static std::pair< int64_t, int64_t >	getResultsPositionInLoopsToShapeMap (LinalgOp &op)

Function Documentation

◆ findPermutationsIndexingOperand()

llvm::SmallDenseSet< int64_t > findPermutationsIndexingOperand	(	AffineMap	indexingMap,
		ArrayRef< utils::IteratorType >	iterators,
		utils::IteratorType	iter )

static

Given an indexingMap and its corresponding iterators, returns the positions of the iterators of type iter that are indexed by the indexingMap as a permutation.

This is useful to infer various subcomputations on a LinalgOp. This is performed by looking up each result in the indexingMap and determining whether:

It is a single AffineDimExpr.
It is the only result involving this AffineDimExpr.

Definition at line 400 of file LinalgInterfaces.cpp.

References mlir::AffineMap::getNumDims(), and mlir::AffineMap::getResults().

Referenced by inferContractionDimsImpl(), and inferConvolutionDimsImpl().

◆ getAffineExprOfType()

template<typename T>

T getAffineExprOfType	(	AffineExpr	lhs,
		AffineExpr	rhs )

static

Of the given two expressions returns one that is of type T (lhs gets preference over rhs)

Definition at line 609 of file LinalgInterfaces.cpp.

References lhs, and rhs.

◆ getConstantsFromExprList()

SmallVector< int64_t, 2 > getConstantsFromExprList ( const SmallVector< AffineExpr, 2 > & exprs )

static

Definition at line 731 of file LinalgInterfaces.cpp.

Referenced by inferConvolutionDimsImpl().

◆ getPreservedDims()

llvm::SmallDenseSet< int64_t > getPreservedDims ( AffineMap map )

static

Definition at line 721 of file LinalgInterfaces.cpp.

References mlir::AffineMap::getResults(), and mlir::AffineMap::isProjectedPermutation().

Referenced by mlir::linalg::detail::isConvolutionInterfaceImpl().

◆ getResultsPositionInLoopsToShapeMap()

std::pair< int64_t, int64_t > getResultsPositionInLoopsToShapeMap ( LinalgOp & op )

static

Definition at line 1199 of file LinalgInterfaces.cpp.

◆ getSourceSkipUnary()

Value getSourceSkipUnary ( Value value )

static

If the value is defined by a chain of unary side effect-free, go up the use-def chain until the first value that isn't defined by such an op.

Definition at line 293 of file LinalgInterfaces.cpp.

References mlir::Value::getDefiningOp(), mlir::Operation::getNumOperands(), and mlir::Operation::getOperand().

Referenced by mlir::linalg::detail::isContractionBody().

◆ inferContractionDimsImpl()

FailureOr< ContractionDimensions > inferContractionDimsImpl	(	ArrayRef< AffineMap >	indexingMaps,
		ArrayRef< utils::IteratorType >	iterators )

static

Find 2 parallel (m and n) and 1 reduction (k) dimension candidates that form a matmul subcomputation within linalgOp.

These dimensions are such that:

The m dimension is involved in an outer-product along LHS (i.e. it is a permutation on RES and LHS and does not appear in RHS).
The n dimension is involved in an outer-product along RHS (i.e. it is a permutation on RES and RHS and does not appear in LHS).
The k dimension appears as a permutation on LHS and RHS.
m, n and k appear only once in any given indexing.
Optional batch dimensions that appear in all operands are captured. This allows e.g. detecting that some contraction is embedded within linalgOp with some orthogonal heuristic.

Definition at line 449 of file LinalgInterfaces.cpp.

References b, mlir::linalg::ContractionDimensions::batch, findPermutationsIndexingOperand(), mlir::linalg::ContractionDimensions::k, mlir::linalg::ContractionDimensions::m, and mlir::linalg::ContractionDimensions::n.

Referenced by mlir::linalg::inferContractionDims(), and mlir::linalg::inferContractionDims().

◆ inferConvolutionDimsImpl()

FailureOr< ConvolutionDimensions > inferConvolutionDimsImpl	(	LinalgOp	linalgOp,
		ConvAccessExprWalker &	inputExprWalker,
		bool	allowEmptyConvolvedDims )

static

Classifies dimensions in the linalgOp used by a convolution subcomputation, as captured by inputExprWalker.

If allowEmptyConvolvedDims is not set this will fail if there is not at least one convolved dimension pair (output image + filter loop).

The returned dimensions are ordered as follows:

outputImage is sorted by dimension index.
filterLoop is ordered to match the pairing with outputImage, i.e., outputImage[i] and filterLoop[i] are paired dimensions from the convolution access pattern (e.g., oh + kh pairs oh with kh).
strides[i] corresponds to outputImage[i].
dilations[i] corresponds to filterLoop[i].
Other dimension sets (batch, outputChannel, etc.) are sorted by index.

Definition at line 755 of file LinalgInterfaces.cpp.

References mlir::linalg::ConvolutionDimensions::batch, mlir::linalg::ConvolutionDimensions::depth, mlir::linalg::ConvolutionDimensions::dilations, mlir::linalg::ConvolutionDimensions::filterLoop, findPermutationsIndexingOperand(), getConstantsFromExprList(), mlir::linalg::ConvolutionDimensions::inputChannel, mlir::linalg::ConvolutionDimensions::outputChannel, mlir::linalg::ConvolutionDimensions::outputImage, and mlir::linalg::ConvolutionDimensions::strides.

Referenced by mlir::linalg::inferConvolutionDims(), and mlir::linalg::detail::isConvolutionInterfaceImpl().

◆ inferIteratorsFromOutMap()

FailureOr< SmallVector< utils::IteratorType > > inferIteratorsFromOutMap ( AffineMap map )

static

Infer the iterator types from the init affine map.

This looks at which dims are present in the map results, and returns an iterator types array with parallel types for dims that are present, and reduction types for dims that are not present.

Definition at line 427 of file LinalgInterfaces.cpp.

References mlir::AffineMap::getNumDims(), mlir::AffineMap::getResults(), and mlir::AffineMap::isProjectedPermutation().

Referenced by mlir::linalg::inferContractionDims().