Feedback

Author: dcaballe

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

@@ -437,8 +437,8 @@ void VectorDialect::initialize() {
 Operation *VectorDialect::materializeConstant(OpBuilder &builder,
                                               Attribute value, Type type,
                                               Location loc) {
-  if (auto poisonAttr = dyn_cast<ub::PoisonAttrInterface>(value))
-    return builder.create<ub::PoisonOp>(loc, type, poisonAttr);
+  if (isa<ub::PoisonAttrInterface>(value))
+    return value.getDialect().materializeConstant(builder, value, type, loc);
 
   return arith::ConstantOp::materialize(builder, value, type, loc);
 }

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

@@ -1250,11 +1250,11 @@ func.func @extract_scalar_from_vec_1d_f32(%arg0: vector<16xf32>) -> f32 {
 
 // -----
 
-func.func @extract_poison_idx(%arg0: vector<16xf32>) -> f32 {
+func.func @extract_scalar_from_vec_1d_f32_poison_idx(%arg0: vector<16xf32>) -> f32 {
   %0 = vector.extract %arg0[-1]: f32 from vector<16xf32>
   return %0 : f32
 }
-// CHECK-LABEL: @extract_poison_idx
+// CHECK-LABEL: @extract_scalar_from_vec_1d_f32_poison_idx
 //       CHECK:   %[[UB:.*]] = ub.poison : f32
 //       CHECK:   return %[[UB]] : f32
 
@@ -1335,6 +1335,16 @@ func.func @extract_vec_2d_from_vec_3d_f32(%arg0: vector<4x3x16xf32>) -> vector<3
 
 // -----
 
+func.func @extract_vec_2d_from_vec_3d_f32_poison_idx(%arg0: vector<4x3x16xf32>) -> vector<3x16xf32> {
+  %0 = vector.extract %arg0[-1]: vector<3x16xf32> from vector<4x3x16xf32>
+  return %0 : vector<3x16xf32>
+}
+// CHECK-LABEL: @extract_vec_2d_from_vec_3d_f32_poison_idx
+//       CHECK:   %[[UB:.*]] = ub.poison : vector<3x16xf32>
+//       CHECK:   return %[[UB]] : vector<3x16xf32>
+
+// -----
+
 func.func @extract_vec_2d_from_vec_3d_f32_scalable(%arg0: vector<4x3x[16]xf32>) -> vector<3x[16]xf32> {
   %0 = vector.extract %arg0[0]: vector<3x[16]xf32> from vector<4x3x[16]xf32>
   return %0 : vector<3x[16]xf32>