[Mlir-commits] [mlir] [mlir][vector] VectorLinearize: `ub.poison` support (PR #128612)

[Diego Caballero]

================
@@ -97,6 +98,35 @@ struct LinearizeConstant final : OpConversionPattern<arith::ConstantOp> {
   unsigned targetVectorBitWidth;
 };
 
+struct LinearizePoison final : OpConversionPattern<ub::PoisonOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LinearizePoison(
+      const TypeConverter &typeConverter, MLIRContext *context,
+      unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
+      PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit),
+        targetVectorBitWidth(targetVectBitWidth) {}
+  LogicalResult
+  matchAndRewrite(ub::PoisonOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    auto resType = getTypeConverter()->convertType<VectorType>(op.getType());
+
+    if (!resType)
+      return rewriter.notifyMatchFailure(loc, "can't convert return type");
+
+    if (!isLessThanTargetBitWidth(op, targetVectorBitWidth))
+      return rewriter.notifyMatchFailure(
+          loc, "Can't flatten since targetBitWidth <= OpSize");
+
+    rewriter.replaceOpWithNewOp<ub::PoisonOp>(op, resType);
+    return success();
+  }
+
+private:
+  unsigned targetVectorBitWidth;
+};
+
----------------
dcaballe wrote:

That's even better :)

https://github.com/llvm/llvm-project/pull/128612