MLIR  19.0.0git
Namespaces | Macros | Functions
FoldMemRefAliasOps.cpp File Reference
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"

Go to the source code of this file.

Namespaces

 mlir
 Include the generated interface declarations.
 
 mlir::memref
 

Macros

#define DEBUG_TYPE   "fold-memref-alias-ops"
 
#define DBGS()   (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 
#define GEN_PASS_DEF_FOLDMEMREFALIASOPS
 

Functions

static LogicalResult resolveSourceIndicesExpandShape (Location loc, PatternRewriter &rewriter, memref::ExpandShapeOp expandShapeOp, ValueRange indices, SmallVectorImpl< Value > &sourceIndices)
 Given the 'indices' of a load/store operation where the memref is a result of a expand_shape op, returns the indices w.r.t to the source memref of the expand_shape op. More...
 
static LogicalResult resolveSourceIndicesCollapseShape (Location loc, PatternRewriter &rewriter, memref::CollapseShapeOp collapseShapeOp, ValueRange indices, SmallVectorImpl< Value > &sourceIndices)
 Given the 'indices' of a load/store operation where the memref is a result of a collapse_shape op, returns the indices w.r.t to the source memref of the collapse_shape op. More...
 
template<typename LoadOrStoreOpTy >
static Value getMemRefOperand (LoadOrStoreOpTy op)
 Helpers to access the memref operand for each op. More...
 
static Value getMemRefOperand (vector::TransferReadOp op)
 
static Value getMemRefOperand (nvgpu::LdMatrixOp op)
 
static Value getMemRefOperand (vector::LoadOp op)
 
static Value getMemRefOperand (vector::StoreOp op)
 
static Value getMemRefOperand (vector::MaskedLoadOp op)
 
static Value getMemRefOperand (vector::MaskedStoreOp op)
 
static Value getMemRefOperand (vector::TransferWriteOp op)
 
static Value getMemRefOperand (gpu::SubgroupMmaLoadMatrixOp op)
 
static Value getMemRefOperand (gpu::SubgroupMmaStoreMatrixOp op)
 
static SmallVector< ValuecalculateExpandedAccessIndices (AffineMap affineMap, const SmallVector< Value > &indices, Location loc, PatternRewriter &rewriter)
 
template<typename XferOp >
static LogicalResult preconditionsFoldSubViewOpImpl (RewriterBase &rewriter, XferOp xferOp, memref::SubViewOp subviewOp)
 
static LogicalResult preconditionsFoldSubViewOp (RewriterBase &rewriter, Operation *op, memref::SubViewOp subviewOp)
 
static LogicalResult preconditionsFoldSubViewOp (RewriterBase &rewriter, vector::TransferReadOp readOp, memref::SubViewOp subviewOp)
 
static LogicalResult preconditionsFoldSubViewOp (RewriterBase &rewriter, vector::TransferWriteOp writeOp, memref::SubViewOp subviewOp)
 

Macro Definition Documentation

◆ DBGS

#define DBGS ( )    (llvm::dbgs() << "[" DEBUG_TYPE "]: ")

Definition at line 33 of file FoldMemRefAliasOps.cpp.

◆ DEBUG_TYPE

#define DEBUG_TYPE   "fold-memref-alias-ops"

Definition at line 32 of file FoldMemRefAliasOps.cpp.

◆ GEN_PASS_DEF_FOLDMEMREFALIASOPS

#define GEN_PASS_DEF_FOLDMEMREFALIASOPS

Definition at line 37 of file FoldMemRefAliasOps.cpp.

Function Documentation

◆ calculateExpandedAccessIndices()

static SmallVector<Value> calculateExpandedAccessIndices ( AffineMap  affineMap,
const SmallVector< Value > &  indices,
Location  loc,
PatternRewriter rewriter 
)
static

◆ getMemRefOperand() [1/10]

static Value getMemRefOperand ( gpu::SubgroupMmaLoadMatrixOp  op)
static

Definition at line 200 of file FoldMemRefAliasOps.cpp.

◆ getMemRefOperand() [2/10]

static Value getMemRefOperand ( gpu::SubgroupMmaStoreMatrixOp  op)
static

Definition at line 204 of file FoldMemRefAliasOps.cpp.

◆ getMemRefOperand() [3/10]

template<typename LoadOrStoreOpTy >
static Value getMemRefOperand ( LoadOrStoreOpTy  op)
static

Helpers to access the memref operand for each op.

Definition at line 176 of file FoldMemRefAliasOps.cpp.

◆ getMemRefOperand() [4/10]

static Value getMemRefOperand ( nvgpu::LdMatrixOp  op)
static

Definition at line 184 of file FoldMemRefAliasOps.cpp.

◆ getMemRefOperand() [5/10]

static Value getMemRefOperand ( vector::LoadOp  op)
static

