| //===- SDBMExpr.cpp - MLIR SDBM Expression implementation -----------------===// |
| // |
| // Copyright 2019 The MLIR Authors. |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| // ============================================================================= |
| // |
| // A striped difference-bound matrix (SDBM) expression is a constant expression, |
| // an identifier, a binary expression with constant RHS and +, stripe operators |
| // or a difference expression between two identifiers. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/SDBM/SDBMExpr.h" |
| #include "SDBMExprDetail.h" |
| #include "mlir/IR/AffineExpr.h" |
| #include "mlir/IR/AffineExprVisitor.h" |
| #include "mlir/SDBM/SDBMDialect.h" |
| |
| #include "llvm/Support/raw_ostream.h" |
| |
| using namespace mlir; |
| |
| namespace { |
| /// A simple compositional matcher for AffineExpr |
| /// |
| /// Example usage: |
| /// |
| /// ```c++ |
| /// AffineExprMatcher x, C, m; |
| /// AffineExprMatcher pattern1 = ((x % C) * m) + x; |
| /// AffineExprMatcher pattern2 = x + ((x % C) * m); |
| /// if (pattern1.match(expr) || pattern2.match(expr)) { |
| /// ... |
| /// } |
| /// ``` |
| class AffineExprMatcherStorage; |
| class AffineExprMatcher { |
| public: |
| AffineExprMatcher(); |
| AffineExprMatcher(const AffineExprMatcher &other); |
| |
| AffineExprMatcher operator+(AffineExprMatcher other) { |
| return AffineExprMatcher(AffineExprKind::Add, *this, other); |
| } |
| AffineExprMatcher operator*(AffineExprMatcher other) { |
| return AffineExprMatcher(AffineExprKind::Mul, *this, other); |
| } |
| AffineExprMatcher floorDiv(AffineExprMatcher other) { |
| return AffineExprMatcher(AffineExprKind::FloorDiv, *this, other); |
| } |
| AffineExprMatcher ceilDiv(AffineExprMatcher other) { |
| return AffineExprMatcher(AffineExprKind::CeilDiv, *this, other); |
| } |
| AffineExprMatcher operator%(AffineExprMatcher other) { |
| return AffineExprMatcher(AffineExprKind::Mod, *this, other); |
| } |
| |
| AffineExpr match(AffineExpr expr); |
| AffineExpr matched(); |
| Optional<int> getMatchedConstantValue(); |
| |
| private: |
| AffineExprMatcher(AffineExprKind k, AffineExprMatcher a, AffineExprMatcher b); |
| AffineExprKind kind; // only used to match in binary op cases. |
| // A shared_ptr allows multiple references to same matcher storage without |
| // worrying about ownership or dealing with an arena. To be cleaned up if we |
| // go with this. |
| std::shared_ptr<AffineExprMatcherStorage> storage; |
| }; |
| |
| class AffineExprMatcherStorage { |
| public: |
| AffineExprMatcherStorage() {} |
| AffineExprMatcherStorage(const AffineExprMatcherStorage &other) |
| : subExprs(other.subExprs.begin(), other.subExprs.end()), |
| matched(other.matched) {} |
| AffineExprMatcherStorage(ArrayRef<AffineExprMatcher> exprs) |
| : subExprs(exprs.begin(), exprs.end()) {} |
| AffineExprMatcherStorage(AffineExprMatcher &a, AffineExprMatcher &b) |
| : subExprs({a, b}) {} |
| llvm::SmallVector<AffineExprMatcher, 0> subExprs; |
| AffineExpr matched; |
| }; |
| } // namespace |
| |
| AffineExprMatcher::AffineExprMatcher() |
| : kind(AffineExprKind::Constant), storage(new AffineExprMatcherStorage()) {} |
| |
| AffineExprMatcher::AffineExprMatcher(const AffineExprMatcher &other) |
| : kind(other.kind), storage(other.storage) {} |
| |
| Optional<int> AffineExprMatcher::getMatchedConstantValue() { |
| if (auto cst = storage->matched.dyn_cast<AffineConstantExpr>()) |
| return cst.getValue(); |
| return None; |
| } |
| |
| AffineExpr AffineExprMatcher::match(AffineExpr expr) { |
| if (kind > AffineExprKind::LAST_AFFINE_BINARY_OP) { |
| if (storage->matched) |
| if (storage->matched != expr) |
| return AffineExpr(); |
| storage->matched = expr; |
| return storage->matched; |
| } |
| if (kind != expr.getKind()) { |
| return AffineExpr(); |
| } |
| if (auto bin = expr.dyn_cast<AffineBinaryOpExpr>()) { |
| if (!storage->subExprs.empty() && |
| !storage->subExprs[0].match(bin.getLHS())) { |
| return AffineExpr(); |
| } |
| if (!storage->subExprs.empty() && |
| !storage->subExprs[1].match(bin.getRHS())) { |
| return AffineExpr(); |
| } |
| if (storage->matched) |
| if (storage->matched != expr) |
| return AffineExpr(); |
| storage->matched = expr; |
| return storage->matched; |
| } |
| llvm_unreachable("binary expected"); |
| } |
| |
| AffineExpr AffineExprMatcher::matched() { return storage->matched; } |
| |
| AffineExprMatcher::AffineExprMatcher(AffineExprKind k, AffineExprMatcher a, |
| AffineExprMatcher b) |
| : kind(k), storage(new AffineExprMatcherStorage(a, b)) { |
| storage->subExprs.push_back(a); |
| storage->subExprs.push_back(b); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // SDBMExpr |
| //===----------------------------------------------------------------------===// |
| |
| SDBMExprKind SDBMExpr::getKind() const { return impl->getKind(); } |
| |
| MLIRContext *SDBMExpr::getContext() const { |
| return impl->dialect->getContext(); |
| } |
| |
| SDBMDialect *SDBMExpr::getDialect() const { return impl->dialect; } |
| |
| void SDBMExpr::print(raw_ostream &os) const { |
| struct Printer : public SDBMVisitor<Printer> { |
| Printer(raw_ostream &ostream) : prn(ostream) {} |
| |
| void visitSum(SDBMSumExpr expr) { |
| visitVarying(expr.getLHS()); |
| prn << " + "; |
| visitConstant(expr.getRHS()); |
| } |
| void visitDiff(SDBMDiffExpr expr) { |
| visitPositive(expr.getLHS()); |
| prn << " - "; |
| visitPositive(expr.getRHS()); |
| } |
| void visitDim(SDBMDimExpr expr) { prn << 'd' << expr.getPosition(); } |
| void visitSymbol(SDBMSymbolExpr expr) { prn << 's' << expr.getPosition(); } |
| void visitStripe(SDBMStripeExpr expr) { |
| visitPositive(expr.getVar()); |
| prn << " # "; |
| visitConstant(expr.getStripeFactor()); |
| } |
| void visitNeg(SDBMNegExpr expr) { |
| prn << '-'; |
| visitPositive(expr.getVar()); |
| } |
| void visitConstant(SDBMConstantExpr expr) { prn << expr.getValue(); } |
| |
| raw_ostream &prn; |
| }; |
| Printer printer(os); |
| printer.visit(*this); |
| } |
| |
| void SDBMExpr::dump() const { |
| print(llvm::errs()); |
| llvm::errs() << '\n'; |
| } |
| |
| namespace { |
| // Helper class to perform negation of an SDBM expression. |
| struct SDBMNegator : public SDBMVisitor<SDBMNegator, SDBMExpr> { |
| // Any positive expression is wrapped into a negation expression. |
| // -(x) = -x |
| SDBMExpr visitPositive(SDBMPositiveExpr expr) { |
| return SDBMNegExpr::get(expr); |
| } |
| // A negation expression is unwrapped. |
| // -(-x) = x |
| SDBMExpr visitNeg(SDBMNegExpr expr) { return expr.getVar(); } |
| // The value of the constant is negated. |
| SDBMExpr visitConstant(SDBMConstantExpr expr) { |
| return SDBMConstantExpr::get(expr.getDialect(), -expr.getValue()); |
| } |
| // Both terms of the sum are negated recursively. |
| SDBMExpr visitSum(SDBMSumExpr expr) { |
| return SDBMSumExpr::get(visit(expr.getLHS()).cast<SDBMVaryingExpr>(), |
| visit(expr.getRHS()).cast<SDBMConstantExpr>()); |
| } |
| // Terms of a difference are interchanged. |
| // -(x - y) = y - x |
| SDBMExpr visitDiff(SDBMDiffExpr expr) { |
| return SDBMDiffExpr::get(expr.getRHS(), expr.getLHS()); |
| } |
| }; |
| } // namespace |
| |
| SDBMExpr SDBMExpr::operator-() { return SDBMNegator().visit(*this); } |
| |
| //===----------------------------------------------------------------------===// |
| // SDBMSumExpr |
| //===----------------------------------------------------------------------===// |
| |
| SDBMSumExpr SDBMSumExpr::get(SDBMVaryingExpr lhs, SDBMConstantExpr rhs) { |
| assert(lhs && "expected SDBM variable expression"); |
| assert(rhs && "expected SDBM constant"); |
| |
| // If LHS of a sum is another sum, fold the constant RHS parts. |
| if (auto lhsSum = lhs.dyn_cast<SDBMSumExpr>()) { |
| lhs = lhsSum.getLHS(); |
| rhs = SDBMConstantExpr::get(rhs.getDialect(), |
| rhs.getValue() + lhsSum.getRHS().getValue()); |
| } |
| |
| StorageUniquer &uniquer = lhs.getDialect()->getUniquer(); |
| return uniquer.get<detail::SDBMBinaryExprStorage>( |
| /*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Add), lhs, rhs); |
| } |
| |
| SDBMVaryingExpr SDBMSumExpr::getLHS() const { |
| return static_cast<ImplType *>(impl)->lhs; |
| } |
| |
| SDBMConstantExpr SDBMSumExpr::getRHS() const { |
| return static_cast<ImplType *>(impl)->rhs; |
| } |
| |
| AffineExpr SDBMExpr::getAsAffineExpr() const { |
| struct Converter : public SDBMVisitor<Converter, AffineExpr> { |
| AffineExpr visitSum(SDBMSumExpr expr) { |
| AffineExpr lhs = visit(expr.getLHS()), rhs = visit(expr.getRHS()); |
| return lhs + rhs; |
| } |
| |
| AffineExpr visitStripe(SDBMStripeExpr expr) { |
| AffineExpr lhs = visit(expr.getVar()), |
| rhs = visit(expr.getStripeFactor()); |
| return lhs - (lhs % rhs); |
| } |
| |
| AffineExpr visitDiff(SDBMDiffExpr expr) { |
| AffineExpr lhs = visit(expr.getLHS()), rhs = visit(expr.getRHS()); |
| return lhs - rhs; |
| } |
| |
| AffineExpr visitDim(SDBMDimExpr expr) { |
| return getAffineDimExpr(expr.getPosition(), expr.getContext()); |
| } |
| |
| AffineExpr visitSymbol(SDBMSymbolExpr expr) { |
| return getAffineSymbolExpr(expr.getPosition(), expr.getContext()); |
| } |
| |
| AffineExpr visitNeg(SDBMNegExpr expr) { |
| return getAffineBinaryOpExpr(AffineExprKind::Mul, |
| getAffineConstantExpr(-1, expr.getContext()), |
| visit(expr.getVar())); |
| } |
| |
| AffineExpr visitConstant(SDBMConstantExpr expr) { |
| return getAffineConstantExpr(expr.getValue(), expr.getContext()); |
| } |
| } converter; |
| return converter.visit(*this); |
| } |
| |
| Optional<SDBMExpr> SDBMExpr::tryConvertAffineExpr(AffineExpr affine) { |
| struct Converter : public AffineExprVisitor<Converter, SDBMExpr> { |
| SDBMExpr visitAddExpr(AffineBinaryOpExpr expr) { |
| // Attempt to recover a stripe expression. Because AffineExprs don't have |
| // a first-class difference kind, we check for both x + -1 * (x mod C) and |
| // -1 * (x mod C) + x cases. |
| AffineExprMatcher x, C, m; |
| AffineExprMatcher pattern1 = ((x % C) * m) + x; |
| AffineExprMatcher pattern2 = x + ((x % C) * m); |
| if ((pattern1.match(expr) && m.getMatchedConstantValue() == -1) || |
| (pattern2.match(expr) && m.getMatchedConstantValue() == -1)) { |
| if (auto convertedLHS = visit(x.matched())) { |
| // TODO(ntv): return convertedLHS.stripe(C); |
| return SDBMStripeExpr::get( |
| convertedLHS.cast<SDBMPositiveExpr>(), |
| visit(C.matched()).cast<SDBMConstantExpr>()); |
| } |
| } |
| auto lhs = visit(expr.getLHS()), rhs = visit(expr.getRHS()); |
| if (!lhs || !rhs) |
| return {}; |
| |
| // In a "add" AffineExpr, the constant always appears on the right. If |
| // there were two constants, they would have been folded away. |
| assert(!lhs.isa<SDBMConstantExpr>() && "non-canonical affine expression"); |
| auto rhsConstant = rhs.dyn_cast<SDBMConstantExpr>(); |
| |
| // SDBM accepts LHS variables and RHS constants in a sum. |
| auto lhsVar = lhs.dyn_cast<SDBMVaryingExpr>(); |
| auto rhsVar = rhs.dyn_cast<SDBMVaryingExpr>(); |
| if (rhsConstant && lhsVar) |
| return SDBMSumExpr::get(lhsVar, rhsConstant); |
| |
| // The sum of a negated variable and a non-negated variable is a |
| // difference, supported as a special kind in SDBM. Because AffineExprs |
| // don't have first-class difference kind, check both LHS and RHS for |
| // negation. |
| auto lhsPos = lhs.dyn_cast<SDBMPositiveExpr>(); |
| auto rhsPos = rhs.dyn_cast<SDBMPositiveExpr>(); |
| auto lhsNeg = lhs.dyn_cast<SDBMNegExpr>(); |
| auto rhsNeg = rhs.dyn_cast<SDBMNegExpr>(); |
| if (lhsNeg && rhsVar) |
| return SDBMDiffExpr::get(rhsPos, lhsNeg.getVar()); |
| if (rhsNeg && lhsVar) |
| return SDBMDiffExpr::get(lhsPos, rhsNeg.getVar()); |
| |
| // Other cases don't fit into SDBM. |
| return {}; |
| } |
| |
| SDBMExpr visitMulExpr(AffineBinaryOpExpr expr) { |
| // Attempt to recover a stripe expression "x # C = (x floordiv C) * C". |
| AffineExprMatcher x, C; |
| AffineExprMatcher pattern = (x.floorDiv(C)) * C; |
| if (pattern.match(expr)) { |
| if (SDBMExpr converted = visit(x.matched())) { |
| if (auto varConverted = converted.dyn_cast<SDBMPositiveExpr>()) |
| // TODO(ntv): return varConverted.stripe(C.getConstantValue()); |
| return SDBMStripeExpr::get( |
| varConverted, |
| SDBMConstantExpr::get(dialect, |
| C.getMatchedConstantValue().getValue())); |
| } |
| } |
| |
| auto lhs = visit(expr.getLHS()), rhs = visit(expr.getRHS()); |
| if (!lhs || !rhs) |
| return {}; |
| |
| // In a "mul" AffineExpr, the constant always appears on the right. If |
| // there were two constants, they would have been folded away. |
| assert(!lhs.isa<SDBMConstantExpr>() && "non-canonical affine expression"); |
| auto rhsConstant = rhs.dyn_cast<SDBMConstantExpr>(); |
| if (!rhsConstant) |
| return {}; |
| |
| // The only supported "multiplication" expression is an SDBM is dimension |
| // negation, that is a product of dimension and constant -1. |
| auto lhsVar = lhs.dyn_cast<SDBMPositiveExpr>(); |
| if (lhsVar && rhsConstant.getValue() == -1) |
| return SDBMNegExpr::get(lhsVar); |
| |
| // Other multiplications are not allowed in SDBM. |
| return {}; |
| } |
| |
| SDBMExpr visitModExpr(AffineBinaryOpExpr expr) { |
| auto lhs = visit(expr.getLHS()), rhs = visit(expr.getRHS()); |
| if (!lhs || !rhs) |
| return {}; |
| |
| // 'mod' can only be converted to SDBM if its LHS is a variable |
| // and its RHS is a constant. Then it `x mod c = x - x stripe c`. |
| auto rhsConstant = rhs.dyn_cast<SDBMConstantExpr>(); |
| auto lhsVar = rhs.dyn_cast<SDBMPositiveExpr>(); |
| if (!lhsVar || !rhsConstant) |
| return {}; |
| return SDBMDiffExpr::get(lhsVar, |
| SDBMStripeExpr::get(lhsVar, rhsConstant)); |
| } |
| |
| // `a floordiv b = (a stripe b) / b`, but we have no division in SDBM |
| SDBMExpr visitFloorDivExpr(AffineBinaryOpExpr expr) { return {}; } |
| SDBMExpr visitCeilDivExpr(AffineBinaryOpExpr expr) { return {}; } |
| |
| // Dimensions, symbols and constants are converted trivially. |
| SDBMExpr visitConstantExpr(AffineConstantExpr expr) { |
| return SDBMConstantExpr::get(dialect, expr.getValue()); |
| } |
| SDBMExpr visitDimExpr(AffineDimExpr expr) { |
| return SDBMDimExpr::get(dialect, expr.getPosition()); |
| } |
| SDBMExpr visitSymbolExpr(AffineSymbolExpr expr) { |
| return SDBMSymbolExpr::get(dialect, expr.getPosition()); |
| } |
| |
| SDBMDialect *dialect; |
| } converter; |
| converter.dialect = affine.getContext()->getRegisteredDialect<SDBMDialect>(); |
| |
| if (auto result = converter.visit(affine)) |
| return result; |
| return None; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // SDBMDiffExpr |
| //===----------------------------------------------------------------------===// |
| |
| SDBMDiffExpr SDBMDiffExpr::get(SDBMPositiveExpr lhs, SDBMPositiveExpr rhs) { |
| assert(lhs && "expected SDBM dimension"); |
| assert(rhs && "expected SDBM dimension"); |
| |
| StorageUniquer &uniquer = lhs.getDialect()->getUniquer(); |
| return uniquer.get<detail::SDBMDiffExprStorage>( |
| /*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Diff), lhs, rhs); |
| } |
| |
| SDBMPositiveExpr SDBMDiffExpr::getLHS() const { |
| return static_cast<ImplType *>(impl)->lhs; |
| } |
| |
| SDBMPositiveExpr SDBMDiffExpr::getRHS() const { |
| return static_cast<ImplType *>(impl)->rhs; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // SDBMStripeExpr |
| //===----------------------------------------------------------------------===// |
| |
| SDBMStripeExpr SDBMStripeExpr::get(SDBMPositiveExpr var, |
| SDBMConstantExpr stripeFactor) { |
| assert(var && "expected SDBM variable expression"); |
| assert(stripeFactor && "expected non-null stripe factor"); |
| if (stripeFactor.getValue() <= 0) |
| llvm::report_fatal_error("non-positive stripe factor"); |
| |
| StorageUniquer &uniquer = var.getDialect()->getUniquer(); |
| return uniquer.get<detail::SDBMBinaryExprStorage>( |
| /*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Stripe), var, |
| stripeFactor); |
| } |
| |
| SDBMPositiveExpr SDBMStripeExpr::getVar() const { |
| if (SDBMVaryingExpr lhs = static_cast<ImplType *>(impl)->lhs) |
| return lhs.cast<SDBMPositiveExpr>(); |
| return {}; |
| } |
| |
| SDBMConstantExpr SDBMStripeExpr::getStripeFactor() const { |
| return static_cast<ImplType *>(impl)->rhs; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // SDBMInputExpr |
| //===----------------------------------------------------------------------===// |
| |
| unsigned SDBMInputExpr::getPosition() const { |
| return static_cast<ImplType *>(impl)->position; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // SDBMDimExpr |
| //===----------------------------------------------------------------------===// |
| |
| SDBMDimExpr SDBMDimExpr::get(SDBMDialect *dialect, unsigned position) { |
| assert(dialect && "expected non-null dialect"); |
| |
| auto assignDialect = [dialect](detail::SDBMPositiveExprStorage *storage) { |
| storage->dialect = dialect; |
| }; |
| |
| StorageUniquer &uniquer = dialect->getUniquer(); |
| return uniquer.get<detail::SDBMPositiveExprStorage>( |
| assignDialect, static_cast<unsigned>(SDBMExprKind::DimId), position); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // SDBMSymbolExpr |
| //===----------------------------------------------------------------------===// |
| |
| SDBMSymbolExpr SDBMSymbolExpr::get(SDBMDialect *dialect, unsigned position) { |
| assert(dialect && "expected non-null dialect"); |
| |
| auto assignDialect = [dialect](detail::SDBMPositiveExprStorage *storage) { |
| storage->dialect = dialect; |
| }; |
| |
| StorageUniquer &uniquer = dialect->getUniquer(); |
| return uniquer.get<detail::SDBMPositiveExprStorage>( |
| assignDialect, static_cast<unsigned>(SDBMExprKind::SymbolId), position); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // SDBMConstantExpr |
| //===----------------------------------------------------------------------===// |
| |
| SDBMConstantExpr SDBMConstantExpr::get(SDBMDialect *dialect, int64_t value) { |
| assert(dialect && "expected non-null dialect"); |
| |
| auto assignCtx = [dialect](detail::SDBMConstantExprStorage *storage) { |
| storage->dialect = dialect; |
| }; |
| |
| StorageUniquer &uniquer = dialect->getUniquer(); |
| return uniquer.get<detail::SDBMConstantExprStorage>( |
| assignCtx, static_cast<unsigned>(SDBMExprKind::Constant), value); |
| } |
| |
| int64_t SDBMConstantExpr::getValue() const { |
| return static_cast<ImplType *>(impl)->constant; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // SDBMNegExpr |
| //===----------------------------------------------------------------------===// |
| |
| SDBMNegExpr SDBMNegExpr::get(SDBMPositiveExpr var) { |
| assert(var && "expected non-null SDBM variable expression"); |
| |
| StorageUniquer &uniquer = var.getDialect()->getUniquer(); |
| return uniquer.get<detail::SDBMNegExprStorage>( |
| /*initFn=*/{}, static_cast<unsigned>(SDBMExprKind::Neg), var); |
| } |
| |
| SDBMPositiveExpr SDBMNegExpr::getVar() const { |
| return static_cast<ImplType *>(impl)->dim; |
| } |
| |
| namespace mlir { |
| namespace ops_assertions { |
| |
| SDBMExpr operator+(SDBMExpr lhs, SDBMExpr rhs) { |
| // If one of the operands is a negation, take a difference rather than a sum. |
| auto lhsNeg = lhs.dyn_cast<SDBMNegExpr>(); |
| auto rhsNeg = rhs.dyn_cast<SDBMNegExpr>(); |
| assert(!(lhsNeg && rhsNeg) && "a sum of negated expressions is a negation of " |
| "a sum of variables and not a correct SDBM"); |
| if (lhsNeg) |
| return rhs - lhsNeg.getVar(); |
| if (rhsNeg) |
| return lhs - rhsNeg.getVar(); |
| |
| // If LHS is a constant and RHS is not, swap the order to get into a supported |
| // sum case. From now on, RHS must be a constant. |
| auto lhsConstant = lhs.dyn_cast<SDBMConstantExpr>(); |
| auto rhsConstant = rhs.dyn_cast<SDBMConstantExpr>(); |
| if (!rhsConstant && lhsConstant) { |
| std::swap(lhs, rhs); |
| std::swap(lhsConstant, rhsConstant); |
| } |
| assert(rhsConstant && "at least one operand must be a constant"); |
| |
| // If LHS is another sum, first compute the sum of its variable |
| // part with the other argument and then add the constant part to enable |
| // constant folding (the variable part may, e.g., be a negation that requires |
| // to enter this function again). |
| auto lhsSum = lhs.dyn_cast<SDBMSumExpr>(); |
| if (lhsSum) |
| return lhsSum.getLHS() + |
| (lhsSum.getRHS().getValue() + rhsConstant.getValue()); |
| |
| // Constant-fold if LHS is a constant. |
| if (lhsConstant) |
| return SDBMConstantExpr::get(lhs.getDialect(), lhsConstant.getValue() + |
| rhsConstant.getValue()); |
| |
| // Fold x + 0 == x. |
| if (rhsConstant.getValue() == 0) |
| return lhs; |
| |
| return SDBMSumExpr::get(lhs.cast<SDBMVaryingExpr>(), |
| rhs.cast<SDBMConstantExpr>()); |
| } |
| |
| SDBMExpr operator-(SDBMExpr lhs, SDBMExpr rhs) { |
| // Fold x - x == 0. |
| if (lhs == rhs) |
| return SDBMConstantExpr::get(lhs.getDialect(), 0); |
| |
| // LHS and RHS may be constants. |
| auto lhsConstant = lhs.dyn_cast<SDBMConstantExpr>(); |
| auto rhsConstant = rhs.dyn_cast<SDBMConstantExpr>(); |
| |
| // Constant fold if both LHS and RHS are constants. |
| if (lhsConstant && rhsConstant) |
| return SDBMConstantExpr::get(lhs.getDialect(), lhsConstant.getValue() - |
| rhsConstant.getValue()); |
| |
| // Replace a difference with a sum with a negated value if one of LHS and RHS |
| // is a constant: |
| // x - C == x + (-C); |
| // C - x == -x + C. |
| // This calls into operator+ for further simplification. |
| if (rhsConstant) |
| return lhs + (-rhsConstant); |
| if (lhsConstant) |
| return -rhs + lhsConstant; |
| |
| // Hoist constant factors outside the difference if any of sides is a sum: |
| // (x + A) - (y - B) == x - y + (A - B). |
| // If either LHS or RHS is a sum, collect the constant values separately and |
| // update LHS and RHS to point to the variable part of the sum. |
| auto lhsSum = lhs.dyn_cast<SDBMSumExpr>(); |
| auto rhsSum = rhs.dyn_cast<SDBMSumExpr>(); |
| int64_t value = 0; |
| if (lhsSum) { |
| value += lhsSum.getRHS().getValue(); |
| lhs = lhsSum.getLHS(); |
| } |
| if (rhsSum) { |
| value -= rhsSum.getRHS().getValue(); |
| rhs = rhsSum.getLHS(); |
| } |
| |
| // This calls into operator+ for futher simplification in case value == 0. |
| return SDBMDiffExpr::get(lhs.cast<SDBMPositiveExpr>(), |
| rhs.cast<SDBMPositiveExpr>()) + |
| value; |
| } |
| |
| SDBMExpr stripe(SDBMExpr expr, SDBMExpr factor) { |
| auto constantFactor = factor.cast<SDBMConstantExpr>(); |
| assert(constantFactor.getValue() > 0 && "non-positive stripe"); |
| |
| // Fold x # 1 = x. |
| if (constantFactor.getValue() == 1) |
| return expr; |
| |
| return SDBMStripeExpr::get(expr.cast<SDBMPositiveExpr>(), constantFactor); |
| } |
| |
| } // namespace ops_assertions |
| } // namespace mlir |