44
55#include " lib/Dialect/MathExt/IR/MathExtOps.h"
66#include " llvm/include/llvm/ADT/APFloat.h" // from @llvm-project
7+ #include " mlir/include/mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
8+ #include " mlir/include/mlir/Dialect/Linalg/IR/Linalg.h" // from @llvm-project
79#include " mlir/include/mlir/IR/Attributes.h" // from @llvm-project
810#include " mlir/include/mlir/IR/BuiltinAttributes.h" // from @llvm-project
911#include " mlir/include/mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
@@ -40,6 +42,93 @@ static bool IsOne(mlir::Attribute attr) {
4042// populateWithGenerated, which can conflict with other generated patterns.
4143#include " lib/Transforms/ActivationCanonicalizations/Rewrites.cpp.inc"
4244
45+ // select(a > c, a, c) = max(a, c) for floats. This replaces the DRR
46+ // `SelectGreaterThanEqualFloat` pattern and folds in
47+ // two attr-forwarding behaviors so the polynomial-approximation domain survives
48+ // regardless of where torch-mlir attached it:
49+ // (a) copy discardable attrs off the select itself (the old DRR behavior),
50+ // (b) if the domain is still missing, copy `domain_lower`/`domain_upper` from
51+ // an enclosing `linalg.generic` (a torch ReLU imports as a generic
52+ // carrying those attrs, with `cmpf+select` in its body) The
53+ // generic's copy is dropped afterwards by stripForwardedDomains().
54+ // Either way the domain lands on the `arith.maximumf` that
55+ // PolynomialApproximation / ReluViaCompositeSign read; without it they fall
56+ // back to [-1, 1].
57+ struct SelectGreaterThanEqualFloatPattern
58+ : public OpRewritePattern<arith::SelectOp> {
59+ using OpRewritePattern<arith::SelectOp>::OpRewritePattern;
60+
61+ LogicalResult matchAndRewrite (arith::SelectOp op,
62+ PatternRewriter& rewriter) const override {
63+ auto cmpOp = op.
getCondition ().
getDefiningOp <arith::CmpFOp>();
64+ if (!cmpOp)
65+ return rewriter.notifyMatchFailure (op, " condition is not arith.cmpf"
66+
67+ auto pred = cmpOp.getPredicate ();
68+ if (pred != arith::CmpFPredicate::UGT && pred != arith::CmpFPredicate::UGE )
69+ return rewriter.notifyMatchFailure (op, " predicate is not ugt/uge"
70+
71+ // Must be the ReLU/max shape: select(a >? c, a, c).
72+ if (cmpOp.getLhs () != op.getTrueValue () ||
73+ cmpOp.getRhs () != op.getFalseValue ())
74+ return rewriter.notifyMatchFailure (op,
75+ " operands are not select(a>c,a,c)"
76+
77+ auto maxOp =
78+ arith::MaximumFOp::create (rewriter, op.getLoc (), op.getTrueValue (),
79+ op.getFalseValue (), cmpOp.getFastmathAttr ());
80+
81+ // (a) Forward any discardable attrs annotated on the select op itself onto
82+ // the maximumf (the old DRR `SelectGreaterThanEqualFloat` behavior). Covers
83+ // IR where the domain is attached directly to the select.
84+ for (auto attr : op->getDiscardableAttrs ())
85+ maxOp->setAttr (attr.getName (), attr.getValue ());
86+
87+ // (b) If the domain still isn't on the maximumf, it lives on an enclosing
88+ // linalg.generic instead (where torch-mlir's importer attaches the ReLU
89+ // domain). Copy it down onto the maximumf — the op PolynomialApproximation
90+ // actually reads. We only COPY here; the generic's own bounds are dropped
91+ // afterwards by stripForwardedDomains(). A single generic can hold several
92+ // ReLUs, so stripping as soon as the first one is rewritten would starve
93+ // the rest and silently leave them on the default [-1, 1] domain.
94+ if (auto generic = dyn_cast<linalg::GenericOp>(op->getParentOp ())) {
95+ Attribute lo = generic->getAttr (" domain_lower"
96+ Attribute hi = generic->getAttr (" domain_upper"
97+ if (lo && !maxOp->hasAttr (" domain_lower"
98+ maxOp->setAttr (" domain_lower"
99+ if (hi && !maxOp->hasAttr (" domain_upper"
100+ maxOp->setAttr (" domain_upper"
101+ }
102+
103+ rewriter.replaceOp (op, maxOp.getResult ());
104+ return success ();
105+ }
106+ };
107+
108+ // The domain bounds must end up on exactly one op. Once the patterns above have
109+ // copied an enclosing generic's bounds down onto every ReLU in its body, the
110+ // generic's own copy is redundant, so drop it: leaving the bounds on BOTH makes
111+ // a later activation-lifting pass merge two `domain_lower` entries into one
112+ // dictionary, tripping DictionaryAttr's uniqueness assertion.
113+ static void stripForwardedDomains (Operation* root) {
114+ root->walk ([](linalg::GenericOp generic) {
115+ if (!generic->hasAttr (" domain_lower" hasAttr (" domain_upper"
116+ return ;
117+ // Only strip if the bounds actually made it onto an op inside the body;
118+ // otherwise there was no ReLU to forward them to and they are still the
119+ // only record of the domain.
120+ bool forwarded = false ;
121+ generic->getRegion (0 ).walk ([&](Operation* inner) {
122+ if (inner->hasAttr (" domain_lower" hasAttr (" domain_upper"
123+ forwarded = true ;
124+ });
125+ if (forwarded) {
126+ generic->removeAttr (" domain_lower"
127+ generic->removeAttr (" domain_upper"
128+ }
129+ });
130+ }
131+
43132struct ActivationCanonicalizations
44133 : impl::ActivationCanonicalizationsBase<ActivationCanonicalizations> {
45134 using ActivationCanonicalizationsBase::ActivationCanonicalizationsBase;
@@ -48,8 +137,11 @@ struct ActivationCanonicalizations
48137 MLIRContext* context = &getContext ();
49138 RewritePatternSet patterns (context);
50139 populateWithGenerated (patterns);
140+ patterns.add <SelectGreaterThanEqualFloatPattern>(context);
51141
52142 (void )walkAndApplyPatterns (getOperation (), std::move (patterns));
143+
144+ stripForwardedDomains (getOperation ());
53145 }
54146};
55147
0 commit comments