diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp index 8a08a157b25d7..4dc7d6bc9c593 100644 --- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp +++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp @@ -5992,8 +5992,29 @@ OpFoldResult ShapeCastOp::fold(FoldAdaptor adaptor) { // shape_cast(constant) -> constant if (auto splatAttr = - llvm::dyn_cast_if_present(adaptor.getSource())) - return splatAttr.reshape(getType()); + llvm::dyn_cast_if_present(adaptor.getSource())) { + + // The shape and 'scalable dims' of the new attribute must match the result + // of the shape_cast: + auto newShape = resultType.getShape(); + auto newScalableDims = resultType.getScalableDims(); + + // The element type must be retained. Note that this is to handle currently + // valid IR like + // + // ``` + // %0 = llvm.mlir.constant(dense<0.> : vector<1xf8E4M3FN>) : vector<1xi8> + // %1 = vector.shape_cast %0 : vector<1xi8> to vector<1x1xi8> + // ``` + // + // where the element types of the attribute and result do not match. + auto newElementType = splatAttr.getElementType(); + + auto newAttr = VectorType::get(newShape, newElementType, newScalableDims); + + return DenseElementsAttr::get(newAttr, + splatAttr.getSplatValue()); + } // shape_cast(poison) -> poison if (llvm::dyn_cast_if_present(adaptor.getSource())) { diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir index 12187dd18012b..7d0f2c7814b74 100644 --- a/mlir/test/Dialect/Vector/canonicalize.mlir +++ b/mlir/test/Dialect/Vector/canonicalize.mlir @@ -1037,6 +1037,18 @@ func.func @fold_broadcast_shapecast(%arg0: vector<4xf32>) -> vector<4xf32> { // ----- +// CHECK-LABEL: func @canonicalize_extract_shapecast_different_element_type +// CHECK: %[[CONST:.*]] = llvm.mlir.constant +// CHECK-NEXT: return %[[CONST]] +func.func @canonicalize_extract_shapecast_different_element_type()->vector<12xi8> { + %0 = llvm.mlir.constant(dense<0.000000e+00> : vector<12xf8E4M3FN>) : vector<12xi8> + %1 = vector.shape_cast %0 : vector<12xi8> to vector<1x12xi8> + %2 = vector.extract %1[0] : vector<12xi8> from vector<1x12xi8> + return %2 : vector<12xi8> +} + +// ----- + // CHECK-LABEL: func @canonicalize_broadcast_shapecast_scalar // CHECK: vector.broadcast // CHECK-NOT: vector.shape_cast