[Mlir-commits] [mlir] [mlir][ub] Add `m_Poison()` matcher (PR #185022)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Mar 6 07:16:02 PST 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-arith
Author: Jakub Kuderski (kuhar)
<details>
<summary>Changes</summary>
Add a dedicated matcher for poison values in the UB dialect, similar to `m_Constant()` for general constants. The matcher uses `PoisonAttrInterface` for future extensibility.
Replace existing checks against `ub::PoisonAttr` with the new matcher.
Assisted-by: claude
---
Full diff: https://github.com/llvm/llvm-project/pull/185022.diff
8 Files Affected:
- (added) mlir/include/mlir/Dialect/UB/IR/UBMatchers.h (+52)
- (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+5-5)
- (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+3-3)
- (modified) mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp (+2-2)
- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+11-10)
- (modified) mlir/unittests/Dialect/CMakeLists.txt (+1)
- (added) mlir/unittests/Dialect/UB/CMakeLists.txt (+8)
- (added) mlir/unittests/Dialect/UB/UBMatchersTest.cpp (+75)
``````````diff
diff --git a/mlir/include/mlir/Dialect/UB/IR/UBMatchers.h b/mlir/include/mlir/Dialect/UB/IR/UBMatchers.h
new file mode 100644
index 0000000000000..9dc29ac37691c
--- /dev/null
+++ b/mlir/include/mlir/Dialect/UB/IR/UBMatchers.h
@@ -0,0 +1,52 @@
+//===- UBMatchers.h - UB Dialect matchers -----------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file provides matchers for the UB dialect, in particular for matching
+// poison values and attributes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_UB_IR_UBMATCHERS_H
+#define MLIR_DIALECT_UB_IR_UBMATCHERS_H
+
+#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/IR/Matchers.h"
+
+namespace mlir::ub {
+namespace detail {
+
+/// Matches a poison attribute (any attribute implementing PoisonAttrInterface).
+/// Supports matching against both Attribute and Operation* (via constant
+/// folding).
+struct poison_attr_matcher {
+ bool match(Attribute attr) { return isa<PoisonAttrInterface>(attr); }
+
+ bool match(Operation *op) {
+ Attribute attr;
+ if (!::mlir::detail::constant_op_binder<Attribute>(&attr).match(op))
+ return false;
+ return match(attr);
+ }
+};
+
+} // namespace detail
+
+/// Matches a poison constant (any attribute implementing PoisonAttrInterface).
+/// Works with `matchPattern` on Value, Operation*, and Attribute.
+///
+/// Examples:
+/// matchPattern(value, ub::m_Poison()) // Matches ub.poison op via Value.
+/// matchPattern(op, ub::m_Poison()) // Matches ub.poison op directly.
+/// matchPattern(attr, ub::m_Poison()) // Matches PoisonAttr(Interface).
+inline detail::poison_attr_matcher m_Poison() {
+ return detail::poison_attr_matcher();
+}
+
+} // namespace mlir::ub
+
+#endif // MLIR_DIALECT_UB_IR_UBMATCHERS_H
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index b99f77fdc8b30..73fa0df80bcba 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -13,7 +13,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/CommonFolders.h"
-#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/Dialect/UB/IR/UBMatchers.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributeInterfaces.h"
#include "mlir/IR/BuiltinAttributes.h"
@@ -460,7 +460,7 @@ arith::AddUIExtendedOp::fold(FoldAdaptor adaptor,
adaptor.getOperands(),
[](APInt a, const APInt &b) { return std::move(a) + b; })) {
// If any operand is poison, propagate poison to both results.
- if (isa<ub::PoisonAttr>(sumAttr)) {
+ if (matchPattern(sumAttr, ub::m_Poison())) {
results.push_back(sumAttr);
results.push_back(sumAttr);
return success();
@@ -1919,7 +1919,7 @@ OpFoldResult arith::BitcastOp::fold(FoldAdaptor adaptor) {
return {};
/// Bitcast poison.
- if (llvm::isa<ub::PoisonAttr>(operand))
+ if (matchPattern(operand, ub::m_Poison()))
return ub::PoisonAttr::get(getContext());
/// Bitcast integer or float to integer or float.
@@ -2503,10 +2503,10 @@ OpFoldResult arith::SelectOp::fold(FoldAdaptor adaptor) {
return falseVal;
// If either operand is fully poisoned, return the other.
- if (isa_and_nonnull<ub::PoisonAttr>(adaptor.getTrueValue()))
+ if (matchPattern(adaptor.getTrueValue(), ub::m_Poison()))
return falseVal;
- if (isa_and_nonnull<ub::PoisonAttr>(adaptor.getFalseValue()))
+ if (matchPattern(adaptor.getFalseValue(), ub::m_Poison()))
return trueVal;
// select %x, true, false => %x
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index ae65afac6b54c..d84408c024e25 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -27,7 +27,7 @@
#include "mlir/Dialect/Transform/IR/TransformTypes.h"
#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
#include "mlir/Dialect/Transform/Utils/Utils.h"
-#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/Dialect/UB/IR/UBMatchers.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
@@ -2199,7 +2199,7 @@ transform::PadOp::apply(transform::TransformRewriter &rewriter,
for (auto const &[untypedAttr, elementOrTensorType] :
llvm::zip(getPaddingValues(), linalgTarget->getOperandTypes())) {
- if (isa<ub::PoisonAttr>(untypedAttr)) {
+ if (matchPattern(untypedAttr, ub::m_Poison())) {
paddingValues.push_back(untypedAttr);
continue;
}
@@ -2452,7 +2452,7 @@ transform::PadTilingInterfaceOp::apply(transform::TransformRewriter &rewriter,
auto attr = dyn_cast<TypedAttr>(untypedAttr);
Type elementType = getElementTypeOrSelf(elementOrTensorType);
- if (isa<ub::PoisonAttr>(untypedAttr)) {
+ if (matchPattern(untypedAttr, ub::m_Poison())) {
paddingValues.push_back(untypedAttr);
continue;
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
index 52ab92f180575..ab2629b41a463 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
@@ -11,7 +11,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/Dialect/UB/IR/UBMatchers.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/BuiltinAttributes.h"
@@ -235,7 +235,7 @@ static Value padOperand(OpBuilder &builder, TilingInterface opToPad,
paddingValue = complex::ConstantOp::create(builder, opToPad.getLoc(),
complexTy, complexAttr);
}
- } else if (isa<ub::PoisonAttr>(paddingValueAttr)) {
+ } else if (matchPattern(paddingValueAttr, ub::m_Poison())) {
paddingValue = ub::PoisonOp::create(builder, opToPad.getLoc(),
getElementTypeOrSelf(v.getType()));
} else if (auto typedAttr = dyn_cast<TypedAttr>(paddingValueAttr)) {
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index a3077d7313b93..7afcd7d3f88fb 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -20,7 +20,7 @@
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/Dialect/UB/IR/UBMatchers.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Utils/VerificationUtils.h"
@@ -497,7 +497,7 @@ void VectorDialect::initialize() {
Operation *VectorDialect::materializeConstant(OpBuilder &builder,
Attribute value, Type type,
Location loc) {
- if (isa<ub::PoisonAttrInterface>(value))
+ if (matchPattern(value, ub::m_Poison()))
return value.getDialect().materializeConstant(builder, value, type, loc);
return arith::ConstantOp::materialize(builder, value, type, loc);
@@ -2112,7 +2112,7 @@ static Attribute foldPoisonIndexInsertExtractOp(MLIRContext *context,
/// Fold a vector extract from is a poison source.
static Attribute foldPoisonSrcExtractOp(Attribute srcAttr) {
- if (isa_and_nonnull<ub::PoisonAttr>(srcAttr))
+ if (matchPattern(srcAttr, ub::m_Poison()))
return srcAttr;
return {};
@@ -2719,7 +2719,7 @@ static OpFoldResult foldFromElementsToConstant(FromElementsOp fromElementsOp,
ArrayRef<Attribute> elements) {
// Check for null or poison attributes before any processing.
if (llvm::any_of(elements, [](Attribute attr) {
- return !attr || isa<ub::PoisonAttrInterface>(attr);
+ return !attr || matchPattern(attr, ub::m_Poison());
}))
return {};
@@ -3168,7 +3168,7 @@ OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) {
}
if (auto attr = llvm::dyn_cast<SplatElementsAttr>(adaptor.getSource()))
return DenseElementsAttr::get(vectorType, attr.getSplatValue<Attribute>());
- if (llvm::dyn_cast<ub::PoisonAttr>(adaptor.getSource()))
+ if (matchPattern(adaptor.getSource(), ub::m_Poison()))
return ub::PoisonAttr::get(getContext());
return {};
}
@@ -3323,8 +3323,8 @@ OpFoldResult vector::ShuffleOp::fold(FoldAdaptor adaptor) {
return {};
// Fold shuffle poison, poison -> poison.
- bool isV1Poison = isa<ub::PoisonAttr>(v1Attr);
- bool isV2Poison = isa<ub::PoisonAttr>(v2Attr);
+ bool isV1Poison = matchPattern(v1Attr, ub::m_Poison());
+ bool isV2Poison = matchPattern(v2Attr, ub::m_Poison());
if (isV1Poison && isV2Poison)
return ub::PoisonAttr::get(getContext());
@@ -4092,7 +4092,8 @@ class InsertStridedSliceConstantFolder final
return failure();
// TODO: Support poison.
- if (isa<ub::PoisonAttr>(vectorDestCst) || isa<ub::PoisonAttr>(sourceCst))
+ if (matchPattern(vectorDestCst, ub::m_Poison()) ||
+ matchPattern(sourceCst, ub::m_Poison()))
return failure();
// TODO: Handle non-unit strides when they become available.
@@ -6584,7 +6585,7 @@ OpFoldResult ShapeCastOp::fold(FoldAdaptor adaptor) {
return denseAttr.reshape(getType());
// shape_cast(poison) -> poison
- if (llvm::dyn_cast_if_present<ub::PoisonAttr>(adaptor.getSource()))
+ if (matchPattern(adaptor.getSource(), ub::m_Poison()))
return ub::PoisonAttr::get(getContext());
return {};
@@ -6940,7 +6941,7 @@ OpFoldResult vector::TransposeOp::fold(FoldAdaptor adaptor) {
return splat.reshape(getResultVectorType());
// Eliminate poison transpose ops.
- if (llvm::dyn_cast_if_present<ub::PoisonAttr>(adaptor.getVector()))
+ if (matchPattern(adaptor.getVector(), ub::m_Poison()))
return ub::PoisonAttr::get(getContext());
// Eliminate identity transposes, and more generally any transposes that
diff --git a/mlir/unittests/Dialect/CMakeLists.txt b/mlir/unittests/Dialect/CMakeLists.txt
index 269eccb1f93c3..6c05ecc9855e4 100644
--- a/mlir/unittests/Dialect/CMakeLists.txt
+++ b/mlir/unittests/Dialect/CMakeLists.txt
@@ -18,4 +18,5 @@ add_subdirectory(SparseTensor)
add_subdirectory(SPIRV)
add_subdirectory(SMT)
add_subdirectory(Transform)
+add_subdirectory(UB)
add_subdirectory(Utils)
diff --git a/mlir/unittests/Dialect/UB/CMakeLists.txt b/mlir/unittests/Dialect/UB/CMakeLists.txt
new file mode 100644
index 0000000000000..4da94e09e0c99
--- /dev/null
+++ b/mlir/unittests/Dialect/UB/CMakeLists.txt
@@ -0,0 +1,8 @@
+add_mlir_unittest(MLIRUBTests
+ UBMatchersTest.cpp
+)
+mlir_target_link_libraries(MLIRUBTests
+ PRIVATE
+ MLIRArithDialect
+ MLIRUBDialect
+)
diff --git a/mlir/unittests/Dialect/UB/UBMatchersTest.cpp b/mlir/unittests/Dialect/UB/UBMatchersTest.cpp
new file mode 100644
index 0000000000000..92564d8f77dc4
--- /dev/null
+++ b/mlir/unittests/Dialect/UB/UBMatchersTest.cpp
@@ -0,0 +1,75 @@
+//===- UBMatchersTest.cpp - Unit tests for UB matchers --------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/UB/IR/UBMatchers.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+namespace {
+
+class UBMatchersTest : public testing::Test {
+public:
+ UBMatchersTest() { ctx.getOrLoadDialect<ub::UBDialect>(); }
+
+protected:
+ MLIRContext ctx;
+ OpBuilder b{&ctx};
+};
+
+TEST_F(UBMatchersTest, MatchPoisonOp) {
+ auto loc = UnknownLoc::get(&ctx);
+ OwningOpRef<ub::PoisonOp> poisonOp =
+ ub::PoisonOp::create(b, loc, b.getF32Type());
+ Value poisonVal = poisonOp->getResult();
+
+ // Match via Operation*.
+ EXPECT_TRUE(matchPattern(poisonOp->getOperation(), ub::m_Poison()));
+ // Match via Value.
+ EXPECT_TRUE(matchPattern(poisonVal, ub::m_Poison()));
+}
+
+TEST_F(UBMatchersTest, MatchPoisonAttr) {
+ auto poisonAttr = ub::PoisonAttr::get(&ctx);
+ // Match via Attribute.
+ EXPECT_TRUE(matchPattern(poisonAttr, ub::m_Poison()));
+}
+
+TEST_F(UBMatchersTest, MatchVectorPoisonOp) {
+ auto loc = UnknownLoc::get(&ctx);
+ auto vectorType = VectorType::get({4}, b.getI32Type());
+ OwningOpRef<ub::PoisonOp> poisonOp = ub::PoisonOp::create(b, loc, vectorType);
+
+ EXPECT_TRUE(matchPattern(poisonOp->getOperation(), ub::m_Poison()));
+}
+
+TEST_F(UBMatchersTest, NonPoisonAttrDoesNotMatch) {
+ IntegerAttr intAttr = b.getI32IntegerAttr(42);
+ EXPECT_FALSE(matchPattern(intAttr, ub::m_Poison()));
+
+ FloatAttr floatAttr = b.getF32FloatAttr(1.0f);
+ EXPECT_FALSE(matchPattern(floatAttr, ub::m_Poison()));
+}
+
+TEST_F(UBMatchersTest, NonPoisonOpDoesNotMatch) {
+ ctx.getOrLoadDialect<arith::ArithDialect>();
+ auto loc = UnknownLoc::get(&ctx);
+ OwningOpRef<arith::ConstantOp> constOp =
+ arith::ConstantOp::create(b, loc, b.getI32IntegerAttr(42));
+ EXPECT_FALSE(matchPattern(constOp->getOperation(), ub::m_Poison()));
+}
+
+TEST_F(UBMatchersTest, NullAttrDoesNotMatch) {
+ Attribute nullAttr;
+ EXPECT_FALSE(matchPattern(nullAttr, ub::m_Poison()));
+}
+
+} // namespace
``````````
</details>
https://github.com/llvm/llvm-project/pull/185022
More information about the Mlir-commits
mailing list