MLIR  16.0.0git
Namespaces | Macros | Functions
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/Passes.h"
#include "mlir/Dialect/Tensor/IR/Tensor.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"
+ Include dependency graph for DropUnitDims.cpp:

Go to the source code of this file.

Namespaces

 mlir
 Include the generated interface declarations.
 

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 39 of file DropUnitDims.cpp.

◆ GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMS

#define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMS

Definition at line 35 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 157 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 108 of file DropUnitDims.cpp.

References mlir::concatAffineMaps(), 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 368 of file DropUnitDims.cpp.

References mlir::Type::cast(), mlir::Type::dyn_cast(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::AffineMap::get(), mlir::getAffineConstantExpr(), mlir::getAffineDimExpr(), mlir::getElementTypeOrSelf(), mlir::AffineMap::getNumDims(), mlir::AffineMap::getNumSymbols(), mlir::AffineMap::getResults(), mlir::Value::getType(), and mlir::Type::isa().