diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp index 0c0a7bc98d8b5..ad709813c6216 100644 --- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp +++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp @@ -1756,11 +1756,6 @@ static Value foldExtractFromShapeCast(ExtractOp extractOp) { if (!shapeCastOp) return Value(); - // 0-D vectors not supported. - assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported"); - if (hasZeroDimVectors(shapeCastOp)) - return Value(); - // Get the nth dimension size starting from lowest dimension. auto getDimReverse = [](VectorType type, int64_t n) { return type.getShape().take_back(n + 1).front(); diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir index 5ae769090dac6..89af0f7332f5c 100644 --- a/mlir/test/Dialect/Vector/canonicalize.mlir +++ b/mlir/test/Dialect/Vector/canonicalize.mlir @@ -782,23 +782,23 @@ func.func @fold_extract_shapecast(%arg0 : vector<5x1x3x2xf32>, // ----- -// CHECK-LABEL: fold_extract_shapecast_negative -// CHECK: %[[V:.*]] = vector.shape_cast %{{.*}} : vector<16xf32> to vector<2x4x2xf32> -// CHECK: %[[R:.*]] = vector.extract %[[V]][1] : vector<4x2xf32> from vector<2x4x2xf32> -// CHECK: return %[[R]] : vector<4x2xf32> -func.func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>) -> vector<4x2xf32> { - %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<2x4x2xf32> - %r = vector.extract %0[1] : vector<4x2xf32> from vector<2x4x2xf32> - return %r : vector<4x2xf32> +// CHECK-LABEL: fold_extract_shapecast_0d_result +// CHECK-SAME: %[[IN:.*]]: vector<1x1x1xf32> +// CHECK: %[[R:.*]] = vector.extract %[[IN]][0, 0, 0] : f32 from vector<1x1x1xf32> +// CHECK: return %[[R]] : f32 +func.func @fold_extract_shapecast_0d_result(%arg0 : vector<1x1x1xf32>) -> f32 { + %0 = vector.shape_cast %arg0 : vector<1x1x1xf32> to vector + %r = vector.extract %0[] : f32 from vector + return %r : f32 } // ----- -// CHECK-LABEL: dont_fold_0d_extract_shapecast -// CHECK: %[[V:.*]] = vector.shape_cast %{{.*}} : vector to vector<1xf32> -// CHECK: %[[R:.*]] = vector.extract %[[V]][0] : f32 from vector<1xf32> +// CHECK-LABEL: fold_extract_shapecast_0d_source +// CHECK-SAME: %[[IN:.*]]: vector +// CHECK: %[[R:.*]] = vector.extract %[[IN]][] : f32 from vector // CHECK: return %[[R]] : f32 -func.func @dont_fold_0d_extract_shapecast(%arg0 : vector) -> f32 { +func.func @fold_extract_shapecast_0d_source(%arg0 : vector) -> f32 { %0 = vector.shape_cast %arg0 : vector to vector<1xf32> %r = vector.extract %0[0] : f32 from vector<1xf32> return %r : f32 @@ -806,6 +806,18 @@ func.func @dont_fold_0d_extract_shapecast(%arg0 : vector) -> f32 { // ----- +// CHECK-LABEL: fold_extract_shapecast_negative +// CHECK: %[[V:.*]] = vector.shape_cast %{{.*}} : vector<16xf32> to vector<2x4x2xf32> +// CHECK: %[[R:.*]] = vector.extract %[[V]][1] : vector<4x2xf32> from vector<2x4x2xf32> +// CHECK: return %[[R]] : vector<4x2xf32> +func.func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>) -> vector<4x2xf32> { + %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<2x4x2xf32> + %r = vector.extract %0[1] : vector<4x2xf32> from vector<2x4x2xf32> + return %r : vector<4x2xf32> +} + +// ----- + // CHECK-LABEL: fold_extract_shapecast_to_shapecast // CHECK-SAME: (%[[ARG:.+]]: vector<3x4xf32>) // CHECK: %[[R:.+]] = vector.shape_cast %[[ARG]] : vector<3x4xf32> to vector<12xf32>