Skip to content

[mlir][tensor] Fold when source is const #71643

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Nov 9, 2023
Merged

[mlir][tensor] Fold when source is const #71643

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4c021b0
[mlir][tensor] Fold when source is const
rikhuijzer Nov 8, 2023
6b819fb
Apply `clang-format`
rikhuijzer Nov 8, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -659,6 +659,7 @@ def Tensor_GatherOp : Tensor_Op<"gather", [
}
}];
let hasVerifier = 1;
let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -986,6 +987,7 @@ def Tensor_ReshapeOp: Tensor_Op<"reshape", [
$source `(` $shape `)` attr-dict `:` functional-type(operands, results)
}];
let hasVerifier = 1;
let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -1867,6 +1869,8 @@ def Tensor_PackOp : Tensor_RelayoutOp<"pack", [
}];

let hasCanonicalizeMethod = 1;

let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -1948,6 +1952,8 @@ def Tensor_UnPackOp : Tensor_RelayoutOp<"unpack"> {
}];

let hasCanonicalizeMethod = 1;

let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
Expand Down
54 changes: 48 additions & 6 deletions mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -834,6 +834,17 @@ void EmptyOp::getCanonicalizationPatterns(RewritePatternSet &results,
ReplaceEmptyTensorStaticShapeDims>(context);
}

/// Try to remove a tensor operation if it would only reshape a constant.
/// Removes the op and replaces the constant with a new constant of the result
/// shape.
static OpFoldResult reshapeConstantSource(DenseElementsAttr source,
TensorType result) {
if (source && source.isSplat() && result.hasStaticShape())
return source.resizeSplat(result);

return {};
}

//===----------------------------------------------------------------------===//
// ExtractOp
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -1089,6 +1100,14 @@ LogicalResult GatherOp::verify() {
return success();
}

OpFoldResult GatherOp::fold(FoldAdaptor adaptor) {
if (OpFoldResult reshapedSource = reshapeConstantSource(
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getSource()),
getResult().getType()))
return reshapedSource;
return {};
}

//===----------------------------------------------------------------------===//
// InsertOp
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -1367,6 +1386,14 @@ LogicalResult ReshapeOp::verify() {
return success();
}

OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
if (OpFoldResult reshapedSource = reshapeConstantSource(
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getSource()),
getResult().getType()))
return reshapedSource;
return {};
}

//===----------------------------------------------------------------------===//
// Reassociative reshape ops
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -2153,12 +2180,10 @@ static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) {
}

OpFoldResult ExtractSliceOp::fold(FoldAdaptor adaptor) {
if (auto splat =
llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getSource())) {
auto resultType = llvm::cast<ShapedType>(getResult().getType());
if (resultType.hasStaticShape())
return splat.resizeSplat(resultType);
}
if (OpFoldResult reshapedSource = reshapeConstantSource(
llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getSource()),
getResult().getType()))
return reshapedSource;
if (getSourceType() == getType() &&
succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType())))
return this->getSource();
Expand Down Expand Up @@ -3823,6 +3848,14 @@ bool PackOp::isLikePad() {
return isLikePadUnPad(*this, packedTensorType);
}

OpFoldResult PackOp::fold(FoldAdaptor adaptor) {
if (OpFoldResult reshapedSource = reshapeConstantSource(
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getSource()),
getResult().getType()))
return reshapedSource;
return {};
}

//===----------------------------------------------------------------------===//
// UnPackOp
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -3951,6 +3984,15 @@ bool UnPackOp::isLikeUnPad() {
RankedTensorType packedTensorType = getSourceType();
return isLikePadUnPad(*this, packedTensorType);
}

OpFoldResult UnPackOp::fold(FoldAdaptor adaptor) {
if (OpFoldResult reshapedSource = reshapeConstantSource(
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getSource()),
getResult().getType()))
return reshapedSource;
return {};
}

//===----------------------------------------------------------------------===//
// Common Canonicalizers and Folders.
//===----------------------------------------------------------------------===//
Expand Down
48 changes: 48 additions & 0 deletions mlir/test/Dialect/Tensor/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -672,6 +672,30 @@ func.func @fold_extract_insert(%input : tensor<?x?x?xf32>, %slice: tensor<4x?x8x

// -----

// CHECK-LABEL: func @fold_gather_constant_splat
// CHECK-NOT: tensor.gather
// CHECK: arith.constant dense<1.000000e-01> : tensor<1x2x1x1x1xf32>
func.func @fold_gather_constant_splat(%indices : tensor<1x2x3xindex>) -> tensor<1x2x1x1x1xf32> {
%cst = arith.constant dense<1.000000e-01> : tensor<4x4x4xf32>
%0 = tensor.gather %cst[%indices] gather_dims([0, 1, 2]) :
(tensor<4x4x4xf32>, tensor<1x2x 3xindex>) -> tensor<1x2x 1x1x1xf32>
return %0 : tensor<1x2x 1x1x1xf32>
}

// -----

// CHECK-LABEL: func @fold_reshape_constant_splat
// CHECK-NOT: tensor.reshape
// CHECK: arith.constant dense<1.000000e-01> : tensor<4xf32>
func.func @fold_reshape_constant_splat(%shape : tensor<1xi32>) -> tensor<4xf32> {
%cst = arith.constant dense<1.000000e-01> : tensor<4x1xf32>
%0 = tensor.reshape %cst(%shape)
: (tensor<4x1xf32>, tensor<1xi32>) -> tensor<4xf32>
return %0 : tensor<4xf32>
}

// -----

// CHECK-LABEL: func @fold_extract_constant_splat
// CHECK-NOT: tensor.extract_slice
// CHECK: arith.constant dense<42> : tensor<4x4xi32>
Expand All @@ -683,6 +707,30 @@ func.func @fold_extract_constant_splat() -> (tensor<4x4xi32>) {

// -----

// CHECK-LABEL: func @fold_pack_constant_splat
// CHECK-NOT: tensor.pack
// CHECK: arith.constant dense<1.000000e-01> : tensor<8x16x8x32xf32>
func.func @fold_pack_constant_splat(%dest : tensor<8x16x8x32xf32>) -> tensor<8x16x8x32xf32> {
%cst = arith.constant dense<1.000000e-01> : tensor<64x128xf32>
%0 = tensor.pack %cst outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
inner_tiles = [8, 32] into %dest : tensor<64x128xf32> -> tensor<8x16x8x32xf32>
return %0 : tensor<8x16x8x32xf32>
}

// -----

// CHECK-LABEL: func @fold_unpack_constant_splat
// CHECK-NOT: tensor.unpack
// CHECK: arith.constant dense<1.000000e-01> : tensor<128x256xf32>
func.func @fold_unpack_constant_splat(%dest : tensor<128x256xf32>) -> tensor<128x256xf32> {
%cst = arith.constant dense<1.000000e-01> : tensor<16x8x8x32xf32>
%0 = tensor.unpack %cst inner_dims_pos = [0, 1]
inner_tiles = [8, 32] into %dest : tensor<16x8x8x32xf32> -> tensor<128x256xf32>
return %0 : tensor<128x256xf32>
}

// -----

// CHECK-LABEL: func @fold_overlapping_insert
// CHECK-SAME: %[[INPUT:.+]]: tensor<?x?x?xf32>, %{{.+}}: tensor<4x?x8xf32>, %[[SLICE2:.+]]: tensor<4x?x8xf32>
func.func @fold_overlapping_insert(%input : tensor<?x?x?xf32>, %slice1: tensor<4x?x8xf32>, %slice2: tensor<4x?x8xf32>, %i: index, %size: index) -> (tensor<?x?x?xf32>) {
Expand Down