This is an archive of the discontinued LLVM Phabricator instance.

Differential D101637

[mlir][vector] Extend vector transfer unrolling to support permutations and broadcast
ClosedPublic

Authored by ThomasRaoux on Apr 30 2021, 7:58 AM.

Details

Reviewers
aartbik
nicolasvasilache
mravishankar
Commits
rGf44c76d6e919: [mlir][vector] Extend vector transfer unrolling to support permutations and…

Diff Detail

Event Timeline

ThomasRaoux created this revision.Apr 30 2021, 7:58 AM
Herald added a reviewer: mravishankar. · View Herald TranscriptApr 30 2021, 7:58 AM
Herald added subscribers: dcaballe, cota, teijeong and 16 others. · View Herald Transcript
ThomasRaoux requested review of this revision.Apr 30 2021, 7:58 AM
Herald added a project: Restricted Project. · View Herald TranscriptApr 30 2021, 7:58 AM
Herald added a subscriber: stephenneuendorffer. · View Herald Transcript
ThomasRaoux retitled this revision from [mlir][vector] Extend vector transfer unrolling to support permuations and broadcast to [mlir][vector] Extend vector transfer unrolling to support permutations and broadcast.Apr 30 2021, 8:18 AM
Harbormaster completed remote builds in B101932: Diff 341908.Apr 30 2021, 9:33 AM
nicolasvasilache accepted this revision.Apr 30 2021, 2:16 PM

stamping but we should really rewrite all this ..

This revision is now accepted and ready to land.Apr 30 2021, 2:16 PM
ThomasRaoux updated this revision to Diff 342434.May 3 2021, 9:57 AM

Fix lit test

This revision was landed with ongoing or failed builds.May 3 2021, 10:49 AM
Closed by commit rGf44c76d6e919: [mlir][vector] Extend vector transfer unrolling to support permutations and… (authored by ThomasRaoux). · Explain Why
This revision was automatically updated to reflect the committed changes.
ThomasRaoux added a commit: rGf44c76d6e919: [mlir][vector] Extend vector transfer unrolling to support permutations and….
Harbormaster completed remote builds in B102315: Diff 342434.May 3 2021, 11:06 AM

Revision Contents

PathSize
mlir/
lib/
Dialect/
Vector/
82 lines
test/
Dialect/
Vector/
93 lines

Diff 342463

mlir/lib/Dialect/Vector/VectorTransforms.cpp

