[mlir] Support attributes in `matchPattern`

[kuhar]

Summary

The primary motivation is to we have a simple mechanism to extract
values from attributes in folders and canon patterns without having to
re-fold constants or write nested conditions over attribute types.

Matching over attributes composes especially well with fold adaptors.

Update folds in Arith and SPIRV dialects to match over attributes, where
applicable.

[mehdi_amini]

FYI: we're using GitHub pull-requests now.

mlir/include/mlir/IR/Matchers.h

 // tree-based pattern matching over MLIR. This mechanism is inspired by LLVM's
 // include/llvm/IR/PatternMatch.h.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_MATCHERS_H
 #define MLIR_IR_MATCHERS_H
 
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 
 namespace mlir {
 
 namespace detail {
 
 /// The matcher that matches a certain kind of Attribute and binds the value
 /// inside the Attribute.
 template <
     typename AttrClass,
     // Require AttrClass to be a derived class from Attribute and get its
     // value type
     typename ValueType = typename std::enable_if_t<
         std::is_base_of<Attribute, AttrClass>::value, AttrClass>::ValueType,
     // Require the ValueType is not void
     typename = std::enable_if_t<!std::is_void<ValueType>::value>>
 struct attr_value_binder {
   ValueType *bind_value;
 
   /// Creates a matcher instance that binds the value to bv if match succeeds.
   attr_value_binder(ValueType *bv) : bind_value(bv) {}
 
-  bool match(const Attribute &attr) {
+  bool match(Attribute attr) {
     if (auto intAttr = llvm::dyn_cast<AttrClass>(attr)) {
       *bind_value = intAttr.getValue();
       return true;
     }
     return false;
   }
 };
 
bool match(Operation *op) {
     }
     return false;
   }
   StringRef attrName;
   AttrT *bindValue;
 };
 
 /// The matcher that matches a constant scalar / vector splat / tensor splat
-/// float operation and binds the constant float value.
-struct constant_float_op_binder {
+/// float Attribute or Operation and binds the constant float value.
+struct constant_float_value_binder {
   FloatAttr::ValueType *bind_value;
 
   /// Creates a matcher instance that binds the value to bv if match succeeds.
-  constant_float_op_binder(FloatAttr::ValueType *bv) : bind_value(bv) {}
+  constant_float_value_binder(FloatAttr::ValueType *bv) : bind_value(bv) {}
+
+  bool match(Attribute attr) {
+    attr_value_binder<FloatAttr> matcher(bind_value);
+    if (matcher.match(attr))
+      return true;
+
+    if (auto splatAttr = dyn_cast<SplatElementsAttr>(attr))
+      return matcher.match(splatAttr.getSplatValue<Attribute>());
+
+    return false;
+  }
 
   bool match(Operation *op) {
     Attribute attr;
     if (!constant_op_binder<Attribute>(&attr).match(op))
       return false;
-    auto type = op->getResult(0).getType();
 
-    if (llvm::isa<FloatType>(type))
-      return attr_value_binder<FloatAttr>(bind_value).match(attr);
-    if (llvm::isa<VectorType, RankedTensorType>(type)) {
-      if (auto splatAttr = llvm::dyn_cast<SplatElementsAttr>(attr)) {
-        return attr_value_binder<FloatAttr>(bind_value)
-            .match(splatAttr.getSplatValue<Attribute>());
-      }
-    }
+    Type type = op->getResult(0).getType();
+    if (isa<FloatType, VectorType, RankedTensorType>(type))
+      return match(attr);
+
     return false;
   }
 };
 
 /// The matcher that matches a given target constant scalar / vector splat /
 /// tensor splat float value that fulfills a predicate.
 struct constant_float_predicate_matcher {
   bool (*predicate)(const APFloat &);
 
+  bool match(Attribute attr) {
+    APFloat value(APFloat::Bogus());
+    return constant_float_value_binder(&value).match(attr) && predicate(value);
+  }
+
   bool match(Operation *op) {
     APFloat value(APFloat::Bogus());
-    return constant_float_op_binder(&value).match(op) && predicate(value);
+    return constant_float_value_binder(&value).match(op) && predicate(value);
   }
 };
 
 /// The matcher that matches a constant scalar / vector splat / tensor splat
-/// integer operation and binds the constant integer value.
-struct constant_int_op_binder {
+/// integer Attribute or Operation and binds the constant integer value.
+struct constant_int_value_binder {
   IntegerAttr::ValueType *bind_value;
 
   /// Creates a matcher instance that binds the value to bv if match succeeds.
-  constant_int_op_binder(IntegerAttr::ValueType *bv) : bind_value(bv) {}
+  constant_int_value_binder(IntegerAttr::ValueType *bv) : bind_value(bv) {}
+
+  bool match(Attribute attr) {
+    attr_value_binder<IntegerAttr> matcher(bind_value);
+    if (matcher.match(attr))
+      return true;
+
+    if (auto splatAttr = dyn_cast<SplatElementsAttr>(attr))
+      return matcher.match(splatAttr.getSplatValue<Attribute>());
+
+    return false;
+  }
 
   bool match(Operation *op) {
     Attribute attr;
     if (!constant_op_binder<Attribute>(&attr).match(op))
       return false;
-    auto type = op->getResult(0).getType();
 
-    if (llvm::isa<IntegerType, IndexType>(type))
-      return attr_value_binder<IntegerAttr>(bind_value).match(attr);
-    if (llvm::isa<VectorType, RankedTensorType>(type)) {
-      if (auto splatAttr = llvm::dyn_cast<SplatElementsAttr>(attr)) {
-        return attr_value_binder<IntegerAttr>(bind_value)
-            .match(splatAttr.getSplatValue<Attribute>());
-      }
-    }
+    Type type = op->getResult(0).getType();
+    if (isa<IntegerType, IndexType, VectorType, RankedTensorType>(type))
+      return match(attr);
+
     return false;
   }
 };
 
 /// The matcher that matches a given target constant scalar / vector splat /
 /// tensor splat integer value that fulfills a predicate.
 struct constant_int_predicate_matcher {
   bool (*predicate)(const APInt &);
 
+  bool match(Attribute attr) {
+    APInt value;
+    return constant_int_value_binder(&value).match(attr) && predicate(value);
+  }
+
   bool match(Operation *op) {
     APInt value;
-    return constant_int_op_binder(&value).match(op) && predicate(value);
+    return constant_int_value_binder(&value).match(op) && predicate(value);
   }
 };
 
 /// The matcher that matches a certain kind of op.
 template <typename OpClass>
 struct op_matcher {
   bool match(Operation *op) { return isa<OpClass>(op); }
 };
 
 /// Trait to check whether T provides a 'match' method with type
-/// `OperationOrValue`.
-template <typename T, typename OperationOrValue>
-using has_operation_or_value_matcher_t =
-    decltype(std::declval<T>().match(std::declval<OperationOrValue>()));
+/// `MatchTarget` (Value, Operation, or Attribute).
+template <typename T, typename MatchTarget>
+using has_compatible_matcher_t =
+    decltype(std::declval<T>().match(std::declval<MatchTarget>()));
 
 /// Statically switch to a Value matcher.
 template <typename MatcherClass>
-std::enable_if_t<llvm::is_detected<detail::has_operation_or_value_matcher_t,
+std::enable_if_t<llvm::is_detected<detail::has_compatible_matcher_t,
                                    MatcherClass, Value>::value,
                  bool>
 matchOperandOrValueAtIndex(Operation *op, unsigned idx, MatcherClass &matcher) {
   return matcher.match(op->getOperand(idx));
 }
 
 /// Statically switch to an Operation matcher.
 template <typename MatcherClass>
-std::enable_if_t<llvm::is_detected<detail::has_operation_or_value_matcher_t,
+std::enable_if_t<llvm::is_detected<detail::has_compatible_matcher_t,
                                    MatcherClass, Operation *>::value,
                  bool>
 matchOperandOrValueAtIndex(Operation *op, unsigned idx, MatcherClass &matcher) {
   if (auto *defOp = op->getOperand(idx).getDefiningOp())
     return matcher.match(defOp);
   return false;
 }
 

 template <typename OpClass>
 inline detail::op_matcher<OpClass> m_Op() {
   return detail::op_matcher<OpClass>();
 }
 
 /// Entry point for matching a pattern over a Value.
 template <typename Pattern>
 inline bool matchPattern(Value value, const Pattern &pattern) {
+  assert(value);
   // TODO: handle other cases
   if (auto *op = value.getDefiningOp())
     return const_cast<Pattern &>(pattern).match(op);
   return false;
 }
 
 /// Entry point for matching a pattern over an Operation.
 template <typename Pattern>
 inline bool matchPattern(Operation *op, const Pattern &pattern) {
+  assert(op);
   return const_cast<Pattern &>(pattern).match(op);
 }
 
+/// Entry point for matching a pattern over an Attribute. Returns `false`
+/// when `attr` is null.
+template <typename Pattern>
+inline bool matchPattern(Attribute attr, const Pattern &pattern) {
+  static_assert(llvm::is_detected<detail::has_compatible_matcher_t, Pattern,
+                                  Attribute>::value,
+                "Pattern does not support matching Attributes");
+  if (!attr)
+    return false;
+  return const_cast<Pattern &>(pattern).match(attr);
+}
+
 /// Matches a constant holding a scalar/vector/tensor float (splat) and
 /// writes the float value to bind_value.
-inline detail::constant_float_op_binder
+inline detail::constant_float_value_binder
 m_ConstantFloat(FloatAttr::ValueType *bind_value) {
-  return detail::constant_float_op_binder(bind_value);
+  return detail::constant_float_value_binder(bind_value);
 }
 
 /// Matches a constant holding a scalar/vector/tensor integer (splat) and
 /// writes the integer value to bind_value.
-inline detail::constant_int_op_binder
+inline detail::constant_int_value_binder
 m_ConstantInt(IntegerAttr::ValueType *bind_value) {
-  return detail::constant_int_op_binder(bind_value);
+  return detail::constant_int_value_binder(bind_value);
 }
 
 template <typename OpType, typename... Matchers>
 auto m_Op(Matchers... matchers) {
   return detail::RecursivePatternMatcher<OpType, Matchers...>(matchers...);
 }
 
 namespace matchers {
 inline auto m_Any() { return detail::AnyValueMatcher(); }
 inline auto m_Any(Value *val) { return detail::AnyCapturedValueMatcher(val); }
 inline auto m_Val(Value v) { return detail::PatternMatcherValue(v); }
 } // namespace matchers
 
 } // namespace mlir
 
 #endif // MLIR_IR_MATCHERS_H

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

static int64_t getScalarOrElementWidth(Type type) {
   return -1;
 }
 
 static int64_t getScalarOrElementWidth(Value value) {
   return getScalarOrElementWidth(value.getType());
 }
 
 static FailureOr<APInt> getIntOrSplatIntValue(Attribute attr) {
-  if (auto intAttr = llvm::dyn_cast<IntegerAttr>(attr))
-    return intAttr.getValue();
-
-  if (auto splatAttr = llvm::dyn_cast<SplatElementsAttr>(attr))
-    if (llvm::isa<IntegerType>(splatAttr.getElementType()))
-      return splatAttr.getSplatValue<APInt>();
+  APInt value;
+  if (matchPattern(attr, m_ConstantInt(&value)))
+    return value;
 
   return failure();
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd canonicalization patterns
 //===----------------------------------------------------------------------===//
 

 }
 
 //===----------------------------------------------------------------------===//
 // AddIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::AddIOp::fold(FoldAdaptor adaptor) {
   // addi(x, 0) -> x
-  if (matchPattern(getRhs(), m_Zero()))
+  if (matchPattern(adaptor.getRhs(), m_Zero()))
     return getLhs();
 
   // addi(subi(a, b), b) -> a
   if (auto sub = getLhs().getDefiningOp<SubIOp>())
     if (getRhs() == sub.getRhs())
       return sub.getLhs();
 
   // addi(b, subi(a, b)) -> a

 // SubIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::SubIOp::fold(FoldAdaptor adaptor) {
   // subi(x,x) -> 0
   if (getOperand(0) == getOperand(1))
     return Builder(getContext()).getZeroAttr(getType());
   // subi(x,0) -> x
-  if (matchPattern(getRhs(), m_Zero()))
+  if (matchPattern(adaptor.getRhs(), m_Zero()))
     return getLhs();
 
   if (auto add = getLhs().getDefiningOp<AddIOp>()) {
     // subi(addi(a, b), b) -> a
     if (getRhs() == add.getRhs())
       return add.getLhs();
     // subi(addi(a, b), a) -> b
     if (getRhs() == add.getLhs())

 }
 
 //===----------------------------------------------------------------------===//
 // MulIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::MulIOp::fold(FoldAdaptor adaptor) {
   // muli(x, 0) -> 0
-  if (matchPattern(getRhs(), m_Zero()))
+  if (matchPattern(adaptor.getRhs(), m_Zero()))
     return getRhs();
   // muli(x, 1) -> x
-  if (matchPattern(getRhs(), m_One()))
-    return getOperand(0);
+  if (matchPattern(adaptor.getRhs(), m_One()))
+    return getLhs();
   // TODO: Handle the overflow case.
 
   // default folder
   return constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(),
       [](const APInt &a, const APInt &b) { return a * b; });
 }
 
