MLIR  15.0.0git
Macros | Functions
DropUnitDims.cpp File Reference
#include "PassDetail.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.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/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
+ Include dependency graph for DropUnitDims.cpp:

Go to the source code of this file.

Macros

#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 llvm::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 30 of file DropUnitDims.cpp.

Function Documentation

◆ replaceUnitDimIndexOps()

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

◆ 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 (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 (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 99 of file DropUnitDims.cpp.

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

Referenced by replaceUnitDimIndexOps().

◆ replaceUnitExtents()

static llvm::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 didimensionmension 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 240 of file DropUnitDims.cpp.

References mlir::concatAffineMaps(), mlir::OpBuilder::create(), mlir::Type::dyn_cast(), mlir::failure(), mlir::AffineMap::get(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::getAffineConstantExpr(), mlir::getAffineDimExpr(), mlir::Builder::getContext(), mlir::getElementTypeOrSelf(), mlir::AffineMap::getNumDims(), mlir::AffineMap::getNumSymbols(), mlir::AffineMap::getResults(), mlir::Value::getType(), mlir::inversePermutation(), and mlir::Type::isa().