MLIR  16.0.0git
Classes | Functions
TensorInferTypeOpInterfaceImpl.cpp File Reference
#include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
+ Include dependency graph for TensorInferTypeOpInterfaceImpl.cpp:

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< 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< 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::applyMapToValues(), mlir::OpBuilder::createOrFold(), mlir::AffineMap::get(), mlir::Builder::getAffineSymbolExpr(), and mlir::AffineMap::getResults().

◆ 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 62 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 22 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 74 of file TensorInferTypeOpInterfaceImpl.cpp.

References mlir::applyMapToValues(), mlir::OpBuilder::create(), mlir::detail::enumerate(), mlir::AffineExpr::floorDiv(), mlir::AffineMap::get(), mlir::Builder::getAffineSymbolExpr(), and mlir::Builder::getI64IntegerAttr().

◆ 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 112 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