Definition at line 188 of file FoldMemRefAliasOps.cpp.

◆ getMemRefOperand() [6/10]

static Value getMemRefOperand ( vector::MaskedLoadOp  op)
static

Definition at line 192 of file FoldMemRefAliasOps.cpp.

◆ getMemRefOperand() [7/10]

static Value getMemRefOperand ( vector::MaskedStoreOp  op)
static

Definition at line 194 of file FoldMemRefAliasOps.cpp.

◆ getMemRefOperand() [8/10]

static Value getMemRefOperand ( vector::StoreOp  op)
static

Definition at line 190 of file FoldMemRefAliasOps.cpp.

◆ getMemRefOperand() [9/10]

static Value getMemRefOperand ( vector::TransferReadOp  op)
static

Definition at line 180 of file FoldMemRefAliasOps.cpp.

◆ getMemRefOperand() [10/10]

static Value getMemRefOperand ( vector::TransferWriteOp  op)
static

Definition at line 196 of file FoldMemRefAliasOps.cpp.

◆ preconditionsFoldSubViewOp() [1/3]

static LogicalResult preconditionsFoldSubViewOp ( RewriterBase rewriter,
Operation op,
memref::SubViewOp  subviewOp 
)
static

Definition at line 360 of file FoldMemRefAliasOps.cpp.

◆ preconditionsFoldSubViewOp() [2/3]

static LogicalResult preconditionsFoldSubViewOp ( RewriterBase rewriter,
vector::TransferReadOp  readOp,
memref::SubViewOp  subviewOp 
)
static

Definition at line 366 of file FoldMemRefAliasOps.cpp.

References preconditionsFoldSubViewOpImpl().

◆ preconditionsFoldSubViewOp() [3/3]

static LogicalResult preconditionsFoldSubViewOp ( RewriterBase rewriter,
vector::TransferWriteOp  writeOp,
memref::SubViewOp  subviewOp 
)
static

Definition at line 372 of file FoldMemRefAliasOps.cpp.

References preconditionsFoldSubViewOpImpl().

◆ preconditionsFoldSubViewOpImpl()

template<typename XferOp >
static LogicalResult preconditionsFoldSubViewOpImpl ( RewriterBase rewriter,
XferOp  xferOp,
memref::SubViewOp  subviewOp 
)
static

◆ resolveSourceIndicesCollapseShape()

static LogicalResult resolveSourceIndicesCollapseShape ( Location  loc,
PatternRewriter rewriter,
memref::CollapseShapeOp  collapseShapeOp,
ValueRange  indices,
SmallVectorImpl< Value > &  sourceIndices 
)
static

Given the 'indices' of a load/store operation where the memref is a result of a collapse_shape op, returns the indices w.r.t to the source memref of the collapse_shape op.

For example

%0 = ... : memref<2x6x42xf32> %1 = memref.collapse_shape %0 [[0, 1], [2]] : memref<2x6x42xf32> into memref<12x42xf32> %2 = load %1[i1, i2] : memref<12x42xf32>

could be folded into

%2 = load %0[i1 / 6, i1 % 6, i2] : memref<2x6x42xf32>

Definition at line 119 of file FoldMemRefAliasOps.cpp.

References mlir::computeSuffixProduct(), mlir::delinearize(), mlir::failure(), mlir::AffineMap::get(), mlir::Builder::getAffineDimExpr(), mlir::Builder::getConstantAffineMap(), mlir::getValueOrCreateConstantIndexOp(), mlir::affine::makeComposedFoldedAffineApply(), and mlir::success().

◆ resolveSourceIndicesExpandShape()

static LogicalResult resolveSourceIndicesExpandShape ( Location  loc,
PatternRewriter rewriter,
memref::ExpandShapeOp  expandShapeOp,
ValueRange  indices,
SmallVectorImpl< Value > &  sourceIndices 
)
static

Given the 'indices' of a load/store operation where the memref is a result of a expand_shape op, returns the indices w.r.t to the source memref of the expand_shape op.

For example

%0 = ... : memref<12x42xf32> %1 = memref.expand_shape %0 [[0, 1], [2]] : memref<12x42xf32> into memref<2x6x42xf32> %2 = load %1[i1, i2, i3] : memref<2x6x42xf32

could be folded into

%2 = load %0[6 * i1 + i2, i3] : memref<12x42xf32>

Apply permutation and create AffineApplyOp.

Definition at line 62 of file FoldMemRefAliasOps.cpp.

References mlir::bindDimsList(), mlir::computeSuffixProduct(), mlir::failure(), mlir::AffineMap::get(), mlir::Builder::getContext(), mlir::getValueOrCreateConstantIndexOp(), mlir::linearize(), mlir::affine::makeComposedFoldedAffineApply(), and mlir::success().