This is an archive of the discontinued LLVM Phabricator instance.

Differential D101942

[mlir][Linalg] Fix element type of results when folding reshapes.
ClosedPublic

Authored by mravishankar on May 5 2021, 1:54 PM.

Details

Reviewers
nicolasvasilache
ThomasRaoux
Commits
rGb6060b76731d: [mlir][Linalg] Fix element type of results when folding reshapes.
Summary

Fixing a minor bug which lead to element type of the output being
modified when folding reshapes with generic op.

Diff Detail

Event Timeline

mravishankar created this revision.May 5 2021, 1:54 PM
Herald added subscribers: dcaballe, cota, teijeong and 16 others. · View Herald TranscriptMay 5 2021, 1:54 PM
mravishankar requested review of this revision.May 5 2021, 1:54 PM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptMay 5 2021, 1:54 PM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: limo1996, stephenneuendorffer, nicolasvasilache. · View Herald Transcript
mravishankar added a reviewer: ThomasRaoux.May 5 2021, 1:55 PM
ThomasRaoux accepted this revision.May 5 2021, 1:56 PM

Thanks for catching that.

This revision is now accepted and ready to land.May 5 2021, 1:56 PM
Harbormaster completed remote builds in B102835: Diff 343179.May 5 2021, 2:30 PM
Closed by commit rGb6060b76731d: [mlir][Linalg] Fix element type of results when folding reshapes. (authored by mravishankar). · Explain WhyMay 5 2021, 3:41 PM
This revision was automatically updated to reflect the committed changes.
mravishankar added a commit: rGb6060b76731d: [mlir][Linalg] Fix element type of results when folding reshapes..

Revision Contents

PathSize
mlir/
lib/
Dialect/
Linalg/
Transforms/
7 lines
test/
Dialect/
Linalg/
37 lines

Diff 343215

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

Show First 20 Lines • Show All 1,123 Lines • ▼ Show 20 Lines for (auto map : op.getIndexingMaps()) {
} }
newMaps.push_back(AffineMap::get(newRank, 0, newExprs, op.getContext())); newMaps.push_back(AffineMap::get(newRank, 0, newExprs, op.getContext()));
} }
// 4. Reshape the output tensors. // 4. Reshape the output tensors.
SmallVector<Value> newOutputs; SmallVector<Value> newOutputs;
SmallVector<Type> newOutputTypes; SmallVector<Type> newOutputTypes;
for (auto output : op.outputs()) { for (auto output : op.outputs()) {
auto newOutputType = RankedTensorType::get(
reshapeFound.getSrcType().getShape(),
output.getType().template cast<RankedTensorType>().getElementType());
Value newOutput = rewriter.create<TensorReshapeOp>( Value newOutput = rewriter.create<TensorReshapeOp>(
op->getLoc(), reshapeFound.getSrcType(), output, reassociation); op->getLoc(), newOutputType, output, reassociation);
newOutputTypes.push_back(newOutput.getType()); newOutputTypes.push_back(newOutputType);
newOutputs.push_back(newOutput); newOutputs.push_back(newOutput);
} }
// 5. Create a new generic op with lowerer rank. // 5. Create a new generic op with lowerer rank.
SmallVector<StringRef, 4> iteratorTypes(newRank, SmallVector<StringRef, 4> iteratorTypes(newRank,
getParallelIteratorTypeName()); getParallelIteratorTypeName());
auto newOp = auto newOp =
rewriter.create<GenericOpTy>(op->getLoc(), newOutputTypes, newOperands, rewriter.create<GenericOpTy>(op->getLoc(), newOutputTypes, newOperands,
newOutputs, newMaps, iteratorTypes); newOutputs, newMaps, iteratorTypes);
▲ Show 20 Lines • Show All 363 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/fusion-push-reshape.mlir

Show First 20 Lines • Show All 82 Lines • ▼ Show 20 Linesfunc @reshape_negative(%A: tensor<12544x16xf32>, %B: tensor<112xf32>) -> tensor<112x112x16xf32> {
ins(%20, %B : tensor<112x112x16xf32>, tensor<112xf32>) ins(%20, %B : tensor<112x112x16xf32>, tensor<112xf32>)
outs(%21 : tensor<112x112x16xf32>) { outs(%21 : tensor<112x112x16xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): // no predecessors ^bb0(%arg1: f32, %arg2: f32, %arg3: f32): // no predecessors
%s = subf %arg1, %arg2 : f32 %s = subf %arg1, %arg2 : f32
linalg.yield %s : f32 linalg.yield %s : f32
} -> tensor<112x112x16xf32> } -> tensor<112x112x16xf32>
return %22 : tensor<112x112x16xf32> return %22 : tensor<112x112x16xf32>
} }
// -----
func @type_correctness(%arg0 : tensor<6x5xi32>, %arg1 : tensor<5xf32>,
%arg2 : tensor<5xf32>) -> tensor<2x3x5xf32> {
%cst_6 = constant 1.000000e+00 : f32
%cst_7 = constant 7.000000e+00 : f32
%cst_8 = constant 1.1920929E-7 : f32
%25 = linalg.tensor_reshape %arg0
[affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d2)>]
: tensor<6x5xi32> into tensor<2x3x5xi32>
%26 = linalg.init_tensor [2, 3, 5] : tensor<2x3x5xf32>
%28 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
affine_map<(d0, d1, d2) -> (d2)>,
affine_map<(d0, d1, d2) -> (d2)>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]}
ins(%25, %arg1, %arg2 : tensor<2x3x5xi32>, tensor<5xf32>, tensor<5xf32>)
outs(%26 : tensor<2x3x5xf32>) {
^bb0(%arg6: i32, %arg7: f32, %arg8: f32, %arg9: f32): // no predecessors
%29 = sitofp %arg6 : i32 to f32
%30 = addf %arg7, %cst_8 : f32
%31 = divf %cst_7, %30 : f32
%32 = divf %cst_6, %31 : f32
%33 = mulf %29, %32 : f32
%34 = addf %33, %arg8 : f32
linalg.yield %34 : f32
} -> tensor<2x3x5xf32>
return %28 : tensor<2x3x5xf32>
}
// CHECK-LABEL: func @type_correctness
// CHECK: %[[OP:.+]] = linalg.generic
// CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}} : tensor<6x5xi32>, tensor<5xf32>, tensor<5xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<6x5xf32>)
// CHECK: linalg.tensor_reshape %[[OP]]
// CHECK-SAME: tensor<6x5xf32> into tensor<2x3x5xf32>