MLIR  19.0.0git
Macros | Functions
VectorTransferSplitRewritePatterns.cpp File Reference
#include <optional>
#include <type_traits>
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/VectorInterfaces.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"

Go to the source code of this file.

Macros

#define DEBUG_TYPE   "vector-transfer-split"
 

Functions

static Value createInBoundsCond (RewriterBase &b, VectorTransferOpInterface xferOp)
 Build the condition to ensure that a particular VectorTransferOpInterface is in-bounds. More...
 
static LogicalResult splitFullAndPartialTransferPrecondition (VectorTransferOpInterface xferOp)
 Split a vector.transfer operation into an in-bounds (i.e., no out-of-bounds masking) fast path and a slow path. More...
 
static MemRefType getCastCompatibleMemRefType (MemRefType aT, MemRefType bT)
 Given two MemRefTypes aT and bT, return a MemRefType to which both can be cast. More...
 
static Value castToCompatibleMemRefType (OpBuilder &b, Value memref, MemRefType compatibleMemRefType)
 Casts the given memref to a compatible memref type. More...
 
static std::pair< Value, ValuecreateSubViewIntersection (RewriterBase &b, VectorTransferOpInterface xferOp, Value alloc)
 Operates under a scoped context to build the intersection between the view xferOp.getSource() @ xferOp.getIndices() and the view alloc. More...
 
static scf::IfOp createFullPartialLinalgCopy (RewriterBase &b, vector::TransferReadOp xferOp, TypeRange returnTypes, Value inBoundsCond, MemRefType compatibleMemRefType, Value alloc)
 Given an xferOp for which: More...
 
static scf::IfOp createFullPartialVectorTransferRead (RewriterBase &b, vector::TransferReadOp xferOp, TypeRange returnTypes, Value inBoundsCond, MemRefType compatibleMemRefType, Value alloc)
 Given an xferOp for which: More...
 
static ValueRange getLocationToWriteFullVec (RewriterBase &b, vector::TransferWriteOp xferOp, TypeRange returnTypes, Value inBoundsCond, MemRefType compatibleMemRefType, Value alloc)
 Given an xferOp for which: More...
 
static void createFullPartialLinalgCopy (RewriterBase &b, vector::TransferWriteOp xferOp, Value inBoundsCond, Value alloc)
 Given an xferOp for which: More...
 
static void createFullPartialVectorTransferWrite (RewriterBase &b, vector::TransferWriteOp xferOp, Value inBoundsCond, Value alloc)
 Given an xferOp for which: More...
 
static OperationgetAutomaticAllocationScope (Operation *op)
 

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE   "vector-transfer-split"

Definition at line 36 of file VectorTransferSplitRewritePatterns.cpp.

Function Documentation

◆ castToCompatibleMemRefType()

static Value castToCompatibleMemRefType ( OpBuilder b,
Value  memref,
MemRefType  compatibleMemRefType 
)
static

Casts the given memref to a compatible memref type.

If the source memref has a different address space than the target type, a memref.memory_space_cast is first inserted, followed by a memref.cast.

Definition at line 172 of file VectorTransferSplitRewritePatterns.cpp.

References mlir::OpBuilder::create(), mlir::get(), mlir::Value::getLoc(), and mlir::Value::getType().

Referenced by createFullPartialLinalgCopy(), createFullPartialVectorTransferRead(), and getLocationToWriteFullVec().

◆ createFullPartialLinalgCopy() [1/2]

static scf::IfOp createFullPartialLinalgCopy ( RewriterBase b,
vector::TransferReadOp  xferOp,
TypeRange  returnTypes,
Value  inBoundsCond,
MemRefType  compatibleMemRefType,
Value  alloc 
)
static

Given an xferOp for which:

  1. inBoundsCond and a compatibleMemRefType have been computed.
  2. a memref of single vector alloc has been allocated. Produce IR resembling:
    %1:3 = scf.if (%inBounds) {
    (memref.memory_space_cast %A: memref<A..., addr_space> to memref<A...>)
    %view = memref.cast %A: memref<A...> to compatibleMemRefType
    scf.yield %view, ... : compatibleMemRefType, index, index
    } else {
    %2 = linalg.fill(%pad, %alloc)
    %3 = subview %view [...][...][...]
    %4 = subview %alloc [0, 0] [...] [...]
    linalg.copy(%3, %4)
    %5 = memref.cast %alloc: memref<B...> to compatibleMemRefType
    scf.yield %5, ... : compatibleMemRefType, index, index
    }
    Return the produced scf::IfOp.