Show First 20 Lines • Show All 510 Lines • ▼ Show 20 Lines for (auto it : llvm::enumerate(op->getOperandTypes()))
vectors[it.index()] = {it.value().cast<VectorType>(), indexMap, vectors[it.index()] = {it.value().cast<VectorType>(), indexMap,
static_cast<int64_t>(it.index()), false}; static_cast<int64_t>(it.index()), false};
vectors[numVectors - 1] = {resultType, indexMap, -1, false}; vectors[numVectors - 1] = {resultType, indexMap, -1, false};
resultIndex = numVectors - 1; resultIndex = numVectors - 1;
} }
/// Generates slices of 'vectorType' according to 'sizes' and 'strides, and /// Generates slices of 'vectorType' according to 'sizes' and 'strides, and
/// calls 'fn' with linear index and indices for each slice. /// calls 'fn' with linear index and indices for each slice.
static void generateTransferOpSlices( static void
Type shapedElementType, VectorType vectorType, TupleType tupleType, generateTransferOpSlices(Type shapedElementType, VectorType vectorType,
ArrayRef<int64_t> sizes, ArrayRef<int64_t> strides, ArrayRef<Value> indices, TupleType tupleType, ArrayRef<int64_t> sizes,
OpBuilder &builder, function_ref<void(unsigned, ArrayRef<Value>)> fn) { ArrayRef<int64_t> strides, ArrayRef<Value> indices,
AffineMap permutationMap, OpBuilder &builder,
function_ref<void(unsigned, ArrayRef<Value>)> fn) {
// Compute strides w.r.t. to slice counts in each dimension. // Compute strides w.r.t. to slice counts in each dimension.
auto maybeDimSliceCounts = shapeRatio(vectorType.getShape(), sizes); auto maybeDimSliceCounts = shapeRatio(vectorType.getShape(), sizes);
assert(maybeDimSliceCounts.hasValue()); assert(maybeDimSliceCounts.hasValue());
auto sliceDimCounts = *maybeDimSliceCounts; auto sliceDimCounts = *maybeDimSliceCounts;
auto sliceStrides = computeStrides(sliceDimCounts); auto sliceStrides = computeStrides(sliceDimCounts);
int64_t numSlices = tupleType.size(); int64_t numSlices = tupleType.size();
unsigned numSliceIndices = indices.size();
// Compute 'indexOffset' at which to update 'indices', which is equal // Compute 'indexOffset' at which to update 'indices', which is equal
// to the memref rank (indices.size) minus the effective 'vectorRank'. // to the memref rank (indices.size) minus the effective 'vectorRank'.
// The effective 'vectorRank', is equal to the rank of the vector type // The effective 'vectorRank', is equal to the rank of the vector type
// minus the rank of the memref vector element type (if it has one). // minus the rank of the memref vector element type (if it has one).
// //
// For example: // For example:
// //
// Given memref type 'memref<6x2x1xvector<2x4xf32>>' and vector // Given memref type 'memref<6x2x1xvector<2x4xf32>>' and vector
// transfer_read/write ops which read/write vectors of type // transfer_read/write ops which read/write vectors of type
// 'vector<2x1x2x4xf32>'. The memref rank is 3, and the effective // 'vector<2x1x2x4xf32>'. The memref rank is 3, and the effective
// vector rank is 4 - 2 = 2, and so 'indexOffset' = 3 - 2 = 1. // vector rank is 4 - 2 = 2, and so 'indexOffset' = 3 - 2 = 1.
// //
unsigned vectorRank = vectorType.getRank(); unsigned vectorRank = vectorType.getRank();
if (auto sourceVectorElementType = shapedElementType.dyn_cast<VectorType>()) { if (auto sourceVectorElementType = shapedElementType.dyn_cast<VectorType>()) {
assert(vectorRank >= sourceVectorElementType.getRank()); assert(vectorRank >= sourceVectorElementType.getRank());
vectorRank -= sourceVectorElementType.getRank(); vectorRank -= sourceVectorElementType.getRank();
} }
unsigned indexOffset = numSliceIndices - vectorRank; auto isBroadcast = [](AffineExpr expr) {
if (auto constExpr = expr.dyn_cast<AffineConstantExpr>())
return constExpr.getValue() == 0;
return false;
};
auto *ctx = builder.getContext(); auto *ctx = builder.getContext();
for (unsigned i = 0; i < numSlices; ++i) { for (unsigned i = 0; i < numSlices; ++i) {
auto vectorOffsets = delinearize(sliceStrides, i); auto vectorOffsets = delinearize(sliceStrides, i);
auto elementOffsets = auto elementOffsets =
computeElementOffsetsFromVectorSliceOffsets(sizes, vectorOffsets); computeElementOffsetsFromVectorSliceOffsets(sizes, vectorOffsets);
// Compute 'sliceIndices' by adding 'sliceOffsets[i]' to 'indices[i]'. // Compute 'sliceIndices' by adding 'sliceOffsets[i]' to 'indices[i]'.
SmallVector<Value, 4> sliceIndices(numSliceIndices); SmallVector<Value, 4> sliceIndices(indices.begin(), indices.end());
for (unsigned j = 0; j < numSliceIndices; ++j) { for (auto dim : llvm::enumerate(permutationMap.getResults())) {
if (j < indexOffset) { if (isBroadcast(dim.value()))
sliceIndices[j] = indices[j]; continue;
} else { unsigned pos = dim.value().cast<AffineDimExpr>().getPosition();
auto expr = getAffineDimExpr(0, ctx) + auto expr = getAffineDimExpr(0, ctx) +
getAffineConstantExpr(elementOffsets[j - indexOffset], ctx); getAffineConstantExpr(elementOffsets[dim.index()], ctx);
auto map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr); auto map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
sliceIndices[j] = builder.create<AffineApplyOp>( sliceIndices[pos] = builder.create<AffineApplyOp>(
indices[j].getLoc(), map, ArrayRef<Value>(indices[j])); indices[pos].getLoc(), map, ArrayRef<Value>(indices[pos]));
}
} }
// Call 'fn' to generate slice 'i' at 'sliceIndices'. // Call 'fn' to generate slice 'i' at 'sliceIndices'.
fn(i, sliceIndices); fn(i, sliceIndices);
} }
} }
/// Returns true if 'map' is a suffix of an identity affine map, false
/// otherwise. Example: affine_map<(d0, d1, d2, d3) -> (d2, d3)>
static bool isIdentitySuffix(AffineMap map) {
if (map.getNumDims() < map.getNumResults())
return false;
ArrayRef<AffineExpr> results = map.getResults();
Optional<int> lastPos;
for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
auto expr = results[i].dyn_cast<AffineDimExpr>();
if (!expr)
return false;
int currPos = static_cast<int>(expr.getPosition());
if (lastPos.hasValue() && currPos != lastPos.getValue() + 1)
return false;
lastPos = currPos;
}
return true;
}
/// Unroll transfer_read ops to the given shape and create an aggregate with all /// Unroll transfer_read ops to the given shape and create an aggregate with all
/// the chunks. /// the chunks.
static Value unrollTransferReadOp(vector::TransferReadOp readOp, static Value unrollTransferReadOp(vector::TransferReadOp readOp,
ArrayRef<int64_t> targetShape, ArrayRef<int64_t> targetShape,
OpBuilder &builder) { OpBuilder &builder) {
if (!isIdentitySuffix(readOp.permutation_map()))
return nullptr;
if (readOp.mask()) if (readOp.mask())
return nullptr; return nullptr;
auto sourceVectorType = readOp.getVectorType(); auto sourceVectorType = readOp.getVectorType();
SmallVector<int64_t, 4> strides(targetShape.size(), 1); SmallVector<int64_t, 4> strides(targetShape.size(), 1);
Location loc = readOp.getLoc(); Location loc = readOp.getLoc();
auto shapedElementType = auto shapedElementType =
readOp.source().getType().cast<ShapedType>().getElementType(); readOp.source().getType().cast<ShapedType>().getElementType();
Show All 11 Lines auto createSlice = [&](unsigned index, ArrayRef<Value> sliceIndices) {
// `in_bounds` attribute propagates conservatively: if the coarse op didn't // `in_bounds` attribute propagates conservatively: if the coarse op didn't
// need out-of-bounds masking, the fine op doesn't either. // need out-of-bounds masking, the fine op doesn't either.
vectorTupleValues[index] = builder.create<vector::TransferReadOp>( vectorTupleValues[index] = builder.create<vector::TransferReadOp>(
loc, sliceVectorType, readOp.source(), sliceIndices, loc, sliceVectorType, readOp.source(), sliceIndices,
readOp.permutation_map(), readOp.padding(), readOp.permutation_map(), readOp.padding(),
readOp.in_bounds() ? *readOp.in_bounds() : ArrayAttr()); readOp.in_bounds() ? *readOp.in_bounds() : ArrayAttr());
}; };
generateTransferOpSlices(shapedElementType, sourceVectorType, tupleType, generateTransferOpSlices(shapedElementType, sourceVectorType, tupleType,
targetShape, strides, indices, builder, createSlice); targetShape, strides, indices,
readOp.permutation_map(), builder, createSlice);
// Create tuple of splice transfer read operations. // Create tuple of splice transfer read operations.
Value tupleOp = Value tupleOp =
builder.create<vector::TupleOp>(loc, tupleType, vectorTupleValues); builder.create<vector::TupleOp>(loc, tupleType, vectorTupleValues);
// Replace 'readOp' with result 'insertSlicesResult'. // Replace 'readOp' with result 'insertSlicesResult'.
Value newVec = builder.create<vector::InsertSlicesOp>( Value newVec = builder.create<vector::InsertSlicesOp>(
loc, sourceVectorType, tupleOp, builder.getI64ArrayAttr(targetShape), loc, sourceVectorType, tupleOp, builder.getI64ArrayAttr(targetShape),
builder.getI64ArrayAttr(strides)); builder.getI64ArrayAttr(strides));
return newVec; return newVec;
} }
// Entry point for unrolling declarative pattern rewrite for transfer_write op. // Entry point for unrolling declarative pattern rewrite for transfer_write op.
LogicalResult LogicalResult
mlir::vector::unrollTransferWriteOp(OpBuilder &builder, Operation *op, mlir::vector::unrollTransferWriteOp(OpBuilder &builder, Operation *op,
ArrayRef<int64_t> targetShape, ArrayRef<int64_t> targetShape,
SmallVector<Value, 1> &result) { SmallVector<Value, 1> &result) {
auto writeOp = cast<vector::TransferWriteOp>(op); auto writeOp = cast<vector::TransferWriteOp>(op);
if (!isIdentitySuffix(writeOp.permutation_map()))
return failure();
if (writeOp.mask()) if (writeOp.mask())
return failure(); return failure();
VectorType sourceVectorType = writeOp.getVectorType(); VectorType sourceVectorType = writeOp.getVectorType();
SmallVector<int64_t, 4> strides(targetShape.size(), 1); SmallVector<int64_t, 4> strides(targetShape.size(), 1);
TupleType tupleType = generateExtractSlicesOpResultType( TupleType tupleType = generateExtractSlicesOpResultType(
sourceVectorType, targetShape, strides, builder); sourceVectorType, targetShape, strides, builder);
Location loc = writeOp.getLoc(); Location loc = writeOp.getLoc();
Value tuple = builder.create<vector::ExtractSlicesOp>( Value tuple = builder.create<vector::ExtractSlicesOp>(
Show All 12 Lines Operation *write = builder.create<vector::TransferWriteOp>(
loc, element.getResult(), loc, element.getResult(),
resultTensor ? resultTensor : writeOp.source(), sliceIndices, resultTensor ? resultTensor : writeOp.source(), sliceIndices,
writeOp.permutation_map(), writeOp.permutation_map(),
writeOp.in_bounds() ? *writeOp.in_bounds() : ArrayAttr()); writeOp.in_bounds() ? *writeOp.in_bounds() : ArrayAttr());
if (!write->getResults().empty()) if (!write->getResults().empty())
resultTensor = write->getResult(0); resultTensor = write->getResult(0);
}; };
generateTransferOpSlices(shapedElementType, sourceVectorType, tupleType, generateTransferOpSlices(shapedElementType, sourceVectorType, tupleType,
targetShape, strides, indices, builder, createSlice); targetShape, strides, indices,
writeOp.permutation_map(), builder, createSlice);
if (resultTensor) if (resultTensor)
result.push_back(resultTensor); result.push_back(resultTensor);
return success(); return success();
} }
// Entry point for unrolling declarative pattern rewrites. // Entry point for unrolling declarative pattern rewrites.
SmallVector<Value, 1> SmallVector<Value, 1>
mlir::vector::unrollSingleResultVectorOp(OpBuilder &builder, Operation *op, mlir::vector::unrollSingleResultVectorOp(OpBuilder &builder, Operation *op,
▲ Show 20 Lines • Show All 41 Lines • ▼ Show 20 Linespublic:
LogicalResult matchAndRewrite(vector::TransferReadOp readOp, LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
PatternRewriter &rewriter) const override { PatternRewriter &rewriter) const override {
if (ignoreFilter && ignoreFilter(readOp)) if (ignoreFilter && ignoreFilter(readOp))
return failure(); return failure();
if (readOp.mask()) if (readOp.mask())
return failure(); return failure();
// TODO: Support splitting TransferReadOp with non-identity permutation
// maps. Repurpose code from MaterializeVectors transformation.
if (!isIdentitySuffix(readOp.permutation_map()))
return failure();
// Return unless there is only one user, and it is an ExtractSlicesOp. // Return unless there is only one user, and it is an ExtractSlicesOp.
Value readResult = readOp.getResult(); Value readResult = readOp.getResult();
if (!readResult.hasOneUse()) if (!readResult.hasOneUse())
return failure(); return failure();
auto extractSlicesOp = auto extractSlicesOp =
dyn_cast<vector::ExtractSlicesOp>(readResult.use_begin()->getOwner()); dyn_cast<vector::ExtractSlicesOp>(readResult.use_begin()->getOwner());
if (!extractSlicesOp) if (!extractSlicesOp)
Show All 28 Linespublic:
LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp, LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
PatternRewriter &rewriter) const override { PatternRewriter &rewriter) const override {
if (ignoreFilter && ignoreFilter(writeOp)) if (ignoreFilter && ignoreFilter(writeOp))
return failure(); return failure();
if (writeOp.mask()) if (writeOp.mask())
return failure(); return failure();
// TODO: Support splitting TransferWriteOp with non-identity permutation
// maps. Repurpose code from MaterializeVectors transformation.
if (!isIdentitySuffix(writeOp.permutation_map()))
return failure();
// Fail to match unless this is writing a vector resulting from an // Fail to match unless this is writing a vector resulting from an
// InsertSlicesOp. // InsertSlicesOp.
auto insertSlicesOp = auto insertSlicesOp =
writeOp.vector().getDefiningOp<vector::InsertSlicesOp>(); writeOp.vector().getDefiningOp<vector::InsertSlicesOp>();
if (!insertSlicesOp) if (!insertSlicesOp)
return failure(); return failure();
// Get the TupleOp operand of the InsertSlicesOp. // Get the TupleOp operand of the InsertSlicesOp.
Show All 22 Lines auto createSlice = [&](unsigned index, ArrayRef<Value> sliceIndices) {
loc, tupleOp.getOperand(index), loc, tupleOp.getOperand(index),
resultTensor ? resultTensor : writeOp.source(), sliceIndices, resultTensor ? resultTensor : writeOp.source(), sliceIndices,
writeOp.permutation_map(), writeOp.permutation_map(),
writeOp.in_bounds() ? *writeOp.in_bounds() : ArrayAttr()); writeOp.in_bounds() ? *writeOp.in_bounds() : ArrayAttr());
if (!write->getResults().empty()) if (!write->getResults().empty())
resultTensor = write->getResult(0); resultTensor = write->getResult(0);
}; };
generateTransferOpSlices(shapedElementType, resultVectorType, generateTransferOpSlices(shapedElementType, resultVectorType,
sourceTupleType, sizes, strides, indices, rewriter, sourceTupleType, sizes, strides, indices,
createSlice); writeOp.permutation_map(), rewriter, createSlice);
if (resultTensor) if (resultTensor)
rewriter.replaceOp(writeOp, ArrayRef<Value>(resultTensor)); rewriter.replaceOp(writeOp, ArrayRef<Value>(resultTensor));
else else
rewriter.eraseOp(writeOp); rewriter.eraseOp(writeOp);
return success(); return success();
} }
▲ Show 20 Lines • Show All 3,018 Lines • Show Last 20 Lines

