mirror of
https://github.com/llvm/llvm-project.git
synced 2025-04-25 07:56:06 +00:00
3fe8ec7906
This reverts commit 7c349c369847dc2f1736efb9c90d03521cd44a90. Per discussion at https://reviews.llvm.org/rG7c349c369847dc2f1736efb9c90d03521cd44a90 and elsewhere, the lowering to LLVM defined here isn't what it should be and the fastmath flag usage isn't correct, so `arith.is_nan` and `arith.is_inf` cannot exist in their current form. It's unclear if those operations should be introduced in the future, since they make the dialect more complex and don't add any expressive power. Further discussion may be moved to an RFC (or I'll drop this patch). Differential Revision: https://reviews.llvm.org/D157543
384 lines
15 KiB
C++
384 lines
15 KiB
C++
//===- ExpandOps.cpp - Pass to legalize Arith 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 "mlir/Dialect/Arith/Transforms/Passes.h"
|
|
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
|
#include "mlir/IR/ImplicitLocOpBuilder.h"
|
|
#include "mlir/IR/TypeUtilities.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
namespace mlir {
|
|
namespace arith {
|
|
#define GEN_PASS_DEF_ARITHEXPANDOPS
|
|
#include "mlir/Dialect/Arith/Transforms/Passes.h.inc"
|
|
} // namespace arith
|
|
} // namespace mlir
|
|
|
|
using namespace mlir;
|
|
|
|
/// Create an integer or index constant.
|
|
static Value createConst(Location loc, Type type, int value,
|
|
PatternRewriter &rewriter) {
|
|
auto attr = rewriter.getIntegerAttr(getElementTypeOrSelf(type), value);
|
|
if (auto shapedTy = dyn_cast<ShapedType>(type)) {
|
|
return rewriter.create<arith::ConstantOp>(
|
|
loc, DenseElementsAttr::get(shapedTy, attr));
|
|
}
|
|
|
|
return rewriter.create<arith::ConstantOp>(loc, attr);
|
|
}
|
|
|
|
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 = createConst(loc, a.getType(), 0, rewriter);
|
|
Value compare =
|
|
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, a, zero);
|
|
Value one = createConst(loc, a.getType(), 1, rewriter);
|
|
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);
|
|
rewriter.replaceOpWithNewOp<arith::SelectOp>(op, compare, zero, plusOne);
|
|
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();
|
|
Type type = op.getType();
|
|
Value a = op.getLhs();
|
|
Value b = op.getRhs();
|
|
Value plusOne = createConst(loc, type, 1, rewriter);
|
|
Value zero = createConst(loc, type, 0, rewriter);
|
|
Value minusOne = createConst(loc, type, -1, rewriter);
|
|
// Compute x = (b>0) ? -1 : 1.
|
|
Value compare =
|
|
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, b, zero);
|
|
Value x = rewriter.create<arith::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);
|
|
// Perform substitution and return success.
|
|
rewriter.replaceOpWithNewOp<arith::SelectOp>(op, compareRes, posRes,
|
|
negRes);
|
|
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();
|
|
Type type = op.getType();
|
|
Value a = op.getLhs();
|
|
Value b = op.getRhs();
|
|
Value plusOne = createConst(loc, type, 1, rewriter);
|
|
Value zero = createConst(loc, type, 0, rewriter);
|
|
Value minusOne = createConst(loc, type, -1, rewriter);
|
|
// Compute x = (b<0) ? 1 : -1.
|
|
Value compare =
|
|
rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, b, zero);
|
|
Value x = rewriter.create<arith::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);
|
|
// Perform substitution and return success.
|
|
rewriter.replaceOpWithNewOp<arith::SelectOp>(op, compareRes, negRes,
|
|
posRes);
|
|
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();
|
|
// If any operand is NaN, 'cmp' will be true (and 'select' returns 'lhs').
|
|
static_assert(pred == arith::CmpFPredicate::UGT ||
|
|
pred == arith::CmpFPredicate::ULT,
|
|
"pred must be either UGT or ULT");
|
|
Value cmp = rewriter.create<arith::CmpFOp>(loc, pred, lhs, rhs);
|
|
Value select = rewriter.create<arith::SelectOp>(loc, cmp, lhs, rhs);
|
|
|
|
// Handle the case where rhs is NaN: 'isNaN(rhs) ? rhs : select'.
|
|
Value isNaN = rewriter.create<arith::CmpFOp>(loc, arith::CmpFPredicate::UNO,
|
|
rhs, rhs);
|
|
rewriter.replaceOpWithNewOp<arith::SelectOp>(op, isNaN, rhs, select);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
struct BFloat16ExtFOpConverter : public OpRewritePattern<arith::ExtFOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
LogicalResult matchAndRewrite(arith::ExtFOp op,
|
|
PatternRewriter &rewriter) const final {
|
|
ImplicitLocOpBuilder b(op.getLoc(), rewriter);
|
|
auto operand = op.getOperand();
|
|
Type operandTy = operand.getType();
|
|
Type resultTy = op.getType();
|
|
Type operandETy = getElementTypeOrSelf(operandTy);
|
|
Type resultETy = getElementTypeOrSelf(resultTy);
|
|
|
|
if (!operandETy.isBF16() || !resultETy.isF32()) {
|
|
return rewriter.notifyMatchFailure(op, "not a ext of bf16 to f32.");
|
|
}
|
|
|
|
Type i16Ty = b.getI16Type();
|
|
Type i32Ty = b.getI32Type();
|
|
if (auto shapedTy = dyn_cast<ShapedType>(operandTy)) {
|
|
i16Ty = shapedTy.clone(i16Ty);
|
|
i32Ty = shapedTy.clone(i32Ty);
|
|
}
|
|
|
|
Value bitcast = b.create<arith::BitcastOp>(i16Ty, operand);
|
|
Value exti = b.create<arith::ExtUIOp>(i32Ty, bitcast);
|
|
|
|
Value c16 = createConst(op.getLoc(), i32Ty, 16, rewriter);
|
|
Value shl = b.create<arith::ShLIOp>(exti, c16);
|
|
Value result = b.create<arith::BitcastOp>(resultTy, shl);
|
|
|
|
rewriter.replaceOp(op, result);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
struct BFloat16TruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
LogicalResult matchAndRewrite(arith::TruncFOp op,
|
|
PatternRewriter &rewriter) const final {
|
|
ImplicitLocOpBuilder b(op.getLoc(), rewriter);
|
|
auto operand = op.getOperand();
|
|
Type operandTy = operand.getType();
|
|
Type resultTy = op.getType();
|
|
Type operandETy = getElementTypeOrSelf(operandTy);
|
|
Type resultETy = getElementTypeOrSelf(resultTy);
|
|
|
|
if (!operandETy.isF32() || !resultETy.isBF16()) {
|
|
return rewriter.notifyMatchFailure(op, "not a trunc of f32 to bf16.");
|
|
}
|
|
|
|
Type i1Ty = b.getI1Type();
|
|
Type i16Ty = b.getI16Type();
|
|
Type i32Ty = b.getI32Type();
|
|
Type f32Ty = b.getF32Type();
|
|
if (auto shapedTy = dyn_cast<ShapedType>(operandTy)) {
|
|
i1Ty = shapedTy.clone(i1Ty);
|
|
i16Ty = shapedTy.clone(i16Ty);
|
|
i32Ty = shapedTy.clone(i32Ty);
|
|
f32Ty = shapedTy.clone(f32Ty);
|
|
}
|
|
|
|
Value bitcast = b.create<arith::BitcastOp>(i32Ty, operand);
|
|
|
|
Value c23 = createConst(op.getLoc(), i32Ty, 23, rewriter);
|
|
Value c31 = createConst(op.getLoc(), i32Ty, 31, rewriter);
|
|
Value c23Mask = createConst(op.getLoc(), i32Ty, (1 << 23) - 1, rewriter);
|
|
Value expMask =
|
|
createConst(op.getLoc(), i32Ty, ((1 << 8) - 1) << 23, rewriter);
|
|
Value expMax =
|
|
createConst(op.getLoc(), i32Ty, ((1 << 8) - 2) << 23, rewriter);
|
|
|
|
// Grab the sign bit.
|
|
Value sign = b.create<arith::ShRUIOp>(bitcast, c31);
|
|
|
|
// Our mantissa rounding value depends on the sign bit and the last
|
|
// truncated bit.
|
|
Value cManRound = createConst(op.getLoc(), i32Ty, (1 << 15), rewriter);
|
|
cManRound = b.create<arith::SubIOp>(cManRound, sign);
|
|
|
|
// Grab out the mantissa and directly apply rounding.
|
|
Value man = b.create<arith::AndIOp>(bitcast, c23Mask);
|
|
Value manRound = b.create<arith::AddIOp>(man, cManRound);
|
|
|
|
// Grab the overflow bit and shift right if we overflow.
|
|
Value roundBit = b.create<arith::ShRUIOp>(manRound, c23);
|
|
Value manNew = b.create<arith::ShRUIOp>(manRound, roundBit);
|
|
|
|
// Grab the exponent and round using the mantissa's carry bit.
|
|
Value exp = b.create<arith::AndIOp>(bitcast, expMask);
|
|
Value expCarry = b.create<arith::AddIOp>(exp, manRound);
|
|
expCarry = b.create<arith::AndIOp>(expCarry, expMask);
|
|
|
|
// If the exponent is saturated, we keep the max value.
|
|
Value expCmp =
|
|
b.create<arith::CmpIOp>(arith::CmpIPredicate::uge, exp, expMax);
|
|
exp = b.create<arith::SelectOp>(expCmp, exp, expCarry);
|
|
|
|
// If the exponent is max and we rolled over, keep the old mantissa.
|
|
Value roundBitBool = b.create<arith::TruncIOp>(i1Ty, roundBit);
|
|
Value keepOldMan = b.create<arith::AndIOp>(expCmp, roundBitBool);
|
|
man = b.create<arith::SelectOp>(keepOldMan, man, manNew);
|
|
|
|
// Assemble the now rounded f32 value (as an i32).
|
|
Value rounded = b.create<arith::ShLIOp>(sign, c31);
|
|
rounded = b.create<arith::OrIOp>(rounded, exp);
|
|
rounded = b.create<arith::OrIOp>(rounded, man);
|
|
|
|
Value c16 = createConst(op.getLoc(), i32Ty, 16, rewriter);
|
|
Value shr = b.create<arith::ShRUIOp>(rounded, c16);
|
|
Value trunc = b.create<arith::TruncIOp>(i16Ty, shr);
|
|
Value result = b.create<arith::BitcastOp>(resultTy, trunc);
|
|
|
|
rewriter.replaceOp(op, result);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
struct ArithExpandOpsPass
|
|
: public arith::impl::ArithExpandOpsBase<ArithExpandOpsPass> {
|
|
ArithExpandOpsPass() = default;
|
|
ArithExpandOpsPass(const arith::ArithExpandOpsOptions& options) {
|
|
this->includeBf16 = options.includeBf16;
|
|
}
|
|
|
|
void runOnOperation() override {
|
|
RewritePatternSet patterns(&getContext());
|
|
ConversionTarget target(getContext());
|
|
|
|
arith::populateArithExpandOpsPatterns(patterns);
|
|
|
|
target.addLegalDialect<arith::ArithDialect>();
|
|
// clang-format off
|
|
target.addIllegalOp<
|
|
arith::CeilDivSIOp,
|
|
arith::CeilDivUIOp,
|
|
arith::FloorDivSIOp,
|
|
arith::MaxFOp,
|
|
arith::MinFOp
|
|
>();
|
|
|
|
if (includeBf16) {
|
|
arith::populateExpandBFloat16Patterns(patterns);
|
|
target.addDynamicallyLegalOp<arith::ExtFOp>(
|
|
[](arith::ExtFOp op) {
|
|
Type inETy = getElementTypeOrSelf(op.getOperand().getType());
|
|
Type outETy = getElementTypeOrSelf(op.getType());
|
|
return !(inETy.isBF16() && outETy.isF32());
|
|
});
|
|
|
|
target.addDynamicallyLegalOp<arith::TruncFOp>(
|
|
[](arith::TruncFOp op) {
|
|
Type inETy = getElementTypeOrSelf(op.getOperand().getType());
|
|
Type outETy = getElementTypeOrSelf(op.getType());
|
|
return !(inETy.isF32() && outETy.isBF16());
|
|
});
|
|
}
|
|
|
|
// clang-format on
|
|
if (failed(applyPartialConversion(getOperation(), target,
|
|
std::move(patterns))))
|
|
signalPassFailure();
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
void mlir::arith::populateCeilFloorDivExpandOpsPatterns(
|
|
RewritePatternSet &patterns) {
|
|
patterns
|
|
.add<CeilDivSIOpConverter, CeilDivUIOpConverter, FloorDivSIOpConverter>(
|
|
patterns.getContext());
|
|
}
|
|
|
|
void mlir::arith::populateExpandBFloat16Patterns(RewritePatternSet &patterns) {
|
|
patterns.add<BFloat16ExtFOpConverter, BFloat16TruncFOpConverter>(
|
|
patterns.getContext());
|
|
}
|
|
|
|
void mlir::arith::populateArithExpandOpsPatterns(RewritePatternSet &patterns) {
|
|
populateCeilFloorDivExpandOpsPatterns(patterns);
|
|
// clang-format off
|
|
patterns.add<
|
|
MaxMinFOpConverter<MaxFOp, arith::CmpFPredicate::UGT>,
|
|
MaxMinFOpConverter<MinFOp, arith::CmpFPredicate::ULT>
|
|
>(patterns.getContext());
|
|
// clang-format on
|
|
}
|
|
|
|
std::unique_ptr<Pass> mlir::arith::createArithExpandOpsPass() {
|
|
return std::make_unique<ArithExpandOpsPass>();
|
|
}
|
|
|
|
std::unique_ptr<Pass> mlir::arith::createArithExpandOpsPass(
|
|
const ArithExpandOpsOptions& options) {
|
|
return std::make_unique<ArithExpandOpsPass>(options);
|
|
}
|