Definition at line 250 of file VectorTransferSplitRewritePatterns.cpp.

References castToCompatibleMemRefType(), mlir::OpBuilder::create(), and createSubViewIntersection().

Referenced by mlir::vector::splitFullAndPartialTransfer().

◆ createFullPartialLinalgCopy() [2/2]

static void createFullPartialLinalgCopy ( RewriterBase b,
vector::TransferWriteOp  xferOp,
Value  inBoundsCond,
Value  alloc 
)
static

Given an xferOp for which:

  1. inBoundsCond has been computed.
  2. a memref of single vector alloc has been allocated.
  3. it originally wrote to view Produce IR resembling:
    %notInBounds = arith.xori %inBounds, %true
    scf.if (%notInBounds) {
    %3 = subview %alloc [...][...][...]
    %4 = subview %view [0, 0][...][...]
    linalg.copy(%3, %4)
    }

Definition at line 394 of file VectorTransferSplitRewritePatterns.cpp.

References mlir::OpBuilder::create(), and createSubViewIntersection().

◆ createFullPartialVectorTransferRead()

static scf::IfOp createFullPartialVectorTransferRead ( RewriterBase b,
vector::TransferReadOp  xferOp,
TypeRange  returnTypes,
Value  inBoundsCond,
MemRefType  compatibleMemRefType,
Value  alloc 
)
static

Given an xferOp for which:

  1. inBoundsCond and a compatibleMemRefType have been computed.
  2. a memref of single vector alloc has been allocated. Produce IR resembling:
    %1:3 = scf.if (%inBounds) {
    (memref.memory_space_cast %A: memref<A..., addr_space> to memref<A...>)
    memref.cast %A: memref<A...> to compatibleMemRefType
    scf.yield %view, ... : compatibleMemRefType, index, index
    } else {
    %2 = vector.transfer_read %view[...], %pad : memref<A...>, vector<...>
    %3 = vector.type_cast %extra_alloc :
    memref<...> to memref<vector<...>>
    store %2, %3[] : memref<vector<...>>
    %4 = memref.cast %alloc: memref<B...> to compatibleMemRefType
    scf.yield %4, ... : compatibleMemRefType, index, index
    }
    Return the produced scf::IfOp.

Definition at line 303 of file VectorTransferSplitRewritePatterns.cpp.

References castToCompatibleMemRefType(), mlir::OpBuilder::clone(), mlir::OpBuilder::create(), mlir::get(), and mlir::Value::getType().

Referenced by mlir::vector::splitFullAndPartialTransfer().

◆ createFullPartialVectorTransferWrite()

static void createFullPartialVectorTransferWrite ( RewriterBase b,
vector::TransferWriteOp  xferOp,
Value  inBoundsCond,
Value  alloc 
)
static

Given an xferOp for which:

  1. inBoundsCond has been computed.
  2. a memref of single vector alloc has been allocated.
  3. it originally wrote to view Produce IR resembling:
    %notInBounds = arith.xori %inBounds, %true
    scf.if (%notInBounds) {
    %2 = load %alloc : memref<vector<...>>
    vector.transfer_write %2, %view[...] : memref<A...>, vector<...>
    }

Definition at line 422 of file VectorTransferSplitRewritePatterns.cpp.

References mlir::OpBuilder::clone(), mlir::OpBuilder::create(), mlir::get(), and mlir::IRMapping::map().

Referenced by mlir::vector::splitFullAndPartialTransfer().

◆ createInBoundsCond()

static Value createInBoundsCond ( RewriterBase b,
VectorTransferOpInterface  xferOp 
)
static

◆ createSubViewIntersection()

static std::pair<Value, Value> createSubViewIntersection ( RewriterBase b,
VectorTransferOpInterface  xferOp,
Value  alloc 
)
static

