blob: e97afac9ca2d72bfc33da19350bd5dc275d7e2f0 [file] [log] [blame]
//===- ExpandOps.cpp - Pass to legalize Arithmetic ops for LLVM lowering --===//
//
// 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 "PassDetail.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/DialectConversion.h"
using namespace mlir;
namespace {
/// Expands CeilDivUIOp (n, m) into
/// n == 0 ? 0 : ((n-1) / m) + 1
struct CeilDivUIOpConverter : public OpRewritePattern<arith::CeilDivUIOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(arith::CeilDivUIOp op,
PatternRewriter &rewriter) const final {
Location loc = op.getLoc();
Value a = op.getLhs();
Value b = op.getRhs();
Value zero = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(a.getType(), 0));
Value compare =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, a, zero);
Value one = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(a.getType(), 1));
Value minusOne = rewriter.create<arith::SubIOp>(loc, a, one);
Value quotient = rewriter.create<arith::DivUIOp>(loc, minusOne, b);
Value plusOne = rewriter.create<arith::AddIOp>(loc, quotient, one);
Value res = rewriter.create<SelectOp>(loc, compare, zero, plusOne);
rewriter.replaceOp(op, {res});
return success();
}
};
/// Expands CeilDivSIOp (n, m) into
/// 1) x = (m > 0) ? -1 : 1
/// 2) (n*m>0) ? ((n+x) / m) + 1 : - (-n / m)
struct CeilDivSIOpConverter : public OpRewritePattern<arith::CeilDivSIOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(arith::CeilDivSIOp op,
PatternRewriter &rewriter) const final {
Location loc = op.getLoc();
auto signedCeilDivIOp = cast<arith::CeilDivSIOp>(op);
Type type = signedCeilDivIOp.getType();
Value a = signedCeilDivIOp.getLhs();
Value b = signedCeilDivIOp.getRhs();
Value plusOne = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(type, 1));
Value zero = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(type, 0));
Value minusOne = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(type, -1));
// Compute x = (b>0) ? -1 : 1.
Value compare =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, b, zero);
Value x = rewriter.create<SelectOp>(loc, compare, minusOne, plusOne);
// Compute positive res: 1 + ((x+a)/b).
Value xPlusA = rewriter.create<arith::AddIOp>(loc, x, a);
Value xPlusADivB = rewriter.create<arith::DivSIOp>(loc, xPlusA, b);
Value posRes = rewriter.create<arith::AddIOp>(loc, plusOne, xPlusADivB);
// Compute negative res: - ((-a)/b).
Value minusA = rewriter.create<arith::SubIOp>(loc, zero, a);
Value minusADivB = rewriter.create<arith::DivSIOp>(loc, minusA, b);
Value negRes = rewriter.create<arith::SubIOp>(loc, zero, minusADivB);
// Result is (a*b>0) ? pos result : neg result.
// Note, we want to avoid using a*b because of possible overflow.
// The case that matters are a>0, a==0, a<0, b>0 and b<0. We do
// not particuliarly care if a*b<0 is true or false when b is zero
// as this will result in an illegal divide. So `a*b<0` can be reformulated
// as `(a<0 && b<0) || (a>0 && b>0)' or `(a<0 && b<0) || (a>0 && b>=0)'.
// We pick the first expression here.
Value aNeg =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, a, zero);
Value aPos =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, a, zero);
Value bNeg =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, b, zero);
Value bPos =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, b, zero);
Value firstTerm = rewriter.create<arith::AndIOp>(loc, aNeg, bNeg);
Value secondTerm = rewriter.create<arith::AndIOp>(loc, aPos, bPos);
Value compareRes =
rewriter.create<arith::OrIOp>(loc, firstTerm, secondTerm);
Value res = rewriter.create<SelectOp>(loc, compareRes, posRes, negRes);
// Perform substitution and return success.
rewriter.replaceOp(op, {res});
return success();
}
};
/// Expands FloorDivSIOp (n, m) into
/// 1) x = (m<0) ? 1 : -1
/// 2) return (n*m<0) ? - ((-n+x) / m) -1 : n / m
struct FloorDivSIOpConverter : public OpRewritePattern<arith::FloorDivSIOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(arith::FloorDivSIOp op,
PatternRewriter &rewriter) const final {
Location loc = op.getLoc();
arith::FloorDivSIOp signedFloorDivIOp = cast<arith::FloorDivSIOp>(op);
Type type = signedFloorDivIOp.getType();
Value a = signedFloorDivIOp.getLhs();
Value b = signedFloorDivIOp.getRhs();
Value plusOne = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(type, 1));
Value zero = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(type, 0));
Value minusOne = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(type, -1));
// Compute x = (b<0) ? 1 : -1.
Value compare =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, b, zero);
Value x = rewriter.create<SelectOp>(loc, compare, plusOne, minusOne);
// Compute negative res: -1 - ((x-a)/b).
Value xMinusA = rewriter.create<arith::SubIOp>(loc, x, a);
Value xMinusADivB = rewriter.create<arith::DivSIOp>(loc, xMinusA, b);
Value negRes = rewriter.create<arith::SubIOp>(loc, minusOne, xMinusADivB);
// Compute positive res: a/b.
Value posRes = rewriter.create<arith::DivSIOp>(loc, a, b);
// Result is (a*b<0) ? negative result : positive result.
// Note, we want to avoid using a*b because of possible overflow.
// The case that matters are a>0, a==0, a<0, b>0 and b<0. We do
// not particuliarly care if a*b<0 is true or false when b is zero
// as this will result in an illegal divide. So `a*b<0` can be reformulated
// as `(a>0 && b<0) || (a>0 && b<0)' or `(a>0 && b<0) || (a>0 && b<=0)'.
// We pick the first expression here.
Value aNeg =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, a, zero);
Value aPos =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, a, zero);
Value bNeg =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, b, zero);
Value bPos =
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, b, zero);
Value firstTerm = rewriter.create<arith::AndIOp>(loc, aNeg, bPos);
Value secondTerm = rewriter.create<arith::AndIOp>(loc, aPos, bNeg);
Value compareRes =
rewriter.create<arith::OrIOp>(loc, firstTerm, secondTerm);
Value res = rewriter.create<SelectOp>(loc, compareRes, negRes, posRes);
// Perform substitution and return success.
rewriter.replaceOp(op, {res});
return success();
}
};
template <typename OpTy, arith::CmpFPredicate pred>
struct MaxMinFOpConverter : public OpRewritePattern<OpTy> {
public:
using OpRewritePattern<OpTy>::OpRewritePattern;
LogicalResult matchAndRewrite(OpTy op,
PatternRewriter &rewriter) const final {
Value lhs = op.getLhs();
Value rhs = op.getRhs();
Location loc = op.getLoc();
Value cmp = rewriter.create<arith::CmpFOp>(loc, pred, lhs, rhs);
Value select = rewriter.create<SelectOp>(loc, cmp, lhs, rhs);
auto floatType = getElementTypeOrSelf(lhs.getType()).cast<FloatType>();
Value isNaN = rewriter.create<arith::CmpFOp>(loc, arith::CmpFPredicate::UNO,
lhs, rhs);
Value nan = rewriter.create<arith::ConstantFloatOp>(
loc, APFloat::getQNaN(floatType.getFloatSemantics()), floatType);
if (VectorType vectorType = lhs.getType().dyn_cast<VectorType>())
nan = rewriter.create<SplatOp>(loc, vectorType, nan);
rewriter.replaceOpWithNewOp<SelectOp>(op, isNaN, nan, select);
return success();
}
};
template <typename OpTy, arith::CmpIPredicate pred>
struct MaxMinIOpConverter : public OpRewritePattern<OpTy> {
public:
using OpRewritePattern<OpTy>::OpRewritePattern;
LogicalResult matchAndRewrite(OpTy op,
PatternRewriter &rewriter) const final {
Value lhs = op.getLhs();
Value rhs = op.getRhs();
Location loc = op.getLoc();
Value cmp = rewriter.create<arith::CmpIOp>(loc, pred, lhs, rhs);
rewriter.replaceOpWithNewOp<SelectOp>(op, cmp, lhs, rhs);
return success();
}
};
struct ArithmeticExpandOpsPass
: public ArithmeticExpandOpsBase<ArithmeticExpandOpsPass> {
void runOnFunction() override {
RewritePatternSet patterns(&getContext());
ConversionTarget target(getContext());
arith::populateArithmeticExpandOpsPatterns(patterns);
target.addLegalDialect<arith::ArithmeticDialect, StandardOpsDialect>();
// clang-format off
target.addIllegalOp<
arith::CeilDivSIOp,
arith::CeilDivUIOp,
arith::FloorDivSIOp,
arith::MaxFOp,
arith::MaxSIOp,
arith::MaxUIOp,
arith::MinFOp,
arith::MinSIOp,
arith::MinUIOp
>();
// clang-format on
if (failed(
applyPartialConversion(getFunction(), target, std::move(patterns))))
signalPassFailure();
}
};
} // end anonymous namespace
void mlir::arith::populateArithmeticExpandOpsPatterns(
RewritePatternSet &patterns) {
// clang-format off
patterns.add<
CeilDivSIOpConverter,
CeilDivUIOpConverter,
FloorDivSIOpConverter,
MaxMinFOpConverter<MaxFOp, arith::CmpFPredicate::OGT>,
MaxMinFOpConverter<MinFOp, arith::CmpFPredicate::OLT>,
MaxMinIOpConverter<MaxSIOp, arith::CmpIPredicate::sgt>,
MaxMinIOpConverter<MaxUIOp, arith::CmpIPredicate::ugt>,
MaxMinIOpConverter<MinSIOp, arith::CmpIPredicate::slt>,
MaxMinIOpConverter<MinUIOp, arith::CmpIPredicate::ult>
>(patterns.getContext());
// clang-format on
}
std::unique_ptr<Pass> mlir::arith::createArithmeticExpandOpsPass() {
return std::make_unique<ArithmeticExpandOpsPass>();
}