TensorInferTypeOpInterfaceImpl.cpp File Reference

#include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"

Go to the source code of this file.

Classes
struct	ReifyExpandOrCollapseShapeOp< OpTy >

Functions
static llvm::DenseMap< int64_t, int64_t >	getExpandedDimToCollapsedDimMap (ArrayRef< AffineMap > reassociation)
	Compute a map that for a given dimension of the expanded type gives the dimension in the collapsed type it maps to. More...
static OpFoldResult	getCollapsedOutputDimFromInputShape (OpBuilder &builder, Location loc, int64_t dimIndex, Value src, ArrayRef< int64_t > dstStaticShape, ArrayRef< AffineMap > reassociationMap)
	For reshape op compute the shape at dimension dimIndex of the output in terms of shape of the src, when the reshape op is a collapsing operation. More...
static SmallVector< OpFoldResult, 4 >	getCollapsedOutputShapeFromInputShape (OpBuilder &builder, Location loc, Value src, ArrayRef< int64_t > dstStaticShape, ArrayRef< AffineMap > reassociation)
	Given the src of a collapsing reshape op and its reassociation maps, compute the shape of the result of the reshape. More...
static OpFoldResult	getExpandedOutputDimFromInputShape (OpBuilder &builder, Location loc, int64_t dimIndex, Value src, ArrayRef< int64_t > dstStaticShape, ArrayRef< AffineMap > reassociation, llvm::DenseMap< int64_t, int64_t > &expandedDimToCollapsedDim)
	For an expanding reshape op, compute the value for a dimension of the output from the shape of the input. More...
static SmallVector< OpFoldResult, 4 >	getExpandedOutputShapeFromInputShape (OpBuilder &builder, Location loc, Value src, ArrayRef< int64_t > dstStaticShape, ArrayRef< AffineMap > reassociation)
	Given the src of an expanding reshape op, the reassociation maps and the result type, compute the shape of the result of the reshape. More...
static SmallVector< OpFoldResult, 4 >	getReshapeOutputShapeFromInputShape (OpBuilder &builder, Location loc, Value src, ArrayRef< int64_t > dstStaticShape, ArrayRef< AffineMap > reassocation)

Function Documentation

getCollapsedOutputDimFromInputShape()

static OpFoldResult getCollapsedOutputDimFromInputShape ( OpBuilder &  builder, Location  loc, int64_t  dimIndex, Value  src, ArrayRef< int64_t >  dstStaticShape, ArrayRef< AffineMap >  reassociationMap  )

static

For reshape op compute the shape at dimension dimIndex of the output in terms of shape of the src, when the reshape op is a collapsing operation.

It is the product of the shape of the collapsed dimensions of the src.

Definition at line 41 of file TensorInferTypeOpInterfaceImpl.cpp.

References mlir::OpBuilder::createOrFold(), mlir::AffineMap::get(), mlir::Builder::getAffineSymbolExpr(), mlir::Builder::getIndexAttr(), mlir::AffineMap::getResults(), and mlir::affine::makeComposedAffineApply().

getCollapsedOutputShapeFromInputShape()

static SmallVector<OpFoldResult, 4> getCollapsedOutputShapeFromInputShape ( OpBuilder &  builder, Location  loc, Value  src, ArrayRef< int64_t >  dstStaticShape, ArrayRef< AffineMap >  reassociation  )

static

Given the src of a collapsing reshape op and its reassociation maps, compute the shape of the result of the reshape.

Definition at line 69 of file TensorInferTypeOpInterfaceImpl.cpp.

Referenced by getReshapeOutputShapeFromInputShape().

getExpandedDimToCollapsedDimMap()

static llvm::DenseMap<int64_t, int64_t> getExpandedDimToCollapsedDimMap ( ArrayRef< AffineMap >  reassociation )

static

Compute a map that for a given dimension of the expanded type gives the dimension in the collapsed type it maps to.

Essentially its the inverse of the reassocation maps.

Definition at line 23 of file TensorInferTypeOpInterfaceImpl.cpp.

References mlir::detail::enumerate().

Referenced by getExpandedOutputShapeFromInputShape().

getExpandedOutputDimFromInputShape()

static OpFoldResult getExpandedOutputDimFromInputShape ( OpBuilder &  builder, Location  loc, int64_t  dimIndex, Value  src, ArrayRef< int64_t >  dstStaticShape, ArrayRef< AffineMap >  reassociation, llvm::DenseMap< int64_t, int64_t > &  expandedDimToCollapsedDim  )

static

For an expanding reshape op, compute the value for a dimension of the output from the shape of the input.

Definition at line 81 of file TensorInferTypeOpInterfaceImpl.cpp.

References mlir::OpBuilder::create(), mlir::detail::enumerate(), mlir::AffineExpr::floorDiv(), mlir::AffineMap::get(), mlir::Builder::getAffineSymbolExpr(), mlir::Builder::getIndexAttr(), and mlir::affine::makeComposedAffineApply().

getExpandedOutputShapeFromInputShape()

static SmallVector<OpFoldResult, 4> getExpandedOutputShapeFromInputShape ( OpBuilder &  builder, Location  loc, Value  src, ArrayRef< int64_t >  dstStaticShape, ArrayRef< AffineMap >  reassociation  )

static

Given the src of an expanding reshape op, the reassociation maps and the result type, compute the shape of the result of the reshape.

Definition at line 121 of file TensorInferTypeOpInterfaceImpl.cpp.

References getExpandedDimToCollapsedDimMap().

Referenced by getReshapeOutputShapeFromInputShape().

getReshapeOutputShapeFromInputShape()

static SmallVector<OpFoldResult, 4> getReshapeOutputShapeFromInputShape ( OpBuilder &  builder, Location  loc, Value  src, ArrayRef< int64_t >  dstStaticShape, ArrayRef< AffineMap >  reassocation  )

static

Definition at line 135 of file TensorInferTypeOpInterfaceImpl.cpp.

References getCollapsedOutputShapeFromInputShape(), getExpandedOutputShapeFromInputShape(), and mlir::Value::getType().

Referenced by ReifyExpandOrCollapseShapeOp< OpTy >::reifyResultShapes().