DropUnitDims.cpp File Reference

#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/Utils/Utils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Linalg/Passes.h.inc"

Go to the source code of this file.

Namespaces
	mlir
	This header declares functions that assit transformations in the MemRef dialect.

Macros
#define	GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMS
#define	DEBUG_TYPE   "linalg-drop-unit-dims"

Functions
static ArrayAttr	replaceUnitDims (DenseSet< unsigned > &unitDims, ArrayRef< AffineMap > indexingMaps, MLIRContext *context)
	Implements a pass that canonicalizes the uses of unit-extent dimensions for broadcasting. More...
static void	replaceUnitDimIndexOps (GenericOp genericOp, const DenseSet< unsigned > &unitDims, PatternRewriter &rewriter)
	Update the index accesses of linalg operations having index semantics. More...
static std::optional< UnitExtentReplacementInfo >	replaceUnitExtents (GenericOp genericOp, OpOperand *opOperand, MLIRContext *context)
	Utility function for replacing operands/results to a linalg generic operation with unit-extent dimensions. More...

Macro Definition Documentation

DEBUG_TYPE

#define DEBUG_TYPE   "linalg-drop-unit-dims"

Definition at line 40 of file DropUnitDims.cpp.

GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMS

#define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMS

Definition at line 36 of file DropUnitDims.cpp.

Function Documentation

replaceUnitDimIndexOps()

static void replaceUnitDimIndexOps ( GenericOp  genericOp, const DenseSet< unsigned > &  unitDims, PatternRewriter &  rewriter  )

static

Update the index accesses of linalg operations having index semantics.

Definition at line 162 of file DropUnitDims.cpp.

References mlir::RewriterBase::replaceOpWithNewOp(), and mlir::OpBuilder::setInsertionPoint().

replaceUnitDims()

static ArrayAttr replaceUnitDims ( DenseSet< unsigned > &  unitDims, ArrayRef< AffineMap >  indexingMaps, MLIRContext *  context  )

static

Implements a pass that canonicalizes the uses of unit-extent dimensions for broadcasting.

For example,

```mlir #accesses = [ affine_map<(d0, d1) -> (0, d1)>, affine_map<(d0, d1) -> (d0, 0)>, affine_map<(d0, d1) -> (d0, d1)> ]

#trait = { args_in = 2, args_out = 1, indexing_maps = #accesses, iterator_types = ["parallel", "parallel"], library_call = "some_external_fn" }

func @broadcast_test(arg0 : tensor<5xf32>, arg1 : tensor<5xf32>) -> tensor<5x5xf32> { %0 = linalg.tensor_reshape arg0 [affine_map<(d0, d1) -> (d0, d1)>] : tensor<5xf32> into tensor<1x5xf32> %1 = linalg.tensor_reshape arg1 [affine_map<(d0, d1) -> (d0, d1)>] : tensor<5xf32> into tensor<5x1xf32> %2 = linalg.generic #trait %0, %1 { ^bb0(arg2: f32, arg3: f32): %3 = arith.addf arg2, arg3 : f32 linalg.yield %3 : f32 } : tensor<1x5xf32>, tensor<5x1xf32> -> tensor<5x5xf32> return %2 : tensor<5x5xf32> }

would canonicalize to

```mlir #accesses = [ affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0, d1)> ]

#trait = { args_in = 2, args_out = 1, indexing_maps = #accesses, iterator_types = ["parallel", "parallel"], library_call = "some_external_fn" }

func @broadcast_test(arg0 : tensor<5xf32>, arg1 : tensor<5xf32>) -> tensor<5x5xf32> { %0 = linalg.generic #trait arg0, arg1 { ^bb0(arg2: f32, arg3: f32): %3 = arith.addf arg2, arg3 : f32 linalg.yield %3 : f32 } : tensor<5xf32>, tensor<5xf32> -> tensor<5x5xf32> return %0 : tensor<5x5xf32> } Given dims of the iteration space of a structured op that are known to be single trip count (unitDims), return the indexing maps to use in the canonicalized op with these dims removed, given the original indexingMaps.

Definition at line 113 of file DropUnitDims.cpp.

References mlir::concatAffineMaps(), mlir::get(), mlir::getAffineConstantExpr(), mlir::getAffineDimExpr(), mlir::getAffineSymbolExpr(), mlir::inversePermutation(), and mlir::simplifyAffineMap().

replaceUnitExtents()

static std::optional<UnitExtentReplacementInfo> replaceUnitExtents ( GenericOp  genericOp, OpOperand *  opOperand, MLIRContext *  context  )

static

Utility function for replacing operands/results to a linalg generic operation with unit-extent dimensions.

These can be replaced with an operand/result with the unit-extent dimension removed. This is only done if the indexing map used to access that dimension has a AffineConstantExpr of value 0. Given the type of an result/operand of a Linalg op, and its indexMap the utility function returns:

• the new type with dimensions of size 1 removed.
• modified index map that can be used to access the replaced result/operand
• the reassociation that converts from the original tensor type to the modified tensor type.

Definition at line 373 of file DropUnitDims.cpp.

References mlir::IROperand< DerivedT, IRValueT >::get(), mlir::AffineMap::get(), mlir::getAffineConstantExpr(), mlir::AffineMap::getNumDims(), mlir::AffineMap::getNumSymbols(), mlir::AffineMap::getResults(), and mlir::Value::getType().