if (auto vt = llvm::dyn_cast<VectorType>(getType(0)))
     return llvm::to_vector<4>(vt.getShape());
   return std::nullopt;
 }
 
 LogicalResult
 arith::MulSIExtendedOp::fold(FoldAdaptor adaptor,
                              SmallVectorImpl<OpFoldResult> &results) {
   // mulsi_extended(x, 0) -> 0, 0
-  if (matchPattern(getRhs(), m_Zero())) {
+  if (matchPattern(adaptor.getRhs(), m_Zero())) {
     Attribute zero = adaptor.getRhs();
     results.push_back(zero);
     results.push_back(zero);
     return success();
   }
 
   // mulsi_extended(cst_a, cst_b) -> cst_low, cst_high
   if (Attribute lowAttr = constFoldBinaryOp<IntegerAttr>(
if (auto vt = llvm::dyn_cast<VectorType>(getType(0)))
     return llvm::to_vector<4>(vt.getShape());
   return std::nullopt;
 }
 
 LogicalResult
 arith::MulUIExtendedOp::fold(FoldAdaptor adaptor,
                              SmallVectorImpl<OpFoldResult> &results) {
   // mului_extended(x, 0) -> 0, 0
-  if (matchPattern(getRhs(), m_Zero())) {
+  if (matchPattern(adaptor.getRhs(), m_Zero())) {
     Attribute zero = adaptor.getRhs();
     results.push_back(zero);
     results.push_back(zero);
     return success();
   }
 
   // mului_extended(x, 1) -> x, 0
-  if (matchPattern(getRhs(), m_One())) {
+  if (matchPattern(adaptor.getRhs(), m_One())) {
     Builder builder(getContext());
     Attribute zero = builder.getZeroAttr(getLhs().getType());
     results.push_back(getLhs());
     results.push_back(zero);
     return success();
   }
 
   // mului_extended(cst_a, cst_b) -> cst_low, cst_high

 }
 
 //===----------------------------------------------------------------------===//
 // DivUIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::DivUIOp::fold(FoldAdaptor adaptor) {
   // divui (x, 1) -> x.
-  if (matchPattern(getRhs(), m_One()))
+  if (matchPattern(adaptor.getRhs(), m_One()))
     return getLhs();
 
   // Don't fold if it would require a division by zero.
   bool div0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(),
                                                [&](APInt a, const APInt &b) {
                                                  if (div0 || !b) {
                                                    div0 = true;

 }
 
 //===----------------------------------------------------------------------===//
 // DivSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::DivSIOp::fold(FoldAdaptor adaptor) {
   // divsi (x, 1) -> x.
-  if (matchPattern(getRhs(), m_One()))
+  if (matchPattern(adaptor.getRhs(), m_One()))
     return getLhs();
 
   // Don't fold if it would overflow or if it requires a division by zero.
   bool overflowOrDiv0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(), [&](APInt a, const APInt &b) {
         if (overflowOrDiv0 || !b) {
           overflowOrDiv0 = true;

 }
 
 //===----------------------------------------------------------------------===//
 // CeilDivUIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::CeilDivUIOp::fold(FoldAdaptor adaptor) {
   // ceildivui (x, 1) -> x.
-  if (matchPattern(getRhs(), m_One()))
+  if (matchPattern(adaptor.getRhs(), m_One()))
     return getLhs();
 
   bool overflowOrDiv0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(), [&](APInt a, const APInt &b) {
         if (overflowOrDiv0 || !b) {
           overflowOrDiv0 = true;
           return a;

 }
 
 //===----------------------------------------------------------------------===//
 // CeilDivSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::CeilDivSIOp::fold(FoldAdaptor adaptor) {
   // ceildivsi (x, 1) -> x.
-  if (matchPattern(getRhs(), m_One()))
+  if (matchPattern(adaptor.getRhs(), m_One()))
     return getLhs();
 
   // Don't fold if it would overflow or if it requires a division by zero.
   bool overflowOrDiv0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(), [&](APInt a, const APInt &b) {
         if (overflowOrDiv0 || !b) {
           overflowOrDiv0 = true;

 }
 
 //===----------------------------------------------------------------------===//
 // FloorDivSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::FloorDivSIOp::fold(FoldAdaptor adaptor) {
   // floordivsi (x, 1) -> x.
-  if (matchPattern(getRhs(), m_One()))
+  if (matchPattern(adaptor.getRhs(), m_One()))
     return getLhs();
 
   // Don't fold if it would overflow or if it requires a division by zero.
   bool overflowOrDiv0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(), [&](APInt a, const APInt &b) {
         if (overflowOrDiv0 || !b) {
           overflowOrDiv0 = true;

 }
 
 //===----------------------------------------------------------------------===//
 // RemUIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::RemUIOp::fold(FoldAdaptor adaptor) {
   // remui (x, 1) -> 0.
-  if (matchPattern(getRhs(), m_One()))
+  if (matchPattern(adaptor.getRhs(), m_One()))
     return Builder(getContext()).getZeroAttr(getType());
 
   // Don't fold if it would require a division by zero.
   bool div0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(),
                                                [&](APInt a, const APInt &b) {
                                                  if (div0 || b.isZero()) {
                                                    div0 = true;
                                                    return a;
                                                  }
                                                  return a.urem(b);
                                                });
 
   return div0 ? Attribute() : result;
 }
 
 //===----------------------------------------------------------------------===//
 // RemSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::RemSIOp::fold(FoldAdaptor adaptor) {
   // remsi (x, 1) -> 0.
-  if (matchPattern(getRhs(), m_One()))
+  if (matchPattern(adaptor.getRhs(), m_One()))
     return Builder(getContext()).getZeroAttr(getType());
 
   // Don't fold if it would require a division by zero.
   bool div0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(),
                                                [&](APInt a, const APInt &b) {
                                                  if (div0 || b.isZero()) {
                                                    div0 = true;
for (bool reversePrev : {false, true}) {
 
     return prev.getResult();
   }
   return {};
 }
 
 OpFoldResult arith::AndIOp::fold(FoldAdaptor adaptor) {
   /// and(x, 0) -> 0
-  if (matchPattern(getRhs(), m_Zero()))
+  if (matchPattern(adaptor.getRhs(), m_Zero()))
     return getRhs();
   /// and(x, allOnes) -> x
   APInt intValue;
-  if (matchPattern(getRhs(), m_ConstantInt(&intValue)) && intValue.isAllOnes())
+  if (matchPattern(adaptor.getRhs(), m_ConstantInt(&intValue)) &&
+      intValue.isAllOnes())
     return getLhs();
   /// and(x, not(x)) -> 0
   if (matchPattern(getRhs(), m_Op<XOrIOp>(matchers::m_Val(getLhs()),
                                           m_ConstantInt(&intValue))) &&
       intValue.isAllOnes())
     return Builder(getContext()).getZeroAttr(getType());
   /// and(not(x), x) -> 0
   if (matchPattern(getLhs(), m_Op<XOrIOp>(matchers::m_Val(getRhs()),
return constFoldBinaryOp<IntegerAttr>(
       [](APInt a, const APInt &b) { return std::move(a) & b; });
 }
 
 //===----------------------------------------------------------------------===//
 // OrIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::OrIOp::fold(FoldAdaptor adaptor) {
+  if (APInt rhsVal; matchPattern(adaptor.getRhs(), m_ConstantInt(&rhsVal))) {
     /// or(x, 0) -> x
-  if (matchPattern(getRhs(), m_Zero()))
+    if (rhsVal.isZero())
       return getLhs();
     /// or(x, <all ones>) -> <all ones>
-  if (auto rhsAttr = llvm::dyn_cast_if_present<IntegerAttr>(adaptor.getRhs()))
-    if (rhsAttr.getValue().isAllOnes())
-      return rhsAttr;
+    if (rhsVal.isAllOnes())
+      return adaptor.getRhs();
+  }
 
   APInt intValue;
   /// or(x, xor(x, 1)) -> 1
   if (matchPattern(getRhs(), m_Op<XOrIOp>(matchers::m_Val(getLhs()),
                                           m_ConstantInt(&intValue))) &&
       intValue.isAllOnes())
     return getRhs().getDefiningOp<XOrIOp>().getRhs();
   /// or(xor(x, 1), x) -> 1
   if (matchPattern(getLhs(), m_Op<XOrIOp>(matchers::m_Val(getRhs()),
                                           m_ConstantInt(&intValue))) &&
       intValue.isAllOnes())
     return getLhs().getDefiningOp<XOrIOp>().getRhs();
 
   return constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(),
       [](APInt a, const APInt &b) { return std::move(a) | b; });
 }
 
 //===----------------------------------------------------------------------===//
 // XOrIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::XOrIOp::fold(FoldAdaptor adaptor) {
   /// xor(x, 0) -> x
-  if (matchPattern(getRhs(), m_Zero()))
+  if (matchPattern(adaptor.getRhs(), m_Zero()))
     return getLhs();
   /// xor(x, x) -> 0
   if (getLhs() == getRhs())
     return Builder(getContext()).getZeroAttr(getType());
   /// xor(xor(x, a), a) -> x
   /// xor(xor(a, x), a) -> x
   if (arith::XOrIOp prev = getLhs().getDefiningOp<arith::XOrIOp>()) {
     if (prev.getRhs() == getRhs())

 }
 
 //===----------------------------------------------------------------------===//
 // AddFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::AddFOp::fold(FoldAdaptor adaptor) {
   // addf(x, -0) -> x
-  if (matchPattern(getRhs(), m_NegZeroFloat()))
+  if (matchPattern(adaptor.getRhs(), m_NegZeroFloat()))
     return getLhs();
 
   return constFoldBinaryOp<FloatAttr>(
       adaptor.getOperands(),
       [](const APFloat &a, const APFloat &b) { return a + b; });
 }
 
 //===----------------------------------------------------------------------===//
 // SubFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::SubFOp::fold(FoldAdaptor adaptor) {
   // subf(x, +0) -> x
-  if (matchPattern(getRhs(), m_PosZeroFloat()))
+  if (matchPattern(adaptor.getRhs(), m_PosZeroFloat()))
     return getLhs();
 
   return constFoldBinaryOp<FloatAttr>(
       adaptor.getOperands(),
       [](const APFloat &a, const APFloat &b) { return a - b; });
 }
 
 //===----------------------------------------------------------------------===//
 // MaxFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::MaxFOp::fold(FoldAdaptor adaptor) {
   // maxf(x,x) -> x
   if (getLhs() == getRhs())
     return getRhs();
 
   // maxf(x, -inf) -> x
-  if (matchPattern(getRhs(), m_NegInfFloat()))
+  if (matchPattern(adaptor.getRhs(), m_NegInfFloat()))
     return getLhs();
 
   return constFoldBinaryOp<FloatAttr>(
       adaptor.getOperands(),
       [](const APFloat &a, const APFloat &b) { return llvm::maximum(a, b); });
 }
 
 //===----------------------------------------------------------------------===//
 // MaxSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MaxSIOp::fold(FoldAdaptor adaptor) {
   // maxsi(x,x) -> x
   if (getLhs() == getRhs())
     return getRhs();
 
-  APInt intValue;
+  if (APInt intValue;
+      matchPattern(adaptor.getRhs(), m_ConstantInt(&intValue))) {
     // maxsi(x,MAX_INT) -> MAX_INT
-  if (matchPattern(getRhs(), m_ConstantInt(&intValue)) &&
-      intValue.isMaxSignedValue())
+    if (intValue.isMaxSignedValue())
       return getRhs();
-
     // maxsi(x, MIN_INT) -> x
-  if (matchPattern(getRhs(), m_ConstantInt(&intValue)) &&
-      intValue.isMinSignedValue())
+    if (intValue.isMinSignedValue())
       return getLhs();
+  }
 
   return constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(),
                                         [](const APInt &a, const APInt &b) {
                                           return llvm::APIntOps::smax(a, b);
                                         });
 }
 
 //===----------------------------------------------------------------------===//
 // MaxUIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MaxUIOp::fold(FoldAdaptor adaptor) {
   // maxui(x,x) -> x
   if (getLhs() == getRhs())
     return getRhs();
 
-  APInt intValue;
+  if (APInt intValue;
+      matchPattern(adaptor.getRhs(), m_ConstantInt(&intValue))) {
     // maxui(x,MAX_INT) -> MAX_INT
-  if (matchPattern(getRhs(), m_ConstantInt(&intValue)) && intValue.isMaxValue())
+    if (intValue.isMaxValue())
       return getRhs();
-
     // maxui(x, MIN_INT) -> x
-  if (matchPattern(getRhs(), m_ConstantInt(&intValue)) && intValue.isMinValue())
+    if (intValue.isMinValue())
       return getLhs();
+  }
 
   return constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(),
                                         [](const APInt &a, const APInt &b) {
                                           return llvm::APIntOps::umax(a, b);
                                         });
 }
 
 //===----------------------------------------------------------------------===//
 // MinFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::MinFOp::fold(FoldAdaptor adaptor) {
   // minf(x,x) -> x
   if (getLhs() == getRhs())
     return getRhs();
 
   // minf(x, +inf) -> x
-  if (matchPattern(getRhs(), m_PosInfFloat()))
+  if (matchPattern(adaptor.getRhs(), m_PosInfFloat()))
     return getLhs();
 
   return constFoldBinaryOp<FloatAttr>(
       adaptor.getOperands(),
       [](const APFloat &a, const APFloat &b) { return llvm::minimum(a, b); });
 }
 
 //===----------------------------------------------------------------------===//
 // MinSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MinSIOp::fold(FoldAdaptor adaptor) {
   // minsi(x,x) -> x
   if (getLhs() == getRhs())
     return getRhs();
 
-  APInt intValue;
+  if (APInt intValue;
+      matchPattern(adaptor.getRhs(), m_ConstantInt(&intValue))) {
     // minsi(x,MIN_INT) -> MIN_INT
-  if (matchPattern(getRhs(), m_ConstantInt(&intValue)) &&
-      intValue.isMinSignedValue())
+    if (intValue.isMinSignedValue())
       return getRhs();
-
     // minsi(x, MAX_INT) -> x
-  if (matchPattern(getRhs(), m_ConstantInt(&intValue)) &&
-      intValue.isMaxSignedValue())
+    if (intValue.isMaxSignedValue())
       return getLhs();
+  }
 
   return constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(),
                                         [](const APInt &a, const APInt &b) {
                                           return llvm::APIntOps::smin(a, b);
                                         });
 }
 
 //===----------------------------------------------------------------------===//
 // MinUIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MinUIOp::fold(FoldAdaptor adaptor) {
   // minui(x,x) -> x
   if (getLhs() == getRhs())
     return getRhs();
 
