(original) (raw)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp index d71a236f62f45..1853ae04f45d9 100644 --- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp +++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp @@ -2951,11 +2951,11 @@ void InsertOp::getCanonicalizationPatterns(RewritePatternSet &results, InsertOpConstantFolder>(context); } -// Eliminates insert operations that produce values identical to their source -// value. This happens when the source and destination vectors have identical -// sizes. OpFoldResult vector::InsertOp::fold(FoldAdaptor adaptor) { - if (getNumIndices() == 0) + // Fold "vector.insert %v, %dest [] : vector<2x2xf32> from vector<2x2xf32>" to + // %v. Note: Do not fold "vector.insert %v, %dest [] : f32 into vector" + // (type mismatch). + if (getNumIndices() == 0 && getSourceType() == getType()) return getSource(); return {}; } diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir index 6d6bc199e601c..c963460e7259f 100644 --- a/mlir/test/Dialect/Vector/canonicalize.mlir +++ b/mlir/test/Dialect/Vector/canonicalize.mlir @@ -800,6 +800,43 @@ func.func @fold_extract_shapecast_to_shapecast(%arg0 : vector<3x4xf32>) -> vecto // ----- +// CHECK-LABEL: func @extract_no_fold_scalar_to_0d( +// CHECK-SAME: %[[v:.*]]: vector) +// CHECK: %[[extract:.*]] = vector.extract %[[v]][] : f32 from vector+// CHECK: return %[[extract]] +func.func @extract_no_fold_scalar_to_0d(%v: vector) -> f32 { + %0 = vector.extract %v[] : f32 from vector+ return %0 : f32 +} + +// ----- + +// CHECK-LABEL: func @insert_fold_same_rank( +// CHECK-SAME: %[[v:.*]]: vector<2x2xf32>) +// CHECK: %[[CST:.+]] = arith.constant +// CHECK-SAME: : vector<2x2xf32> +// CHECK-NOT: vector.insert +// CHECK: return %[[CST]] +func.func @insert_fold_same_rank(%v: vector<2x2xf32>) -> vector<2x2xf32> { + %cst = arith.constant dense<0.000000e+00> : vector<2x2xf32> + %0 = vector.insert %cst, %v [] : vector<2x2xf32> into vector<2x2xf32> + return %0 : vector<2x2xf32> +} + +// ----- + +// CHECK-LABEL: func @insert_no_fold_scalar_to_0d( +// CHECK-SAME: %[[v:.*]]: vector) +// CHECK: %[[extract:.*]] = vector.insert %{{.*}}, %[[v]] [] : f32 into vector+// CHECK: return %[[extract]] +func.func @insert_no_fold_scalar_to_0d(%v: vector) -> vector { + %cst = arith.constant 0.000000e+00 : f32 + %0 = vector.insert %cst, %v [] : f32 into vector+ return %0 : vector+} + +// ----- + // CHECK-LABEL: dont_fold_expand_collapse // CHECK: %[[A:.*]] = vector.shape_cast %{{.*}} : vector<1x1x64xf32> to vector<1x1x8x8xf32> // CHECK: %[[B:.*]] = vector.shape_cast %{{.*}} : vector<1x1x8x8xf32> to vector<8x8xf32> @@ -2606,17 +2643,6 @@ func.func @rank_1_shuffle_to_interleave(%arg0: vector<6xi32>, %arg1: vector<6xi3 // ----- -// CHECK-LABEL: func @extract_from_0d_regression( -// CHECK-SAME: %[[v:.*]]: vector) -// CHECK: %[[extract:.*]] = vector.extract %[[v]][] : f32 from vector-// CHECK: return %[[extract]] -func.func @extract_from_0d_regression(%v: vector) -> f32 { - %0 = vector.extract %v[] : f32 from vector- return %0 : f32 -} - -// ----- - // CHECK-LABEL: func @extract_from_0d_splat_broadcast_regression( // CHECK-SAME: %[[a:.*]]: f32, %[[b:.*]]: vector, %[[c:.*]]: vector<2xf32>) func.func @extract_from_0d_splat_broadcast_regression(%a: f32, %b: vector, %c: vector<2xf32>) -> (f32, f32, f32, f32, f32, vector<6x7xf32>, vector<3xf32>) {