[mlir] Guard sccp pass from crashing with different source type

Vector::BroadCastOp expects the idential element type in folding. It
causes the crash if the different source type is given to the SCCP pass.
We need to guard the pass from crashing if the non-idential element type
is given, but still compatible. (e.g. index vs integer type)

Author: Lewuathe

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -2523,8 +2523,16 @@ OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) {
   if (!adaptor.getSource())
     return {};
   auto vectorType = getResultVectorType();
-  if (llvm::isa<IntegerAttr, FloatAttr>(adaptor.getSource()))
-    return DenseElementsAttr::get(vectorType, adaptor.getSource());
+  if (auto attr = llvm::dyn_cast<IntegerAttr>(adaptor.getSource())) {
+    if (vectorType.getElementType() != attr.getType())
+      return {};
+    return DenseElementsAttr::get(vectorType, attr);
+  }
+  if (auto attr = llvm::dyn_cast<FloatAttr>(adaptor.getSource())) {
+    if (vectorType.getElementType() != attr.getType())
+      return {};
+    return DenseElementsAttr::get(vectorType, attr);
+  }
   if (auto attr = llvm::dyn_cast<SplatElementsAttr>(adaptor.getSource()))
     return DenseElementsAttr::get(vectorType, attr.getSplatValue<Attribute>());
   return {};

mlir/test/Transforms/sccp.mlir

@@ -246,3 +246,16 @@ func.func @op_with_region() -> (i32) {
 ^b:
   return %1 : i32
 }
+
+// CHECK-LABEL: no_crash_with_different_source_type
+func.func @no_crash_with_different_source_type() {
+  // CHECK: llvm.mlir.constant(0 : index) : i64
+  %0 = llvm.mlir.constant(0 : index) : i64
+  llvm.br ^b1(%0 : i64)
+^b1(%1: i64):  
+  llvm.br ^b2
+^b2: 
+  // CHECK: vector.broadcast %[[CST:.*]] : i64 to vector<128xi64>
+  %2 = vector.broadcast %1 : i64 to vector<128xi64>
+  llvm.return
+}