-  APInt intValue;
+  if (APInt intValue;
+      matchPattern(adaptor.getRhs(), m_ConstantInt(&intValue))) {
     // minui(x,MIN_INT) -> MIN_INT
-  if (matchPattern(getRhs(), m_ConstantInt(&intValue)) && intValue.isMinValue())
+    if (intValue.isMinValue())
       return getRhs();
-
     // minui(x, MAX_INT) -> x
-  if (matchPattern(getRhs(), m_ConstantInt(&intValue)) && intValue.isMaxValue())
+    if (intValue.isMaxValue())
       return getLhs();
+  }
 
   return constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(),
                                         [](const APInt &a, const APInt &b) {
                                           return llvm::APIntOps::umin(a, b);
                                         });
 }
 
 //===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::MulFOp::fold(FoldAdaptor adaptor) {
   // mulf(x, 1) -> x
-  if (matchPattern(getRhs(), m_OneFloat()))
+  if (matchPattern(adaptor.getRhs(), m_OneFloat()))
     return getLhs();
 
   return constFoldBinaryOp<FloatAttr>(
       adaptor.getOperands(),
       [](const APFloat &a, const APFloat &b) { return a * b; });
 }
 
 void arith::MulFOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                 MLIRContext *context) {
   patterns.add<MulFOfNegF>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // DivFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::DivFOp::fold(FoldAdaptor adaptor) {
   // divf(x, 1) -> x
-  if (matchPattern(getRhs(), m_OneFloat()))
+  if (matchPattern(adaptor.getRhs(), m_OneFloat()))
     return getLhs();
 
   return constFoldBinaryOp<FloatAttr>(
       adaptor.getOperands(),
       [](const APFloat &a, const APFloat &b) { return a / b; });
 }
 
 void arith::DivFOp::getCanonicalizationPatterns(RewritePatternSet &patterns,

 
 OpFoldResult arith::CmpIOp::fold(FoldAdaptor adaptor) {
   // cmpi(pred, x, x)
   if (getLhs() == getRhs()) {
     auto val = applyCmpPredicateToEqualOperands(getPredicate());
     return getBoolAttribute(getType(), getContext(), val);
   }
 
-  if (matchPattern(getRhs(), m_Zero())) {
+  if (matchPattern(adaptor.getRhs(), m_Zero())) {
     if (auto extOp = getLhs().getDefiningOp<ExtSIOp>()) {
       // extsi(%x : i1 -> iN) != 0  ->  %x
       std::optional<int64_t> integerWidth =
           getIntegerWidth(extOp.getOperand().getType());
       if (integerWidth && integerWidth.value() == 1 &&
           getPredicate() == arith::CmpIPredicate::ne)
         return extOp.getOperand();
     }
OpFoldResult arith::SelectOp::fold(FoldAdaptor adaptor) {
   Value trueVal = getTrueValue();
   Value falseVal = getFalseValue();
   if (trueVal == falseVal)
     return trueVal;
 
   Value condition = getCondition();
 
   // select true, %0, %1 => %0
-  if (matchPattern(condition, m_One()))
+  if (matchPattern(adaptor.getCondition(), m_One()))
     return trueVal;
 
   // select false, %0, %1 => %1
-  if (matchPattern(condition, m_Zero()))
+  if (matchPattern(adaptor.getCondition(), m_Zero()))
     return falseVal;
 
   // If either operand is fully poisoned, return the other.
   if (isa_and_nonnull<ub::PoisonAttr>(adaptor.getTrueValue()))
     return falseVal;
 
   if (isa_and_nonnull<ub::PoisonAttr>(adaptor.getFalseValue()))
     return trueVal;
 
   // select %x, true, false => %x
-  if (getType().isInteger(1) && matchPattern(getTrueValue(), m_One()) &&
-      matchPattern(getFalseValue(), m_Zero()))
+  if (getType().isInteger(1) && matchPattern(adaptor.getTrueValue(), m_One()) &&
+      matchPattern(adaptor.getFalseValue(), m_Zero()))
     return condition;
 
   if (auto cmp = dyn_cast_or_null<arith::CmpIOp>(condition.getDefiningOp())) {
     auto pred = cmp.getPredicate();
     if (pred == arith::CmpIPredicate::eq || pred == arith::CmpIPredicate::ne) {
       auto cmpLhs = cmp.getLhs();
       auto cmpRhs = cmp.getRhs();
 
LogicalResult arith::SelectOp::verify() {
   return success();
 }
 //===----------------------------------------------------------------------===//
 // ShLIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::ShLIOp::fold(FoldAdaptor adaptor) {
   // shli(x, 0) -> x
-  if (matchPattern(getRhs(), m_Zero()))
+  if (matchPattern(adaptor.getRhs(), m_Zero()))
     return getLhs();
   // Don't fold if shifting more than the bit width.
   bool bounded = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(), [&](const APInt &a, const APInt &b) {
         bounded = b.ule(b.getBitWidth());
         return a.shl(b);
       });
   return bounded ? result : Attribute();
 }
 
 //===----------------------------------------------------------------------===//
 // ShRUIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::ShRUIOp::fold(FoldAdaptor adaptor) {
   // shrui(x, 0) -> x
-  if (matchPattern(getRhs(), m_Zero()))
+  if (matchPattern(adaptor.getRhs(), m_Zero()))
     return getLhs();
   // Don't fold if shifting more than the bit width.
   bool bounded = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(), [&](const APInt &a, const APInt &b) {
         bounded = b.ule(b.getBitWidth());
         return a.lshr(b);
       });
   return bounded ? result : Attribute();
 }
 
 //===----------------------------------------------------------------------===//
 // ShRSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::ShRSIOp::fold(FoldAdaptor adaptor) {
   // shrsi(x, 0) -> x
-  if (matchPattern(getRhs(), m_Zero()))
+  if (matchPattern(adaptor.getRhs(), m_Zero()))
     return getLhs();
   // Don't fold if shifting more than the bit width.
   bool bounded = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(), [&](const APInt &a, const APInt &b) {
         bounded = b.ule(b.getBitWidth());
         return a.ashr(b);
       });

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOpTraits.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/FunctionImplementation.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"

 LogicalResult spirv::ShiftRightLogicalOp::verify() {
   return verifyShiftOp(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.BtiwiseAndOp
 //===----------------------------------------------------------------------===//
 
-static std::optional<APInt> extractIntConstant(Attribute attr) {
-  IntegerAttr intAttr;
-  if (auto splat = dyn_cast_if_present<SplatElementsAttr>(attr))
-    intAttr = dyn_cast<IntegerAttr>(splat.getSplatValue<Attribute>());
-  else
-    intAttr = dyn_cast_if_present<IntegerAttr>(attr);
-
-  if (!intAttr)
-    return std::nullopt;
-
-  return intAttr.getValue();
-}
-
 OpFoldResult
 spirv::BitwiseAndOp::fold(spirv::BitwiseAndOp::FoldAdaptor adaptor) {
-  std::optional<APInt> rhsVal = extractIntConstant(adaptor.getOperand2());
-  if (!rhsVal)
+  APInt rhsMask;
+  if (!matchPattern(adaptor.getOperand2(), m_ConstantInt(&rhsMask)))
     return {};
 
-  APInt rhsMask = *rhsVal;
-
   // x & 0 -> 0
   if (rhsMask.isZero())
     return getOperand2();
 
   // x & <all ones> -> x
   if (rhsMask.isAllOnes())
     return getOperand1();
 
   // (UConvert x : iN to iK) & <mask with N low bits set> -> UConvert x
   if (auto zext = getOperand1().getDefiningOp<spirv::UConvertOp>()) {
     int valueBits =
         getElementTypeOrSelf(zext.getOperand()).getIntOrFloatBitWidth();
     if (rhsMask.zextOrTrunc(valueBits).isAllOnes())
       return getOperand1();
   }
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.BtiwiseOrOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult spirv::BitwiseOrOp::fold(spirv::BitwiseOrOp::FoldAdaptor adaptor) {
-  std::optional<APInt> rhsVal = extractIntConstant(adaptor.getOperand2());
-  if (!rhsVal)
+  APInt rhsMask;
+  if (!matchPattern(adaptor.getOperand2(), m_ConstantInt(&rhsMask)))
     return {};
 
-  APInt rhsMask = *rhsVal;
-
   // x | 0 -> x
   if (rhsMask.isZero())
     return getOperand1();
 
   // x | <all ones> -> <all ones>
   if (rhsMask.isAllOnes())
     return getOperand2();