mlir/test/Dialect/Vector/vector-transfer-unroll.mlir

// RUN: mlir-opt %s -test-vector-transfer-unrolling-patterns | FileCheck %s // RUN: mlir-opt %s -test-vector-transfer-unrolling-patterns --split-input-file | FileCheck %s
// CHECK-LABEL: func @transfer_read_unroll // CHECK-LABEL: func @transfer_read_unroll
// CHECK-DAG: %[[C2:.*]] = constant 2 : index // CHECK-DAG: %[[C2:.*]] = constant 2 : index
// CHECK-DAG: %[[C0:.*]] = constant 0 : index // CHECK-DAG: %[[C0:.*]] = constant 0 : index
// CHECK: %[[VTR0:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32> // CHECK: %[[VTR0:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR1:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32> // CHECK-NEXT: %[[VTR1:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32> // CHECK-NEXT: %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32> // CHECK-NEXT: %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
▲ Show 20 Lines • Show All 105 Lines • ▼ Show 20 Lines
func @transfer_readwrite_unroll_tensor(%arg0 : tensor<4x4xf32>, %arg1 : tensor<4x4xf32>) -> func @transfer_readwrite_unroll_tensor(%arg0 : tensor<4x4xf32>, %arg1 : tensor<4x4xf32>) ->
tensor<4x4xf32> { tensor<4x4xf32> {
%c0 = constant 0 : index %c0 = constant 0 : index
%cf0 = constant 0.0 : f32 %cf0 = constant 0.0 : f32
%0 = vector.transfer_read %arg0[%c0, %c0], %cf0 : tensor<4x4xf32>, vector<4x4xf32> %0 = vector.transfer_read %arg0[%c0, %c0], %cf0 : tensor<4x4xf32>, vector<4x4xf32>
%r = vector.transfer_write %0, %arg1[%c0, %c0] : vector<4x4xf32>, tensor<4x4xf32> %r = vector.transfer_write %0, %arg1[%c0, %c0] : vector<4x4xf32>, tensor<4x4xf32>
return %r: tensor<4x4xf32> return %r: tensor<4x4xf32>
} }
// -----
// CHECK-LABEL: func @transfer_read_unroll_permutation
// CHECK-DAG: %[[C4:.*]] = constant 4 : index
// CHECK-DAG: %[[C2:.*]] = constant 2 : index
// CHECK-DAG: %[[C0:.*]] = constant 0 : index
// CHECK: %[[VTR0:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR1:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C4]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C2]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR4:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C2]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR5:.*]] = vector.transfer_read {{.*}}[%[[C4]], %[[C2]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[TUPL:.*]] = vector.tuple %[[VTR0]], %[[VTR1]], %[[VTR2]], %[[VTR3]], %[[VTR4]], %[[VTR5]] : vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VEC:.*]] = vector.insert_slices %[[TUPL]], [2, 2], [1, 1] : tuple<vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>> into vector<4x6xf32>
// CHECK-NEXT: return %[[VEC]] : vector<4x6xf32>
#map0 = affine_map<(d0, d1) -> (d1, d0)>
func @transfer_read_unroll_permutation(%arg0 : memref<6x4xf32>) -> vector<4x6xf32> {
%c0 = constant 0 : index
%cf0 = constant 0.0 : f32
%0 = vector.transfer_read %arg0[%c0, %c0], %cf0 {permutation_map = #map0} : memref<6x4xf32>, vector<4x6xf32>
return %0 : vector<4x6xf32>
}
// -----
// CHECK-LABEL: func @transfer_read_unroll_broadcast
// CHECK-DAG: %[[C2:.*]] = constant 2 : index
// CHECK-DAG: %[[C0:.*]] = constant 0 : index
// CHECK: %[[VTR0:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR1:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C2]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C2]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR4:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR5:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C2]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[TUPL:.*]] = vector.tuple %[[VTR0]], %[[VTR1]], %[[VTR2]], %[[VTR3]], %[[VTR4]], %[[VTR5]] : vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VEC:.*]] = vector.insert_slices %[[TUPL]], [2, 2], [1, 1] : tuple<vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>> into vector<6x4xf32>
// CHECK-NEXT: return %[[VEC]] : vector<6x4xf32>
#map0 = affine_map<(d0, d1) -> (0, d1)>
func @transfer_read_unroll_broadcast(%arg0 : memref<6x4xf32>) -> vector<6x4xf32> {
%c0 = constant 0 : index
%cf0 = constant 0.0 : f32
%0 = vector.transfer_read %arg0[%c0, %c0], %cf0 {permutation_map = #map0} : memref<6x4xf32>, vector<6x4xf32>
return %0 : vector<6x4xf32>
}
// -----
// CHECK-LABEL: func @transfer_read_unroll_broadcast_permuation
// CHECK-DAG: %[[C4:.*]] = constant 4 : index
// CHECK-DAG: %[[C2:.*]] = constant 2 : index
// CHECK-DAG: %[[C0:.*]] = constant 0 : index
// CHECK: %[[VTR0:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR1:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C4]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR4:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR5:.*]] = vector.transfer_read {{.*}}[%[[C4]], %[[C0]]], %{{.*}} : memref<6x4xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[TUPL:.*]] = vector.tuple %[[VTR0]], %[[VTR1]], %[[VTR2]], %[[VTR3]], %[[VTR4]], %[[VTR5]] : vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VEC:.*]] = vector.insert_slices %[[TUPL]], [2, 2], [1, 1] : tuple<vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>> into vector<4x6xf32>
// CHECK-NEXT: return %[[VEC]] : vector<4x6xf32>
#map0 = affine_map<(d0, d1) -> (0, d0)>
func @transfer_read_unroll_broadcast_permuation(%arg0 : memref<6x4xf32>) -> vector<4x6xf32> {
%c0 = constant 0 : index
%cf0 = constant 0.0 : f32
%0 = vector.transfer_read %arg0[%c0, %c0], %cf0 {permutation_map = #map0} : memref<6x4xf32>, vector<4x6xf32>
return %0 : vector<4x6xf32>
}
// -----
// CHECK-LABEL: func @transfer_read_unroll_different_rank
// CHECK-DAG: %[[C4:.*]] = constant 4 : index
// CHECK-DAG: %[[C2:.*]] = constant 2 : index
// CHECK-DAG: %[[C0:.*]] = constant 0 : index
// CHECK: %[[VTR0:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR1:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]], %[[C0]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]], %[[C2]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]], %[[C2]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR4:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]], %[[C4]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VTR5:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]], %[[C4]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[TUPL:.*]] = vector.tuple %[[VTR0]], %[[VTR1]], %[[VTR2]], %[[VTR3]], %[[VTR4]], %[[VTR5]] : vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>
// CHECK-NEXT: %[[VEC:.*]] = vector.insert_slices %[[TUPL]], [2, 2], [1, 1] : tuple<vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>, vector<2x2xf32>> into vector<6x4xf32>
// CHECK-NEXT: return %[[VEC]] : vector<6x4xf32>
#map0 = affine_map<(d0, d1, d2) -> (d2, d0)>
func @transfer_read_unroll_different_rank(%arg0 : memref<?x?x?xf32>) -> vector<6x4xf32> {
%c0 = constant 0 : index
%cf0 = constant 0.0 : f32
%0 = vector.transfer_read %arg0[%c0, %c0, %c0], %cf0 {permutation_map = #map0} : memref<?x?x?xf32>, vector<6x4xf32>
return %0 : vector<6x4xf32>
}