This is an archive of the discontinued LLVM Phabricator instance.

Differential D144989

[MLIR][Linalg] Fix propagation for rank-zero tensor
ClosedPublic

Authored by chelini on Feb 28 2023, 10:44 AM.

Details

Reviewers
hanchung
nicolasvasilache
Commits
rG5885c85fc6a7: [MLIR][Linalg] Fix propagation for rank-zero tensor
Summary

isScalar only returns true if the operand is non-shaped.
But we need to handle also rank zero tensors.

Diff Detail

Event Timeline

chelini created this revision.Feb 28 2023, 10:44 AM
Herald added a reviewer: hanchung. · View Herald TranscriptFeb 28 2023, 10:44 AM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: hanchung, Moerafaat, bzcheeseman and 23 others. · View Herald Transcript
chelini requested review of this revision.Feb 28 2023, 10:44 AM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptFeb 28 2023, 10:44 AM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: limo1996, stephenneuendorffer, nicolasvasilache. · View Herald Transcript
Harbormaster completed remote builds in B216531: Diff 501222.Feb 28 2023, 10:44 AM
hanchung accepted this revision.Feb 28 2023, 1:30 PM
This revision is now accepted and ready to land.Feb 28 2023, 1:30 PM
This revision was landed with ongoing or failed builds.Mar 1 2023, 12:31 AM
Closed by commit rG5885c85fc6a7: [MLIR][Linalg] Fix propagation for rank-zero tensor (authored by chelini). · Explain Why
This revision was automatically updated to reflect the committed changes.
chelini added a commit: rG5885c85fc6a7: [MLIR][Linalg] Fix propagation for rank-zero tensor.

Revision Contents

PathSize
mlir/
lib/
Dialect/
Linalg/
Transforms/
2 lines
test/
Dialect/
Linalg/
27 lines

Diff 501405

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

Show First 20 Lines • Show All 151 Lines • ▼ Show 20 Lines
getOrCreatePackedViewOfOperand(OpBuilder &b, Location loc, PackInfo packInfo, getOrCreatePackedViewOfOperand(OpBuilder &b, Location loc, PackInfo packInfo,
GenericOp genericOp, OpOperand *opOperand) { GenericOp genericOp, OpOperand *opOperand) {
int64_t numOrigLoops = genericOp.getNumLoops(); int64_t numOrigLoops = genericOp.getNumLoops();
int64_t numInnerLoops = packInfo.getNumTiledLoops(); int64_t numInnerLoops = packInfo.getNumTiledLoops();
int64_t numLoops = numOrigLoops + numInnerLoops; int64_t numLoops = numOrigLoops + numInnerLoops;
AffineMap origIndexingMap = genericOp.getMatchingIndexingMap(opOperand); AffineMap origIndexingMap = genericOp.getMatchingIndexingMap(opOperand);
llvm::DenseMap<int64_t, int64_t> domainDimToOperandDim; llvm::DenseMap<int64_t, int64_t> domainDimToOperandDim;
SmallVector<AffineExpr> exprs(origIndexingMap.getResults()); SmallVector<AffineExpr> exprs(origIndexingMap.getResults());
if (genericOp.isScalar(opOperand)) if (genericOp.isScalar(opOperand) || exprs.size() == 0)
return std::make_tuple(opOperand->get(), return std::make_tuple(opOperand->get(),
AffineMap::get(numLoops, 0, exprs, b.getContext())); AffineMap::get(numLoops, 0, exprs, b.getContext()));
// Step 1. Construct the information of packing data dimensions; append inner // Step 1. Construct the information of packing data dimensions; append inner
// dimensions to the indexing maps for the operand. // dimensions to the indexing maps for the operand.
for (auto [index, expr] : llvm::enumerate(exprs)) { for (auto [index, expr] : llvm::enumerate(exprs)) {
if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) { if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) {
int64_t dimPos = dimExpr.getPosition(); int64_t dimPos = dimExpr.getPosition();
▲ Show 20 Lines • Show All 417 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

Show First 20 Lines • Show All 646 Lines • ▼ Show 20 Lines
// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<128x256xi32> // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<128x256xi32>
// CHECK-NEXT: %[[GEN:.+]] = linalg.generic // CHECK-NEXT: %[[GEN:.+]] = linalg.generic
// CHECK-SAME: ins(%[[ARG0]] // CHECK-SAME: ins(%[[ARG0]]
// CHECK-SAME: outs(%[[EMPTY]] // CHECK-SAME: outs(%[[EMPTY]]
// CHECK: %[[ALLOC:.+]] = bufferization.alloc_tensor() : tensor<4x16x16x32xi32> // CHECK: %[[ALLOC:.+]] = bufferization.alloc_tensor() : tensor<4x16x16x32xi32>
// CHECK-NEXT: %{{.+}} = tensor.pack %[[GEN]] // CHECK-NEXT: %{{.+}} = tensor.pack %[[GEN]]
// CHECK-SAME: inner_dims_pos = [1, 0] inner_tiles = [16, 32] // CHECK-SAME: inner_dims_pos = [1, 0] inner_tiles = [16, 32]
// CHECK-SAME: into %[[ALLOC]] // CHECK-SAME: into %[[ALLOC]]
// -----
#map0 = affine_map<(d0, d1, d2, d3) -> ()>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
func.func @scalar_tensor(%arg0 : tensor<f32>) -> tensor<1x32x7x7x32xf32> {
%empty_gen = tensor.empty() : tensor<1x7x7x1024xf32>
%gen = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<f32>) outs(%empty_gen : tensor<1x7x7x1024xf32>) {
^bb0(%in: f32, %out: f32):
linalg.yield %in : f32
} -> tensor<1x7x7x1024xf32>
%empty_pack = tensor.empty() : tensor<1x32x7x7x32xf32>
%pack = tensor.pack %gen outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %empty_pack : tensor<1x7x7x1024xf32> -> tensor<1x32x7x7x32xf32>
return %pack : tensor<1x32x7x7x32xf32>
}
// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> ()>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
// CHECK: func.func @scalar_tensor
// CHECK-SAME: %[[ARG0:.+]]: tensor<f32>)
// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x32x7x7x32xf32>
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP1]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]]
// CHECK-SAME: outs(%[[EMPTY]]