Google Git Google Git
Sign in
opensecura/3p/openxla/iree/bc23e5959d089898715aef215659468ef6b3dd49/./compiler/plugins/input/StableHLO/Conversion/MapStableHLOToScalarOp.h
blob: 54b8e27ac3091bc2dd2ddf6c629ea624a55c971b [file]
// Copyright 2019 The IREE Authors
//
// Licensed 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
#ifndef IREE_COMPILER_PLUGINS_INPUT_STABLEHLO_CONVERSION_MAP_STABLEHLO_TO_SCALAR_OP_H
#define IREE_COMPILER_PLUGINS_INPUT_STABLEHLO_CONVERSION_MAP_STABLEHLO_TO_SCALAR_OP_H
#include <optional>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSwitch.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/IR/TypeUtilities.h"
#include "stablehlo/dialect/StablehloOps.h"
namespace mlir {
namespace stablehlo {
namespace impl {
// A struct to map StableHloBinaryOpTy type to the corresponding floating-point
// and integer scalar operation types.
template <typename StableHloBinaryOpTy>
struct StableHloToScalarOp {
using FOp = void;
using IOp = void;
using UOp = void;
using COp = void;
};
template <>
struct StableHloToScalarOp<stablehlo::AddOp> {
using FOp = ::mlir::arith::AddFOp;
using IOp = ::mlir::arith::AddIOp;
using UOp = ::mlir::arith::AddIOp;
using COp = ::mlir::complex::AddOp;
};
template <>
struct StableHloToScalarOp<stablehlo::AndOp> {
using IOp = ::mlir::arith::AndIOp;
using UOp = ::mlir::arith::AndIOp;
};
template <>
struct StableHloToScalarOp<stablehlo::CbrtOp> {
using FOp = ::mlir::math::CbrtOp;
};
template <>
struct StableHloToScalarOp<stablehlo::CompareOp> {
using FOp = ::mlir::arith::CmpFOp;
using IOp = ::mlir::arith::CmpIOp;
using UOp = ::mlir::arith::CmpIOp;
};
template <>
struct StableHloToScalarOp<stablehlo::CeilOp> {
using FOp = ::mlir::math::CeilOp;
};
template <>
struct StableHloToScalarOp<stablehlo::ClzOp> {
using IOp = ::mlir::math::CountLeadingZerosOp;
using UOp = ::mlir::math::CountLeadingZerosOp;
};
template <>
struct StableHloToScalarOp<stablehlo::CosineOp> {
using FOp = ::mlir::math::CosOp;
using COp = ::mlir::complex::CosOp;
};
template <>
struct StableHloToScalarOp<stablehlo::ExpOp> {
using FOp = ::mlir::math::ExpOp;
using COp = ::mlir::complex::ExpOp;
};
template <>
struct StableHloToScalarOp<stablehlo::Expm1Op> {
using FOp = ::mlir::math::ExpM1Op;
using COp = ::mlir::complex::Expm1Op;
};
template <>
struct StableHloToScalarOp<stablehlo::FloorOp> {
using FOp = ::mlir::math::FloorOp;
};
template <>
struct StableHloToScalarOp<stablehlo::LogOp> {
using FOp = ::mlir::math::LogOp;
using COp = ::mlir::complex::LogOp;
};
template <>
struct StableHloToScalarOp<stablehlo::Log1pOp> {
using FOp = ::mlir::math::Log1pOp;
using COp = ::mlir::complex::Log1pOp;
};
template <>
struct StableHloToScalarOp<stablehlo::MulOp> {
using FOp = ::mlir::arith::MulFOp;
using IOp = ::mlir::arith::MulIOp;
using UOp = ::mlir::arith::MulIOp;
using COp = ::mlir::complex::MulOp;
};
template <>
struct StableHloToScalarOp<stablehlo::OrOp> {
using IOp = ::mlir::arith::OrIOp;
using UOp = ::mlir::arith::OrIOp;
};
template <>
struct StableHloToScalarOp<stablehlo::PopulationCountOp> {
using IOp = ::mlir::math::CtPopOp;
using UOp = ::mlir::math::CtPopOp;
};
template <>
struct StableHloToScalarOp<stablehlo::RsqrtOp> {
using FOp = ::mlir::math::RsqrtOp;
using COp = ::mlir::complex::RsqrtOp;
};
template <>
struct StableHloToScalarOp<stablehlo::RoundNearestEvenOp> {
using FOp = ::mlir::math::RoundEvenOp;
};
template <>
struct StableHloToScalarOp<stablehlo::RoundOp> {
using FOp = ::mlir::math::RoundOp;
};
template <>
struct StableHloToScalarOp<stablehlo::SubtractOp> {
using FOp = ::mlir::arith::SubFOp;
using IOp = ::mlir::arith::SubIOp;
using UOp = ::mlir::arith::SubIOp;
using COp = ::mlir::complex::SubOp;
};
template <>
struct StableHloToScalarOp<stablehlo::SqrtOp> {
using FOp = ::mlir::math::SqrtOp;
using COp = ::mlir::complex::SqrtOp;
};
template <>
struct StableHloToScalarOp<stablehlo::SineOp> {
using FOp = ::mlir::math::SinOp;
using COp = ::mlir::complex::SinOp;
};
// FIXME(Jakub)
/*
template <>
struct StableHloToScalarOp<stablehlo::TanOp> {
using FOp = ::mlir::math::TanOp;
using COp = ::mlir::complex::TanOp;
};
*/
template <>
struct StableHloToScalarOp<stablehlo::Atan2Op> {
using FOp = ::mlir::math::Atan2Op;
using COp = ::mlir::complex::Atan2Op;
};
template <>
struct StableHloToScalarOp<stablehlo::TanhOp> {
using FOp = ::mlir::math::TanhOp;
using COp = ::mlir::complex::TanhOp;
};
template <>
struct StableHloToScalarOp<stablehlo::XorOp> {
using IOp = ::mlir::arith::XOrIOp;
using UOp = ::mlir::arith::XOrIOp;
};
// Alias for the map from StableHLO binary op type to STD floating-point op
// type.
template <typename StableHloOp>
using ScalarFOp = typename StableHloToScalarOp<StableHloOp>::FOp;
// Alias for the map from StableHLO binary op type to STD signed integer op
// type.
template <typename StableHloOp>
using ScalarIOp = typename StableHloToScalarOp<StableHloOp>::IOp;
// Alias for the map from StableHLO binary op type to STD unsigned integer op
// type.
template <typename StableHloOp>
using ScalarUOp = typename StableHloToScalarOp<StableHloOp>::UOp;
// Alias for the map from StableHLO binary op type to STD complex op type.
template <typename StableHloOp>
using ScalarCOp = typename StableHloToScalarOp<StableHloOp>::COp;
template <typename... Args>
struct MapStableHloOpToScalarOpImpl {
Value operator()(Location /*loc*/, ArrayRef<Type> /*ResultTypes*/,
ArrayRef<Type> /*argTypes*/, ValueRange /*args*/,
OpBuilder * /*b*/) {
return nullptr;
}
};
template <typename StdScalarOp>
struct MapStableHloOpToScalarOpImpl<StdScalarOp> {
Value operator()(Location loc, ArrayRef<Type> resultTypes,
ArrayRef<Type> /*argTypes*/, ValueRange args, OpBuilder *b) {
return b->template create<StdScalarOp>(loc, resultTypes, args,
std::nullopt);
}
};
template <typename SupportedType, typename StdScalarOp, typename... Args>
struct MapStableHloOpToScalarOpImpl<SupportedType, StdScalarOp, Args...> {
Value operator()(Location loc, ArrayRef<Type> resultTypes,
ArrayRef<Type> argTypes, ValueRange args, OpBuilder *b) {
Type elementType = getElementTypeOrSelf(argTypes.front());
if (SupportedType{}(elementType)) {
return b->template create<StdScalarOp>(loc, resultTypes, args,
std::nullopt);
}
return MapStableHloOpToScalarOpImpl<Args...>{}(loc, resultTypes, argTypes,
args, b);
}
};
template <typename SupportedType, typename... Args>
struct MapStableHloOpToScalarOpImpl<SupportedType, void, Args...> {
Value operator()(Location loc, ArrayRef<Type> resultTypes,
ArrayRef<Type> argTypes, ValueRange args, OpBuilder *b) {
return MapStableHloOpToScalarOpImpl<Args...>{}(loc, resultTypes, argTypes,
args, b);
}
};
struct IsAnyIntegerType {
bool operator()(Type t) { return isa<IntegerType>(t); }
};
struct IsSignedIntegerType {
bool operator()(Type t) {
// Pretend that signless is signed. This will change eventually.
return isa<IntegerType>(t) && !t.isUnsignedInteger() &&
!t.isSignlessInteger(1);
}
};
struct IsUnsignedIntegerType {
bool operator()(Type t) {
return t.isUnsignedInteger() || t.isSignlessInteger(1);
}
};
struct IsFloatType {
bool operator()(Type t) { return isa<FloatType>(t); }
};
struct IsComplexType {
bool operator()(Type t) { return isa<ComplexType>(t); }
};
template <template <typename T> class MapTy, typename OpTy,
typename PredTy = llvm::is_detected<MapTy, OpTy>>
struct MapableIf {
using type = void;
};
template <template <typename T> class MapTy, typename OpTy>
struct MapableIf<MapTy, OpTy, std::true_type> {
using type = MapTy<OpTy>;
};
// Inserts the computation that corresponds to the body of the loop for lowered
// StableHLO unary/binary op. Returns the value for the result.
template <typename StableHloOpTy>
inline Value mapStableHloOpToStdScalarOp(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
typename StableHloOpTy::Adaptor adaptor, OpBuilder *b) {
using ScalarIOpOrVoid = typename MapableIf<ScalarIOp, StableHloOpTy>::type;
using ScalarUOpOrVoid = typename MapableIf<ScalarUOp, StableHloOpTy>::type;
using ScalarFOpOrVoid = typename MapableIf<ScalarFOp, StableHloOpTy>::type;
using ScalarCOpOrVoid = typename MapableIf<ScalarCOp, StableHloOpTy>::type;
return MapStableHloOpToScalarOpImpl<IsSignedIntegerType, ScalarIOpOrVoid,
IsUnsignedIntegerType, ScalarUOpOrVoid,
IsFloatType, ScalarFOpOrVoid,
IsComplexType, ScalarCOpOrVoid>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::AbsOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::AbsOp::Adaptor adaptor, OpBuilder *b) {
Type elementType = getElementTypeOrSelf(argTypes.front());
if (isa<FloatType>(elementType)) {
return MapStableHloOpToScalarOpImpl<IsFloatType, ::mlir::math::AbsFOp>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
if (isa<ComplexType>(elementType)) {
return MapStableHloOpToScalarOpImpl<IsComplexType,
::mlir::complex::AbsOp>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
if (elementType.isSignlessInteger() || elementType.isSignedInteger()) {
// lmhlo.abs(x, result) -> result = select((x > 0), x, sub(0, x))
Value lhs = adaptor.getOperand();
Value zeroIntval =
b->create<arith::ConstantOp>(loc, b->getZeroAttr(lhs.getType()));
auto lhsGtZero = b->create<ScalarIOp<CompareOp>>(
loc, arith::CmpIPredicate::sge, lhs, zeroIntval);
auto negVal =
b->create<ScalarIOp<stablehlo::SubtractOp>>(loc, zeroIntval, lhs);
return b->create<::mlir::arith::SelectOp>(loc, lhsGtZero, lhs, negVal);
}
return nullptr;
}
// Return a constant for v of type t, splat if t is a vector type.
inline Value getConstantOrSplat(OpBuilder *b, Location loc, Type t,
Attribute v) {
if (VectorType vecType = dyn_cast<VectorType>(t)) {
v = SplatElementsAttr::get(vecType, v);
}
return b->create<arith::ConstantOp>(loc, t, cast<TypedAttr>(v));
}
template <typename PredicateType>
inline std::optional<PredicateType>
getCmpPredicate(stablehlo::ComparisonDirection, bool) {
return std::nullopt;
}
template <>
inline std::optional<arith::CmpFPredicate>
getCmpPredicate<arith::CmpFPredicate>(
stablehlo::ComparisonDirection comparisonDirection, bool isSigned) {
assert(isSigned && "cannot have an unsigned float!");
return llvm::StringSwitch<std::optional<arith::CmpFPredicate>>(
stringifyComparisonDirection(comparisonDirection))
.Case("EQ", arith::CmpFPredicate::OEQ)
.Case("NE", arith::CmpFPredicate::UNE)
.Case("GE", arith::CmpFPredicate::OGE)
.Case("GT", arith::CmpFPredicate::OGT)
.Case("LE", arith::CmpFPredicate::OLE)
.Case("LT", arith::CmpFPredicate::OLT)
.Default(std::nullopt);
}
template <>
inline std::optional<arith::CmpIPredicate>
getCmpPredicate<arith::CmpIPredicate>(
stablehlo::ComparisonDirection comparisonDirection, bool isSigned) {
return llvm::StringSwitch<std::optional<arith::CmpIPredicate>>(
stringifyComparisonDirection(comparisonDirection))
.Case("EQ", arith::CmpIPredicate::eq)
.Case("NE", arith::CmpIPredicate::ne)
.Case("GE",
isSigned ? arith::CmpIPredicate::sge : arith::CmpIPredicate::uge)
.Case("GT",
isSigned ? arith::CmpIPredicate::sgt : arith::CmpIPredicate::ugt)
.Case("LE",
isSigned ? arith::CmpIPredicate::sle : arith::CmpIPredicate::ule)
.Case("LT",
isSigned ? arith::CmpIPredicate::slt : arith::CmpIPredicate::ult)
.Default(std::nullopt);
}
inline Value cmpComplex(Location loc, Value lhs, Value rhs,
stablehlo::ComparisonDirection comparisonDirection,
OpBuilder *b) {
auto complexType = cast<ComplexType>(lhs.getType());
if (isa<FloatType>(complexType.getElementType())) {
if (comparisonDirection == stablehlo::ComparisonDirection::EQ) {
return b->create<complex::EqualOp>(loc, lhs, rhs);
}
if (comparisonDirection == stablehlo::ComparisonDirection::NE) {
return b->create<complex::NotEqualOp>(loc, lhs, rhs);
}
// Perform a lexicographical comparison for the (real, imaginary) pair.
Type complexFloatTy = complexType.getElementType();
Value lhsReal = b->create<complex::ReOp>(loc, complexFloatTy, lhs);
Value rhsReal = b->create<complex::ReOp>(loc, complexFloatTy, rhs);
Value lhsImag = b->create<complex::ImOp>(loc, complexFloatTy, lhs);
Value rhsImag = b->create<complex::ImOp>(loc, complexFloatTy, rhs);
auto predicate = getCmpPredicate<arith::CmpFPredicate>(comparisonDirection,
/*is_signed=*/true);
assert(predicate.has_value() && "expected valid comparison direction");
// if (lhsReal == rhsReal && lhsImag `predicate` rhsImag ||
// lhsReal `predicate` rhsReal)
Value realsAreEq = b->create<arith::CmpFOp>(loc, arith::CmpFPredicate::OEQ,
lhsReal, rhsReal);
Value imagsAreOrdered =
b->create<arith::CmpFOp>(loc, *predicate, lhsImag, rhsImag);
Value realsAreOrdered =
b->create<arith::CmpFOp>(loc, *predicate, lhsReal, rhsReal);
Value orLhs = b->create<arith::AndIOp>(loc, realsAreEq, imagsAreOrdered);
return b->create<arith::OrIOp>(loc, orLhs, realsAreOrdered);
}
return nullptr;
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::CompareOp>(
Location loc, ArrayRef<Type> /*resultTypes*/, ArrayRef<Type> argTypes,
stablehlo::CompareOp::Adaptor adaptor, OpBuilder *b) {
stablehlo::ComparisonDirection comparisonDirection =
adaptor.getComparisonDirection();
const auto &lhs = adaptor.getLhs();
const auto &rhs = adaptor.getRhs();
Type elementType = getElementTypeOrSelf(argTypes.front());
if (isa<IntegerType>(elementType)) {
bool isUnsigned = IsUnsignedIntegerType{}(elementType);
std::optional<arith::CmpIPredicate> predicate =
getCmpPredicate<arith::CmpIPredicate>(comparisonDirection, !isUnsigned);
assert(predicate.has_value() && "expected valid comparison direction");
return b->create<ScalarIOp<stablehlo::CompareOp>>(loc, predicate.value(),
lhs, rhs);
}
if (auto floatType = dyn_cast<FloatType>(elementType)) {
if (adaptor.getCompareType() &&
*adaptor.getCompareType() == stablehlo::ComparisonType::TOTALORDER) {
// The semantics of totalorder fp compare are
// -NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN
auto intType = b->getIntegerType(floatType.getWidth());
auto zero =
b->create<arith::ConstantOp>(loc, intType, b->getZeroAttr(intType));
auto max = b->create<arith::ConstantOp>(
loc, intType,
b->getIntegerAttr(intType,
APInt::getSignedMaxValue(floatType.getWidth())));
// Switch from a floating point value to a integer value in such a way
// that when using the integer value to compare, we get the same result
// for normal values, and -NaN is treated as the smallest value, and NaN
// is treated as the largest value.
// If f is a float, and
// x = bit_cast<int32_t>(f);
// y = x < 0 ? numeric_limits<int32_t>::max() - x : x;
// then y is ordered as an int32_t such that finite values have the
// obvious order, -0 is ordered before 0, and -NaN and NaN appear at the
// beginning and end of the ordering.
auto toIntegral = [&](Value v) {
auto x = b->create<arith::BitcastOp>(loc, intType, v);
auto cmp =
b->create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, x, zero);
auto sub = b->create<arith::SubIOp>(loc, max, x);
return b->create<arith::SelectOp>(loc, cmp, sub, x);
};
auto lhsInt = toIntegral(lhs);
auto rhsInt = toIntegral(rhs);
auto predicate =
getCmpPredicate<arith::CmpIPredicate>(comparisonDirection,
/*is_signed=*/true);
assert(predicate.has_value() && "expected valid comparison direction");
return b->create<arith::CmpIOp>(loc, *predicate, lhsInt, rhsInt);
}
std::optional<arith::CmpFPredicate> predicate =
getCmpPredicate<arith::CmpFPredicate>(comparisonDirection,
/*is_signed=*/true);
assert(predicate.has_value() && "expected valid comparison direction");
return b->create<ScalarFOp<stablehlo::CompareOp>>(loc, predicate.value(),
lhs, rhs);
}
if (auto complexType = dyn_cast<ComplexType>(elementType))
return cmpComplex(loc, lhs, rhs, comparisonDirection, b);
return nullptr;
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::ReducePrecisionOp>(
Location loc, ArrayRef<Type> /*resultTypes*/, ArrayRef<Type> argTypes,
stablehlo::ReducePrecisionOp::Adaptor adaptor, OpBuilder *builder) {
using llvm::APInt;
mlir::ImplicitLocOpBuilder b(loc, *builder);
// Integer and float types for casting and constant generation.
auto floatType =
cast<FloatType>(cast<TensorType>(argTypes.front()).getElementType());
int64_t nbits = floatType.getWidth();
auto intType = mlir::IntegerType::get(loc.getContext(), floatType.getWidth());
Value xAsInt = b.create<arith::BitcastOp>(intType, adaptor.getOperand());
// SignificandWidth includes the implicit extra bit.
auto srcMantissaBits = floatType.getFPMantissaWidth() - 1;
int srcExponentBits = nbits - 1 - srcMantissaBits;
// Clear the sign bit, it does not participate in rounding and we will restore
// it later.
APInt signBitMask(nbits, 1);
signBitMask <<= nbits - 1;
APInt expBitsMask(nbits, 1);
expBitsMask = ((expBitsMask << srcExponentBits) - 1) << srcMantissaBits;
auto createConstant = [&](const APInt &v) {
return b.create<arith::ConstantIntOp>(v.getZExtValue(), intType)
.getResult();
};
Value xAbsBits =
b.create<arith::AndIOp>(xAsInt, createConstant(~signBitMask));
Value xIsNan = b.create<arith::CmpIOp>(arith::CmpIPredicate::ugt, xAbsBits,
createConstant(expBitsMask));
int destMantissaBits = adaptor.getMantissaBits();
if (destMantissaBits < static_cast<int>(srcMantissaBits)) {
// Last remaining mantissa bit.
APInt lastMantissaBitMask(nbits, 1);
lastMantissaBitMask <<= srcMantissaBits - destMantissaBits;
// Compute rounding bias for round-to-nearest with ties to even. This is
// equal to a base value of 0111... plus one bit if the last remaining
// mantissa bit is 1.
APInt baseRoundingBias = lastMantissaBitMask.lshr(1) - 1;
Value mantissaDiff = b.create<arith::ConstantIntOp>(
srcMantissaBits - destMantissaBits, intType);
Value highestMantissaMaskVal = createConstant(lastMantissaBitMask);
Value baseRoundingBiasVal = createConstant(baseRoundingBias);
Value xLastMantissaBit = b.create<arith::ShRUIOp>(
b.create<arith::AndIOp>(xAsInt, highestMantissaMaskVal), mantissaDiff);
Value xRoundingBias =
b.create<arith::AddIOp>(xLastMantissaBit, baseRoundingBiasVal);
// Add rounding bias, and mask out truncated bits. Note that the case
// where adding the rounding bias overflows into the exponent bits is
// correct; the non-masked mantissa bits will all be zero, and the
// exponent will be incremented by one.
APInt truncationMask = ~(lastMantissaBitMask - 1);
Value xRounded = b.create<arith::AddIOp>(xAsInt, xRoundingBias);
xAsInt = b.create<arith::AndIOp>(xRounded, createConstant(truncationMask));
}
int destExponentBits = adaptor.getExponentBits();
if (destExponentBits < srcExponentBits) {
// An exponent of 2^(n-1)-1 -- that is, 0111... with the zero in the most-
// significant bit -- is equal to 1.0f for all exponent sizes. Adding
// 2^(n-1)-1 to this gives us the highest non-infinite exponent for a bit-
// size of n, and subtracting 2^(n-1)-1 from this gives us the lowest'
// exponent (corresponding to 0.0f).
//
// Thus, the f32 exponent corresponding to the highest non-infinite
// exponent for a bit size of n is (2^7-1) + 2^(n-1)-1, and the f32
// exponent corresponding to the lowest exponent for a bit size of n is
// (2^7-1) - 2^(n-1)-1.
//
// Note that we have already checked that exponents_bits >= 1.
APInt exponentBias(nbits, 1);
exponentBias = (exponentBias << (srcExponentBits - 1)) - 1;
APInt reducedExponentBias(nbits, 1);
reducedExponentBias = (reducedExponentBias << (destExponentBits - 1)) - 1;
APInt reducedMaxExponent = exponentBias + reducedExponentBias;
APInt reducedMinExponent = exponentBias - reducedExponentBias;
// Do we overflow or underflow?
Value xExponent =
b.create<arith::AndIOp>(xAsInt, createConstant(expBitsMask));
Value xOverflows = b.create<arith::CmpIOp>(
arith::CmpIPredicate::ugt, xExponent,
createConstant(reducedMaxExponent << srcMantissaBits));
Value xUnderflows = b.create<arith::CmpIOp>(
arith::CmpIPredicate::ule, xExponent,
createConstant(reducedMinExponent << srcMantissaBits));
// Compute appropriately-signed values of zero and infinity.
Value xSignedZero =
b.create<arith::AndIOp>(xAsInt, createConstant(signBitMask));
Value xSignedInf =
b.create<arith::OrIOp>(xSignedZero, createConstant(expBitsMask));
// Force to zero or infinity if overflow or underflow. (Note that this
// truncates all denormal values to zero, rather than rounding them.)
xAsInt = b.create<arith::SelectOp>(xOverflows, xSignedInf, xAsInt);
xAsInt = b.create<arith::SelectOp>(xUnderflows, xSignedZero, xAsInt);
}
Value result = b.create<arith::BitcastOp>(floatType, xAsInt);
return b.create<arith::SelectOp>(xIsNan, adaptor.getOperand(), result);
}
// FIXME(Jakub)
// template <>
// inline Value mapStableHloOpToStdScalarOp<stablehlo::CopyOp>(
// Location /*loc*/, ArrayRef<Type> /*ResultTypes*/,
// ArrayRef<Type> /*argTypes*/, stablehlo::CopyOp::Adaptor adaptor,
// OpBuilder* /*b*/) {
// return adaptor.getOperands().front();
// }
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::ComplexOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::ComplexOp::Adaptor adaptor, OpBuilder *b) {
return MapStableHloOpToScalarOpImpl<complex::CreateOp>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::MaxOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::MaxOp::Adaptor adaptor, OpBuilder *b) {
ValueRange operands = adaptor.getOperands();
Value lhs = operands.front();
Type complexTy = lhs.getType();
if (!isa<ComplexType>(complexTy))
return MapStableHloOpToScalarOpImpl<
IsFloatType, arith::MaximumFOp, IsSignedIntegerType, arith::MaxSIOp,
IsUnsignedIntegerType, arith::MaxUIOp>{}(loc, resultTypes, argTypes,
adaptor.getOperands(), b);
assert(resultTypes.size() == 1 && "MaxOp should return a single result");
assert(operands.size() == 2 && "MaxOp should take exactly two arguments");
Value rhs = operands.back();
// 'max' performs a lexicographical comparison for the (real, imaginary) pair.
Value cond = cmpComplex(loc, lhs, rhs, stablehlo::ComparisonDirection::GE, b);
return b->create<arith::SelectOp>(loc, cond, lhs, rhs).getResult();
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::MinOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::MinOp::Adaptor adaptor, OpBuilder *b) {
ValueRange operands = adaptor.getOperands();
Value lhs = operands.front();
Type complexTy = lhs.getType();
if (!isa<ComplexType>(complexTy))
return MapStableHloOpToScalarOpImpl<
IsFloatType, arith::MinimumFOp, IsSignedIntegerType, arith::MinSIOp,
IsUnsignedIntegerType, arith::MinUIOp>{}(loc, resultTypes, argTypes,
adaptor.getOperands(), b);
assert(resultTypes.size() == 1 && "MinOp should return a single result");
assert(operands.size() == 2 && "MinOp should take exactly two arguments");
Value rhs = operands.back();
// 'min' performs a lexicographical comparison for the (real, imaginary) pair.
Value cond = cmpComplex(loc, lhs, rhs, stablehlo::ComparisonDirection::LE, b);
return b->create<arith::SelectOp>(loc, cond, lhs, rhs).getResult();
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::RealOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::RealOp::Adaptor adaptor, OpBuilder *b) {
if (!isa<ComplexType>(adaptor.getOperand().getType()))
return adaptor.getOperand();
return MapStableHloOpToScalarOpImpl<complex::ReOp>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::ImagOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::ImagOp::Adaptor adaptor, OpBuilder *b) {
if (!isa<ComplexType>(adaptor.getOperand().getType()))
return b->create<arith::ConstantOp>(
loc, b->getZeroAttr(adaptor.getOperand().getType()));
return MapStableHloOpToScalarOpImpl<complex::ImOp>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
// 'target_types' is the unconverted type (signed or unsigned if integer),
// 'ResultTypes' is the converted type (signless if integer).
inline Value mapConvertOpToStdScalarOp(Location loc, ArrayRef<Type> targetTypes,
ArrayRef<Type> resultTypes,
ArrayRef<Type> argTypes, ValueRange args,
OpBuilder *b) {
assert(targetTypes.size() == 1 && "ConvertOp should return a single result");
assert(resultTypes.size() == 1 && "ConvertOp should return a single result");
assert(argTypes.size() == 1 && "ConvertOp should take a single argument");
assert(args.size() == 1 && "ConvertOp should take a single argument");
Type sourceType = getElementTypeOrSelf(argTypes.front());
Type targetType = getElementTypeOrSelf(targetTypes.front());
Type convertedSourceType = getElementTypeOrSelf(args.front());
// A boolean value is considered to be unsigned when converting to
// floating-point. Otherwise, it will become `-1`.
if (IsUnsignedIntegerType{}(sourceType) &&
mlir::arith::UIToFPOp::areCastCompatible(convertedSourceType,
targetType)) {
return b->create<mlir::arith::UIToFPOp>(loc, resultTypes, args,
std::nullopt);
}
if (mlir::arith::SIToFPOp::areCastCompatible(sourceType, targetType)) {
return b->create<mlir::arith::SIToFPOp>(loc, resultTypes, args,
std::nullopt);
}
if (isa<FloatType>(sourceType) && isa<FloatType>(targetType)) {
auto src = cast<FloatType>(sourceType);
auto res = cast<FloatType>(targetType);
if (src.getWidth() > res.getWidth()) {
return b->create<mlir::arith::TruncFOp>(loc, resultTypes, args,
std::nullopt);
}
if (src.getWidth() < res.getWidth()) {
return b->create<mlir::arith::ExtFOp>(loc, resultTypes, args,
std::nullopt);
}
// There's no direct conversion between different 16 bit floating point
// types, so go through 32 bit float.
if (sourceType != targetType) {
assert(sourceType.isBF16() || targetType.isBF16());
Value ext = b->create<arith::ExtFOp>(loc, b->getF32Type(), args);
return b->create<arith::TruncFOp>(loc, resultTypes, ext);
}
// No conversion is needed for identical float types.
return args.front();
}
if (targetType.isInteger(/*width=*/1)) {
// When casting to bool, we need to compare whether the value is equal to
// zero.
if (sourceType.isSignlessInteger() || sourceType.isUnsignedInteger()) {
Value zeroIntval = b->create<arith::ConstantOp>(
loc, b->getZeroAttr(args.front().getType()));
return b->create<mlir::arith::CmpIOp>(loc, arith::CmpIPredicate::ne,
args.front(), zeroIntval);
}
if (isa<FloatType>(sourceType)) {
Value zero = b->create<arith::ConstantOp>(
loc, b->getZeroAttr(args.front().getType()));
return b->create<mlir::arith::CmpFOp>(loc, arith::CmpFPredicate::UNE,
args.front(), zero);
}
}
if (isa<IntegerType>(sourceType) && isa<IntegerType>(targetType)) {
auto src = cast<IntegerType>(sourceType);
auto res = cast<IntegerType>(targetType);
if (src.getWidth() > res.getWidth()) {
return b->create<mlir::arith::TruncIOp>(loc, resultTypes, args,
std::nullopt);
}
if (src.getWidth() < res.getWidth()) {
// Special case boolean values, so they get casted to `1` instead of `-1`.
if (IsUnsignedIntegerType{}(src)) {
return b->create<mlir::arith::ExtUIOp>(loc, resultTypes, args,
std::nullopt);
}
return b->create<mlir::arith::ExtSIOp>(loc, resultTypes, args,
std::nullopt);
}
// No conversion is needed for the same width integers
return args.front();
}
if (targetType.isUnsignedInteger() &&
mlir::arith::FPToUIOp::areCastCompatible(convertedSourceType,
targetType)) {
return b->create<mlir::arith::FPToUIOp>(loc, resultTypes, args,
std::nullopt);
}
if (mlir::arith::FPToSIOp::areCastCompatible(convertedSourceType,
targetType)) {
return b->create<mlir::arith::FPToSIOp>(loc, resultTypes, args,
std::nullopt);
}
if (isa<ComplexType>(targetType)) {
Type targetElementType = cast<ComplexType>(targetType).getElementType();
assert(!isa<ComplexType>(targetElementType) &&
"elements of complex numbers should not be complex");
Value targetReal;
Value targetImag;
if (isa<ComplexType>(sourceType)) {
// We are converting from complex type: convert real and imaginary parts
// separately.
Type sourceElementType = cast<ComplexType>(sourceType).getElementType();
assert(!isa<ComplexType>(sourceElementType) &&
"elements of complex numbers should not be complex");
Value sourceReal =
b->create<mlir::complex::ReOp>(loc, sourceElementType, args.front());
targetReal =
mapConvertOpToStdScalarOp(loc, targetElementType, targetElementType,
sourceElementType, sourceReal, b);
Value sourceImag =
b->create<mlir::complex::ImOp>(loc, sourceElementType, args.front());
targetImag =
mapConvertOpToStdScalarOp(loc, targetElementType, targetElementType,
sourceElementType, sourceImag, b);
} else {
// We are converting from real (float, integer, etc.) type, convert the
// real part and set the imaginary part to 0.
targetReal = mapConvertOpToStdScalarOp(
loc, targetElementType, targetElementType, argTypes, args, b);
targetImag = b->create<mlir::arith::ConstantOp>(
loc, b->getFloatAttr(targetElementType, 0.0));
}
return b->create<mlir::complex::CreateOp>(loc, targetType, targetReal,
targetImag);
}
if (auto sourceComplexType = dyn_cast<ComplexType>(sourceType)) {
auto sourceElementType = sourceComplexType.getElementType();
// When converting from complex to a non-complex type, we take just the real
// part of the complex number.
Value sourceReal =
b->create<mlir::complex::ReOp>(loc, sourceElementType, args.front());
return mapConvertOpToStdScalarOp(loc, targetTypes, resultTypes,
sourceElementType, sourceReal, b);
}
return nullptr;
}
/// Lower bitcast operations where the input and resulting type are the same
/// bitwidth, thus implying that the operation is fully defined by parallel
/// loops and scalar operations without any shape dimension changes.
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::BitcastConvertOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::BitcastConvertOp::Adaptor adaptor, OpBuilder *b) {
Type argType = getElementTypeOrSelf(argTypes.front());
Type resultType = getElementTypeOrSelf(resultTypes.front());
// Skip needless casts.
if (argType == resultType)
return adaptor.getOperand();
if (!isa<FloatType, IntegerType>(resultType) ||
!isa<FloatType, IntegerType>(argType))
return nullptr;
if (resultType.getIntOrFloatBitWidth() != argType.getIntOrFloatBitWidth())
return nullptr;
return b->create<mlir::arith::BitcastOp>(loc, resultTypes,
adaptor.getOperands());
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::DotOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::DotOp::Adaptor adaptor, OpBuilder *b) {
// Dot Op converter from lhlo to affine only accepts float and integer types.
const auto &lhs = adaptor.getOperands()[0];
const auto &rhs = adaptor.getOperands()[1];
const auto &result = adaptor.getOperands()[2];
Type elementType = lhs.getType();
if (isa<FloatType>(elementType)) {
Value floatMul =
MapStableHloOpToScalarOpImpl<IsFloatType, ::mlir::arith::MulFOp>{}(
loc, resultTypes, argTypes, {lhs, rhs}, b);
return MapStableHloOpToScalarOpImpl<IsFloatType, ::mlir::arith::AddFOp>{}(
loc, resultTypes, argTypes, {floatMul, result}, b);
}
if (isa<IntegerType>(elementType)) {
Value intMul =
MapStableHloOpToScalarOpImpl<IsAnyIntegerType, ::mlir::arith::MulIOp>{}(
loc, resultTypes, argTypes, {lhs, rhs}, b);
return MapStableHloOpToScalarOpImpl<IsAnyIntegerType,
::mlir::arith::AddIOp>{}(
loc, resultTypes, argTypes, {intMul, result}, b);
}
return nullptr;
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::IsFiniteOp>(
Location loc, ArrayRef<Type> /*ResultTypes*/, ArrayRef<Type> /*argTypes*/,
stablehlo::IsFiniteOp::Adaptor adaptor, OpBuilder *b) {
if (isa<FloatType>(adaptor.getX().getType())) {
auto posInf = APFloat::getInf(
cast<FloatType>(adaptor.getX().getType()).getFloatSemantics());
auto constPosInf = b->create<arith::ConstantOp>(
loc, b->getFloatAttr(adaptor.getX().getType(), posInf));
Value absX = b->create<::mlir::math::AbsFOp>(loc, adaptor.getX());
return b->create<::mlir::arith::CmpFOp>(loc, arith::CmpFPredicate::ONE,
absX, constPosInf);
}
return nullptr;
}
/// Implements the conversion of HLO op to scalar op (to use within region of a
/// linalg.generic op) for compare-select style operations like min/max.
template <typename... Args>
struct CompareSelectOpToStdScalarOp {
static Value map(Location /*loc*/,
stablehlo::ComparisonDirection /*comparison_direction*/,
ArrayRef<Type> /*ResultTypes*/, ArrayRef<Type> /*argTypes*/,
ValueRange /*args*/, OpBuilder * /*b*/) {
return nullptr;
}
};
/// Specialization which allows converting to a comparison operation in standard
/// dialect with a given predicate based on the element type of the operand.
template <typename SupportedType, typename StdCompareOp, typename Predicate,
typename... Args>
struct CompareSelectOpToStdScalarOp<SupportedType, StdCompareOp, Predicate,
Args...> {
static Value map(Location loc,
stablehlo::ComparisonDirection comparisonDirection,
ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
ValueRange args, OpBuilder *b) {
Type elementType = getElementTypeOrSelf(argTypes.front());
if (isa<SupportedType>(elementType)) {
auto predicate = getCmpPredicate<Predicate>(
comparisonDirection, !elementType.isUnsignedInteger());
assert(predicate.has_value() && "expected valid comparison direction");
auto cmp = b->template create<StdCompareOp>(loc, predicate.getValue(),
args[0], args[1]);
return b->create<::mlir::arith::SelectOp>(loc, cmp, args[0], args[1]);
}
return CompareSelectOpToStdScalarOp<Args...>::map(
loc, comparisonDirection, resultTypes, argTypes, args, b);
}
};
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::ClampOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::ClampOp::Adaptor op, OpBuilder *b) {
// clamp(lb, x, ub) = min(max(lb, x), ub)
Value maxLbX = mapStableHloOpToStdScalarOp<stablehlo::MaxOp>(
loc, resultTypes, argTypes, ValueRange{op.getMin(), op.getOperand()}, b);
return mapStableHloOpToStdScalarOp<stablehlo::MinOp>(
loc, resultTypes, argTypes, ValueRange{maxLbX, op.getMax()}, b);
}
template <typename U, typename S>
inline Value makeSafeIntDiv(ImplicitLocOpBuilder &lb, Type originalType,
Value lhs, Value rhs, Value returnedOnZero,
Value returnedOnSignedOverflow) {
Type type = lhs.getType();
auto elementType = cast<IntegerType>(getElementTypeOrSelf(type));
Value zero = lb.create<arith::ConstantOp>(lb.getZeroAttr(type));
auto makeConstant = [&](const APInt &i) {
return getConstantOrSplat(&lb, lb.getLoc(), type,
lb.getIntegerAttr(elementType, i));
};
Value one = makeConstant(APInt(elementType.getWidth(), 1));
Value rhsIsZero =
lb.create<arith::CmpIOp>(arith::CmpIPredicate::eq, rhs, zero);
// For unsigned just set the divisor to 1 when it would be 0.
if (originalType.isUnsignedInteger()) {
Value safeRhs = lb.create<arith::SelectOp>(rhsIsZero, one, rhs);
Value safeDiv = lb.create<U>(lhs, safeRhs);
return lb.create<arith::SelectOp>(rhsIsZero, returnedOnZero, safeDiv);
}
// For signed also check for INT_MIN / -1.
Value smin = makeConstant(APInt::getSignedMinValue(elementType.getWidth()));
Value lhsIsSmin =
lb.create<arith::CmpIOp>(arith::CmpIPredicate::eq, lhs, smin);
Value minusOne = makeConstant(APInt::getAllOnes(elementType.getWidth()));
Value rhsIsMinusOne =
lb.create<arith::CmpIOp>(arith::CmpIPredicate::eq, rhs, minusOne);
Value hasIntMinOverflow = lb.create<arith::AndIOp>(lhsIsSmin, rhsIsMinusOne);
Value rhsIsUnsafe = lb.create<arith::OrIOp>(rhsIsZero, hasIntMinOverflow);
Value safeRhs = lb.create<arith::SelectOp>(rhsIsUnsafe, one, rhs);
Value safeDiv = lb.create<S>(lhs, safeRhs);
Value safeSmin = lb.create<arith::SelectOp>(
hasIntMinOverflow, returnedOnSignedOverflow, safeDiv);
return lb.create<arith::SelectOp>(rhsIsZero, returnedOnZero, safeSmin);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::DivOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::DivOp::Adaptor adaptor, OpBuilder *b) {
Type originalType = getElementTypeOrSelf(argTypes.front());
if (isa<ComplexType, FloatType>(originalType)) {
return MapStableHloOpToScalarOpImpl<IsFloatType, arith::DivFOp,
IsComplexType, complex::DivOp>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
// Integer division overflow behavior:
//
// X / 0 == -1
// INT_SMIN /s -1 = INT_SMIN
ImplicitLocOpBuilder lb(loc, *b);
Type type = adaptor.getLhs().getType();
auto elementType = cast<IntegerType>(getElementTypeOrSelf(type));
auto makeConstant = [&](const APInt &i) {
return getConstantOrSplat(&lb, lb.getLoc(), type,
lb.getIntegerAttr(elementType, i));
};
Value minusOne = makeConstant(APInt::getAllOnes(elementType.getWidth()));
Value smin = makeConstant(APInt::getSignedMinValue(elementType.getWidth()));
return makeSafeIntDiv<arith::DivUIOp, arith::DivSIOp>(
lb, originalType, adaptor.getLhs(), adaptor.getRhs(),
/*returnedOnZero=*/minusOne,
/*returnedOnSignedOverflow=*/smin);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::RemOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::RemOp::Adaptor adaptor, OpBuilder *b) {
Type originalType = getElementTypeOrSelf(argTypes.front());
if (isa<ComplexType, FloatType>(originalType)) {
return MapStableHloOpToScalarOpImpl<IsFloatType, arith::RemFOp>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
// Integer remainder overflow behavior:
//
// X % 0 == X
// INT_SMIN %s -1 = 0
ImplicitLocOpBuilder lb(loc, *b);
Type type = adaptor.getLhs().getType();
Value zero = lb.create<arith::ConstantOp>(lb.getZeroAttr(type));
return makeSafeIntDiv<arith::RemUIOp, arith::RemSIOp>(
lb, originalType, adaptor.getLhs(), adaptor.getRhs(),
/*returnedOnZero=*/adaptor.getLhs(),
/*returnedOnSignedOverflow=*/zero);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::NegOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::NegOp::Adaptor adaptor, OpBuilder *b) {
Type elementType = getElementTypeOrSelf(adaptor.getOperand().getType());
if (isa<ComplexType, FloatType>(elementType)) {
return MapStableHloOpToScalarOpImpl<IsFloatType, ::mlir::arith::NegFOp,
IsComplexType,
::mlir::complex::NegOp>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
if (isa<IntegerType>(elementType)) {
// lmhlo.neg(x, result) -> result = sub(0, x)
Value lhs = adaptor.getOperand();
Value zeroIntval =
b->create<arith::ConstantOp>(loc, b->getZeroAttr(lhs.getType()));
return b->create<ScalarIOp<stablehlo::SubtractOp>>(loc, zeroIntval, lhs);
}
return nullptr;
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::NotOp>(
Location loc, ArrayRef<Type> /*ResultTypes*/, ArrayRef<Type> /*argTypes*/,
stablehlo::NotOp::Adaptor adaptor, OpBuilder *b) {
Type elementType = getElementTypeOrSelf(adaptor.getOperand().getType());
if (auto integerType = dyn_cast<IntegerType>(elementType)) {
// lmhlo.not(x) -> x ^ -1
Value allOnes = getConstantOrSplat(
b, loc, adaptor.getOperand().getType(),
b->getIntegerAttr(integerType,
APInt::getAllOnes(integerType.getWidth())));
return b->create<::mlir::arith::XOrIOp>(loc, allOnes, adaptor.getOperand());
}
return nullptr;
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::LogisticOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> /*argTypes*/,
stablehlo::LogisticOp::Adaptor adaptor, OpBuilder *b) {
// 1.0 / (1.0 - exp(-x))
Value negX = mapStableHloOpToStdScalarOp<stablehlo::NegOp>(
loc, resultTypes, resultTypes, {adaptor.getOperand()}, b);
Value expNegX = mapStableHloOpToStdScalarOp<stablehlo::ExpOp>(
loc, resultTypes, resultTypes, {{negX}}, b);
Value oneFloat = b->create<arith::ConstantOp>(loc, b->getF32FloatAttr(1.0));
Value one = mapConvertOpToStdScalarOp(loc, resultTypes, resultTypes,
{oneFloat.getType()}, {{oneFloat}}, b);
Value oneAddExprNegX = mapStableHloOpToStdScalarOp<stablehlo::AddOp>(
loc, resultTypes, resultTypes, {{expNegX, one}}, b);
return mapStableHloOpToStdScalarOp<stablehlo::DivOp>(
loc, resultTypes, resultTypes, {{one, oneAddExprNegX}}, b);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::PowOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::PowOp::Adaptor adaptor, OpBuilder *b) {
auto lb = ImplicitLocOpBuilder(loc, *b);
// Floating point can use std::powf
auto resultType = resultTypes.front();
if (isa<ComplexType, FloatType>(resultType)) {
return MapStableHloOpToScalarOpImpl<IsFloatType, math::PowFOp,
IsComplexType, complex::PowOp>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
// Exponentiation by squaring:
// https://en.wikipedia.org/wiki/Exponentiation_by_squaring;
Value negOne =
lb.create<arith::ConstantOp>(lb.getIntegerAttr(resultType, -1));
Value zero = lb.create<arith::ConstantOp>(lb.getIntegerAttr(resultType, 0));
Value one = lb.create<arith::ConstantOp>(lb.getIntegerAttr(resultType, 1));
Value two = lb.create<arith::ConstantOp>(lb.getIntegerAttr(resultType, 2));
Value step = lb.create<arith::ConstantIndexOp>(1);
Value lowerBound = lb.create<arith::ConstantIndexOp>(0);
// Everything else would overflow for any exponent > 1, as 2^64
// is the larget possible exponent for a 64-bit integer, and
// that's 1 << 6.
Value upperBound = lb.create<arith::ConstantIndexOp>(6);
auto originalBase = adaptor.getLhs();
auto originalExponent = adaptor.getRhs();
Value accum =
lb.create<scf::ForOp>(
lowerBound, upperBound, step,
SmallVector<Value>({one, originalBase, originalExponent}),
[&](OpBuilder &b, Location, Value /*v*/, ValueRange iters) {
Value accum = iters[0];
Value base = iters[1];
Value exponent = iters[2];
Value condition = b.create<arith::CmpIOp>(
loc, arith::CmpIPredicate::eq,
b.create<::mlir::arith::AndIOp>(loc, exponent, one), one);
Value multiplied =
b.create<::mlir::arith::MulIOp>(loc, accum, base);
accum = b.create<::mlir::arith::SelectOp>(loc, condition,
multiplied, accum);
base = b.create<::mlir::arith::MulIOp>(loc, base, base);
exponent = b.create<::mlir::arith::ShRUIOp>(loc, exponent, one);
b.create<scf::YieldOp>(
loc, SmallVector<Value>({accum, base, exponent}));
})
.getResult(0);
Value rhsIsEven = lb.create<arith::CmpIOp>(
arith::CmpIPredicate::eq,
lb.create<arith::RemSIOp>(adaptor.getRhs(), two), zero);
Value rhsIsNegative = lb.create<arith::CmpIOp>(arith::CmpIPredicate::slt,
adaptor.getRhs(), zero);
Value lhsIsOne =
lb.create<arith::CmpIOp>(arith::CmpIPredicate::eq, adaptor.getLhs(), one);
Value lhsIsNegOne = lb.create<arith::CmpIOp>(arith::CmpIPredicate::eq,
adaptor.getLhs(), negOne);
// The accum is correct when the rhs is non-negative. When rhs is
// negative, we return 0 for integer, with the exception of lhs values of 1
// and -1 which have integer results for negative exponents. Specifically, the
// calulation is the following:
//
// - Return accum if the rhs is not negative.
// - Return 1 or -1 depending on the parity of rhs when the lhs is -1.
// - Return 1 if lhs is 1.
// - Else return 0.
Value ifLhsIsOne = lb.create<::mlir::arith::SelectOp>(lhsIsOne, one, zero);
Value ifLhsIsNegOne = lb.create<::mlir::arith::SelectOp>(
lhsIsNegOne, lb.create<::mlir::arith::SelectOp>(rhsIsEven, one, negOne),
ifLhsIsOne);
return lb.create<::mlir::arith::SelectOp>(rhsIsNegative, ifLhsIsNegOne,
accum);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::SelectOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
stablehlo::SelectOp::Adaptor adaptor, OpBuilder *b) {
return MapStableHloOpToScalarOpImpl<::mlir::arith::SelectOp>{}(
loc, resultTypes, argTypes, adaptor.getOperands(), b);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::SignOp>(
Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> /*argTypes*/,
stablehlo::SignOp::Adaptor adaptor, OpBuilder *b) {
Value operand = adaptor.getOperand();
Type elementType = getElementTypeOrSelf(operand.getType());
if (auto floatType = dyn_cast<FloatType>(elementType)) {
Value zero =
b->create<arith::ConstantOp>(loc, b->getZeroAttr(operand.getType()));
Value ne0I1 = b->create<::mlir::arith::CmpFOp>(
loc, arith::CmpFPredicate::ONE, operand, zero);
Value ne0Float =
b->create<::mlir::arith::UIToFPOp>(loc, zero.getType(), ne0I1);
Value copySign = b->create<::mlir::math::CopySignOp>(loc, resultTypes,
ne0Float, operand);
auto isNan = b->create<::mlir::arith::CmpFOp>(
loc, arith::CmpFPredicate::UNO, operand, operand);
return b->create<::mlir::arith::SelectOp>(loc, isNan, operand, copySign);
}
if (auto integerType = dyn_cast<IntegerType>(elementType)) {
// sign(x) = x == 0 ? 0 : ((x s>> 31) | 1)
Value zero =
b->create<arith::ConstantOp>(loc, b->getZeroAttr(operand.getType()));
Value bitwidthMinusOne = getConstantOrSplat(
b, loc, operand.getType(),
b->getIntegerAttr(integerType, integerType.getWidth() - 1));
Value one = getConstantOrSplat(b, loc, operand.getType(),
b->getIntegerAttr(integerType, 1));
Value cmp = b->create<::mlir::arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
operand, zero);
Value ashr =
b->create<::mlir::arith::ShRSIOp>(loc, operand, bitwidthMinusOne);
Value orOp = b->create<::mlir::arith::OrIOp>(loc, ashr, one);
return b->create<::mlir::arith::SelectOp>(loc, cmp, zero, orOp);
}
if (isa<ComplexType>(elementType)) {
return b->create<::mlir::complex::SignOp>(loc, elementType, operand);
}
return nullptr;
}
/// Construct operations to select the saturated value if the shift amount is
/// greater than the bitwidth of the type.
inline Value selectShiftedOrSaturated(ImplicitLocOpBuilder &lb, Value rhs,
Value shifted, Value saturated,
Type type) {
Type etype = getElementTypeOrSelf(type);
auto bitWidthInt = etype.getIntOrFloatBitWidth();
Value bitWidth = getConstantOrSplat(&lb, lb.getLoc(), type,
lb.getIntegerAttr(etype, bitWidthInt));
Value cmp = lb.create<mlir::arith::CmpIOp>(mlir::arith::CmpIPredicate::ugt,
bitWidth, rhs);
return lb.create<mlir::arith::SelectOp>(cmp, shifted, saturated);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::ShiftLeftOp>(
Location loc, ArrayRef<Type> /*ResultTypes*/, ArrayRef<Type> /*argTypes*/,
stablehlo::ShiftLeftOp::Adaptor adaptor, OpBuilder *b) {
ImplicitLocOpBuilder lb(loc, *b);
Value lhs = adaptor.getLhs();
Value rhs = adaptor.getRhs();
Type type = lhs.getType();
// "Saturate" if the shift is greater than the bitwidth of the type
Value zero = lb.create<arith::ConstantOp>(lb.getZeroAttr(type));
Value shifted = lb.create<mlir::arith::ShLIOp>(lhs, rhs);
return selectShiftedOrSaturated(lb, rhs, shifted, zero, type);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::ShiftRightLogicalOp>(
Location loc, ArrayRef<Type> /*ResultTypes*/, ArrayRef<Type> /*argTypes*/,
stablehlo::ShiftRightLogicalOp::Adaptor adaptor, OpBuilder *b) {
ImplicitLocOpBuilder lb(loc, *b);
Value lhs = adaptor.getLhs();
Value rhs = adaptor.getRhs();
Type type = lhs.getType();
// "Saturate" if the shift is greater than the bitwidth of the type
Value zero = lb.create<arith::ConstantOp>(b->getZeroAttr(type));
Value shifted = lb.create<mlir::arith::ShRUIOp>(lhs, rhs);
return selectShiftedOrSaturated(lb, rhs, shifted, zero, type);
}
template <>
inline Value mapStableHloOpToStdScalarOp<stablehlo::ShiftRightArithmeticOp>(
Location loc, ArrayRef<Type> /*ResultTypes*/, ArrayRef<Type> /*argTypes*/,
stablehlo::ShiftRightArithmeticOp::Adaptor adaptor, OpBuilder *b) {
ImplicitLocOpBuilder lb(loc, *b);
Value lhs = adaptor.getLhs();
Value rhs = adaptor.getRhs();
Type type = lhs.getType();
Type etype = getElementTypeOrSelf(type);
auto bitWidthInt = etype.getIntOrFloatBitWidth();
// "Saturate" if the shift is greater than the bitwidth of the type
Value maxShift = getConstantOrSplat(
b, loc, type, lb.getIntegerAttr(etype, bitWidthInt - 1));
Value saturatedShifted = lb.create<mlir::arith::ShRSIOp>(lhs, maxShift);
Value shifted = lb.create<mlir::arith::ShRSIOp>(lhs, rhs);
return selectShiftedOrSaturated(lb, rhs, shifted, saturatedShifted, type);
}
} // namespace impl
struct StableHloOpToStdScalarOp {
// Converts stablehlo 'op' to linalg and arith ops.
template <typename StableHloOpTy>
static Value mapOp(StableHloOpTy op, ArrayRef<Type> resultTypes,
ValueRange args, OpBuilder *b) {
auto argTypes = llvm::to_vector(op->getOperandTypes());
return mapOpWithArgTypes(op, resultTypes, argTypes, args, b);
}
// Converts stablehlo 'op' to linalg and arith ops. The types of 'args' may
// already be converted, 'argTypes' are their original types.
template <typename StableHloOpTy>
static Value mapOpWithArgTypes(StableHloOpTy op, ArrayRef<Type> resultTypes,
ArrayRef<Type> argTypes, ValueRange args,
OpBuilder *b) {
static_assert(!std::is_same<StableHloOpTy, stablehlo::ConvertOp>::value);
return mapOpOfType<StableHloOpTy>(
op.getLoc(), resultTypes, argTypes,
typename StableHloOpTy::Adaptor(args, op->getAttrDictionary(),
op.getProperties()),
b);
}
// Overload for stablehlo::ConvertOp.
static Value mapOpWithArgTypes(stablehlo::ConvertOp op,
ArrayRef<Type> resultTypes,
ArrayRef<Type> argTypes, ValueRange args,
OpBuilder *b) {
return impl::mapConvertOpToStdScalarOp(op.getLoc(), op.getType(),
resultTypes, argTypes, args, b);
}
// Converts stablehlo 'op' to linalg and arith ops.
template <typename StableHloOpTy>
static Value
mapOpOfType(Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
typename StableHloOpTy::Adaptor adaptor, OpBuilder *b) {
if (std::is_same<StableHloOpTy, stablehlo::ConvertOp>::value) {
// Note: this assumes that the caller is passing result/arg types with
// appropriate signedness.
return impl::mapConvertOpToStdScalarOp(
loc, resultTypes, resultTypes, argTypes, adaptor.getOperands(), b);
}
return impl::mapStableHloOpToStdScalarOp<StableHloOpTy>(
loc, resultTypes, argTypes, adaptor, b);
}
};
} // namespace stablehlo
} // namespace mlir
#endif // IREE_COMPILER_PLUGINS_INPUT_STABLEHLO_CONVERSION_MAP_STABLEHLO_TO_SCALAR_OP_H