Operates under a scoped context to build the intersection between the view xferOp.getSource() @ xferOp.getIndices() and the view alloc.

Definition at line 191 of file VectorTransferSplitRewritePatterns.cpp.

References mlir::bindDims(), mlir::OpBuilder::create(), mlir::Builder::getContext(), mlir::Builder::getIndexAttr(), mlir::Value::getType(), and mlir::AffineMap::inferFromExprList().

Referenced by createFullPartialLinalgCopy().

◆ getAutomaticAllocationScope()

static Operation* getAutomaticAllocationScope ( Operation op)
static

◆ getCastCompatibleMemRefType()

static MemRefType getCastCompatibleMemRefType ( MemRefType  aT,
MemRefType  bT 
)
static

Given two MemRefTypes aT and bT, return a MemRefType to which both can be cast.

If the MemRefTypes don't have the same rank or are not strided, return null; otherwise:

  1. if aT and bT are cast-compatible, return aT.
  2. else return a new MemRefType obtained by iterating over the shape and strides and: a. keeping the ones that are static and equal across aT and bT. b. using a dynamic shape and/or stride for the dimensions that don't agree.

Definition at line 141 of file VectorTransferSplitRewritePatterns.cpp.

References mlir::failed(), mlir::get(), and mlir::getStridesAndOffset().

Referenced by mlir::vector::splitFullAndPartialTransfer().

◆ getLocationToWriteFullVec()

static ValueRange getLocationToWriteFullVec ( RewriterBase b,
vector::TransferWriteOp  xferOp,
TypeRange  returnTypes,
Value  inBoundsCond,
MemRefType  compatibleMemRefType,
Value  alloc 
)
static

Given an xferOp for which:

  1. inBoundsCond and a compatibleMemRefType have been computed.
  2. a memref of single vector alloc has been allocated. Produce IR resembling:
    %1:3 = scf.if (%inBounds) {
    memref.cast %A: memref<A...> to compatibleMemRefType
    scf.yield %view, ... : compatibleMemRefType, index, index
    } else {
    %3 = vector.type_cast %extra_alloc :
    memref<...> to memref<vector<...>>
    %4 = memref.cast %alloc: memref<B...> to compatibleMemRefType
    scf.yield %4, ... : compatibleMemRefType, index, index
    }

Definition at line 352 of file VectorTransferSplitRewritePatterns.cpp.

References castToCompatibleMemRefType(), and mlir::OpBuilder::create().

Referenced by mlir::vector::splitFullAndPartialTransfer().

◆ splitFullAndPartialTransferPrecondition()

static LogicalResult splitFullAndPartialTransferPrecondition ( VectorTransferOpInterface  xferOp)
static

Split a vector.transfer operation into an in-bounds (i.e., no out-of-bounds masking) fast path and a slow path.

If ifOp is not null and the result is success, theifOp` points to the newly created conditional upon function return. To accommodate for the fact that the original vector.transfer indexing may be arbitrary and the slow path indexes @[0...0] in the temporary buffer, the scf.if op returns a view and values of type index. At this time, only vector.transfer_read case is implemented.

Example (a 2-D vector.transfer_read): ``` %1 = vector.transfer_read %0[...], pad : memref<A...>, vector<...> ``` is transformed into: ``` %1:3 = scf.if (inBounds) { // fast path, direct cast memref.cast A: memref<A...> to compatibleMemRefType scf.yield view : compatibleMemRefType, index, index } else { // slow path, not in-bounds vector.transfer or linalg.copy. memref.cast alloc: memref<B...> to compatibleMemRefType scf.yield %4 : compatibleMemRefType, index, index */ // } /** %0 = vector.transfer_read %1#0[%1#1, %1#2] {in_bounds = [true ... true]} `` wherealloc` is a top of the function alloca'ed buffer of one vector.

Preconditions:

  1. xferOp.getPermutationMap() must be a minor identity map
  2. the rank of the xferOp.memref() and the rank of the xferOp.getVector() must be equal. This will be relaxed in the future but requires rank-reducing subviews.

Definition at line 113 of file VectorTransferSplitRewritePatterns.cpp.

References mlir::failure(), and mlir::success().

Referenced by mlir::vector::splitFullAndPartialTransfer().