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/Utils/ReshapeOpsUtils.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 "mlir/Transforms/WalkPatternRewriteDriver.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Linalg/Passes.h.inc"

Go to the source code of this file.

Classes
struct	UnitExtentReplacementInfo
	Compute the modified metadata for an operands of operation whose unit dims are being dropped. More...

Namespaces
namespace	mlir
	Include the generated interface declarations.

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

Functions
static void	replaceUnitDimIndexOps (GenericOp genericOp, const llvm::SmallDenseSet< unsigned > &unitDims, RewriterBase &rewriter)
	Implements a pass that canonicalizes the uses of unit-extent dimensions for broadcasting.
static UnitExtentReplacementInfo	dropUnitExtentFromOperandMetadata (MLIRContext *context, IndexingMapOpInterface op, OpOperand *opOperand, llvm::SmallDenseMap< unsigned, unsigned > &oldDimsToNewDimsMap, ArrayRef< AffineExpr > dimReplacements)

Macro Definition Documentation

DEBUG_TYPE

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

Definition at line 39 of file DropUnitDims.cpp.

GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS

#define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS

Definition at line 35 of file DropUnitDims.cpp.

Function Documentation

dropUnitExtentFromOperandMetadata()

UnitExtentReplacementInfo dropUnitExtentFromOperandMetadata ( MLIRContext * context, IndexingMapOpInterface op, OpOperand * opOperand, llvm::SmallDenseMap< unsigned, unsigned > & oldDimsToNewDimsMap, ArrayRef< AffineExpr > dimReplacements )

static

Definition at line 345 of file DropUnitDims.cpp.

References mlir::AffineMap::get(), mlir::AffineMap::getNumSymbols(), mlir::AffineMap::getResults(), UnitExtentReplacementInfo::indexMap, UnitExtentReplacementInfo::reassociation, and UnitExtentReplacementInfo::targetShape.

Referenced by mlir::linalg::dropUnitDims().

replaceUnitDimIndexOps()

void replaceUnitDimIndexOps ( GenericOp genericOp, const llvm::SmallDenseSet< unsigned > & unitDims, RewriterBase & rewriter )

static

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

For example,

#accesses = [
  affine_map<(d0, d1) -> (0, d1)>,
  affine_map<(d0, d1) -> (d0, 0)>,
  affine_map<(d0, d1) -> (d0, d1)>
]
 
#trait = {
  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 = { 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> } Update the index accesses of linalg operations having index semantics.

Definition at line 228 of file DropUnitDims.cpp.

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

Referenced by mlir::linalg::dropUnitDims().