252 lines
10 KiB
C++
252 lines
10 KiB
C++
//===- MathToAPFloat.cpp - Mathmetic to APFloat Conversion ----------------===//
|
|
//
|
|
// 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 "Utils.h"
|
|
|
|
#include "mlir/Conversion/ArithAndMathToAPFloat/MathToAPFloat.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/Dialect/Func/Utils/Utils.h"
|
|
#include "mlir/Dialect/Math/IR/Math.h"
|
|
#include "mlir/Dialect/Math/Transforms/Passes.h"
|
|
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/IR/Verifier.h"
|
|
#include "mlir/Transforms/WalkPatternRewriteDriver.h"
|
|
|
|
namespace mlir {
|
|
#define GEN_PASS_DEF_MATHTOAPFLOATCONVERSIONPASS
|
|
#include "mlir/Conversion/Passes.h.inc"
|
|
} // namespace mlir
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::func;
|
|
|
|
struct AbsFOpToAPFloatConversion final : OpRewritePattern<math::AbsFOp> {
|
|
AbsFOpToAPFloatConversion(MLIRContext *context, SymbolOpInterface symTable,
|
|
PatternBenefit benefit = 1)
|
|
: OpRewritePattern<math::AbsFOp>(context, benefit), symTable(symTable) {}
|
|
|
|
LogicalResult matchAndRewrite(math::AbsFOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
if (failed(checkPreconditions(rewriter, op)))
|
|
return failure();
|
|
// Get APFloat function from runtime library.
|
|
auto i32Type = IntegerType::get(symTable->getContext(), 32);
|
|
auto i64Type = IntegerType::get(symTable->getContext(), 64);
|
|
FailureOr<FuncOp> fn = lookupOrCreateFnDecl(
|
|
rewriter, symTable, "_mlir_apfloat_abs", {i32Type, i64Type});
|
|
if (failed(fn))
|
|
return fn;
|
|
Location loc = op.getLoc();
|
|
rewriter.setInsertionPoint(op);
|
|
// Scalarize and convert to APFloat runtime calls.
|
|
Value repl = forEachScalarValue(
|
|
rewriter, loc, op.getOperand(), /*operand2=*/Value(), op.getType(),
|
|
[&](Value operand, Value, Type resultType) {
|
|
auto floatTy = cast<FloatType>(operand.getType());
|
|
auto intWType = rewriter.getIntegerType(floatTy.getWidth());
|
|
Value operandBits = arith::ExtUIOp::create(
|
|
rewriter, loc, i64Type,
|
|
arith::BitcastOp::create(rewriter, loc, intWType, operand));
|
|
// Call APFloat function.
|
|
Value semValue = getAPFloatSemanticsValue(rewriter, loc, floatTy);
|
|
SmallVector<Value> params = {semValue, operandBits};
|
|
Value negatedBits =
|
|
func::CallOp::create(rewriter, loc, TypeRange(i64Type),
|
|
SymbolRefAttr::get(*fn), params)
|
|
->getResult(0);
|
|
// Truncate result to the original width.
|
|
auto truncatedBits =
|
|
arith::TruncIOp::create(rewriter, loc, intWType, negatedBits);
|
|
return arith::BitcastOp::create(rewriter, loc, floatTy,
|
|
truncatedBits);
|
|
});
|
|
|
|
rewriter.replaceOp(op, repl);
|
|
return success();
|
|
}
|
|
|
|
SymbolOpInterface symTable;
|
|
};
|
|
|
|
template <typename OpTy>
|
|
struct IsOpToAPFloatConversion final : OpRewritePattern<OpTy> {
|
|
IsOpToAPFloatConversion(MLIRContext *context, const char *APFloatName,
|
|
SymbolOpInterface symTable,
|
|
PatternBenefit benefit = 1)
|
|
: OpRewritePattern<OpTy>(context, benefit), symTable(symTable),
|
|
APFloatName(APFloatName) {};
|
|
|
|
LogicalResult matchAndRewrite(OpTy op,
|
|
PatternRewriter &rewriter) const override {
|
|
if (failed(checkPreconditions(rewriter, op)))
|
|
return failure();
|
|
// Get APFloat function from runtime library.
|
|
auto i1 = IntegerType::get(symTable->getContext(), 1);
|
|
auto i32Type = IntegerType::get(symTable->getContext(), 32);
|
|
auto i64Type = IntegerType::get(symTable->getContext(), 64);
|
|
std::string funcName =
|
|
(llvm::Twine("_mlir_apfloat_is") + APFloatName).str();
|
|
FailureOr<FuncOp> fn = lookupOrCreateFnDecl(
|
|
rewriter, symTable, funcName, {i32Type, i64Type}, nullptr, i1);
|
|
if (failed(fn))
|
|
return fn;
|
|
Location loc = op.getLoc();
|
|
rewriter.setInsertionPoint(op);
|
|
// Scalarize and convert to APFloat runtime calls.
|
|
Value repl = forEachScalarValue(
|
|
rewriter, loc, op.getOperand(), /*operand2=*/Value(), op.getType(),
|
|
[&](Value operand, Value, Type resultType) {
|
|
auto floatTy = cast<FloatType>(operand.getType());
|
|
auto intWType = rewriter.getIntegerType(floatTy.getWidth());
|
|
Value operandBits = arith::ExtUIOp::create(
|
|
rewriter, loc, i64Type,
|
|
arith::BitcastOp::create(rewriter, loc, intWType, operand));
|
|
|
|
// Call APFloat function.
|
|
Value semValue = getAPFloatSemanticsValue(rewriter, loc, floatTy);
|
|
Value params[] = {semValue, operandBits};
|
|
return func::CallOp::create(rewriter, loc, TypeRange(i1),
|
|
SymbolRefAttr::get(*fn), params)
|
|
.getResult(0);
|
|
});
|
|
rewriter.replaceOp(op, repl);
|
|
return success();
|
|
}
|
|
|
|
SymbolOpInterface symTable;
|
|
const char *APFloatName;
|
|
};
|
|
|
|
struct FmaOpToAPFloatConversion final : OpRewritePattern<math::FmaOp> {
|
|
FmaOpToAPFloatConversion(MLIRContext *context, SymbolOpInterface symTable,
|
|
PatternBenefit benefit = 1)
|
|
: OpRewritePattern<math::FmaOp>(context, benefit), symTable(symTable) {};
|
|
|
|
LogicalResult matchAndRewrite(math::FmaOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
if (failed(checkPreconditions(rewriter, op)))
|
|
return failure();
|
|
// Cast operands to 64-bit integers.
|
|
mlir::Type resType = op.getResult().getType();
|
|
auto floatTy = dyn_cast<FloatType>(resType);
|
|
if (!floatTy) {
|
|
auto vecTy1 = cast<VectorType>(resType);
|
|
floatTy = llvm::cast<FloatType>(vecTy1.getElementType());
|
|
}
|
|
auto i32Type = IntegerType::get(symTable->getContext(), 32);
|
|
auto i64Type = IntegerType::get(symTable->getContext(), 64);
|
|
FailureOr<FuncOp> fn = lookupOrCreateFnDecl(
|
|
rewriter, symTable, "_mlir_apfloat_fused_multiply_add",
|
|
{i32Type, i64Type, i64Type, i64Type});
|
|
if (failed(fn))
|
|
return fn;
|
|
Location loc = op.getLoc();
|
|
rewriter.setInsertionPoint(op);
|
|
|
|
IntegerType intWType = rewriter.getIntegerType(floatTy.getWidth());
|
|
IntegerType int64Type = rewriter.getI64Type();
|
|
|
|
auto scalarFMA = [&rewriter, &loc, &floatTy, &fn, &intWType,
|
|
&int64Type](Value a, Value b, Value c) {
|
|
Value operand = arith::ExtUIOp::create(
|
|
rewriter, loc, int64Type,
|
|
arith::BitcastOp::create(rewriter, loc, intWType, a));
|
|
Value multiplicand = arith::ExtUIOp::create(
|
|
rewriter, loc, int64Type,
|
|
arith::BitcastOp::create(rewriter, loc, intWType, b));
|
|
Value addend = arith::ExtUIOp::create(
|
|
rewriter, loc, int64Type,
|
|
arith::BitcastOp::create(rewriter, loc, intWType, c));
|
|
// Call APFloat function.
|
|
Value semValue = getAPFloatSemanticsValue(rewriter, loc, floatTy);
|
|
SmallVector<Value> params = {semValue, operand, multiplicand, addend};
|
|
auto resultOp =
|
|
func::CallOp::create(rewriter, loc, TypeRange(rewriter.getI64Type()),
|
|
SymbolRefAttr::get(*fn), params);
|
|
|
|
// Truncate result to the original width.
|
|
auto trunc = arith::TruncIOp::create(rewriter, loc, intWType,
|
|
resultOp->getResult(0));
|
|
return arith::BitcastOp::create(rewriter, loc, floatTy, trunc);
|
|
};
|
|
|
|
if (auto vecTy1 = dyn_cast<VectorType>(op.getA().getType())) {
|
|
// Sanity check: Operand types must match.
|
|
assert(vecTy1 == dyn_cast<VectorType>(op.getB().getType()) &&
|
|
"expected same vector types");
|
|
assert(vecTy1 == dyn_cast<VectorType>(op.getC().getType()) &&
|
|
"expected same vector types");
|
|
// Prepare scalar operands.
|
|
ResultRange scalarOperands =
|
|
vector::ToElementsOp::create(rewriter, loc, op.getA())->getResults();
|
|
ResultRange scalarMultiplicands =
|
|
vector::ToElementsOp::create(rewriter, loc, op.getB())->getResults();
|
|
ResultRange scalarAddends =
|
|
vector::ToElementsOp::create(rewriter, loc, op.getC())->getResults();
|
|
// Call the function for each pair of scalar operands.
|
|
SmallVector<Value> results;
|
|
for (auto [operand, multiplicand, addend] : llvm::zip_equal(
|
|
scalarOperands, scalarMultiplicands, scalarAddends)) {
|
|
results.push_back(scalarFMA(operand, multiplicand, addend));
|
|
}
|
|
// Package the results into a vector.
|
|
auto fromElements = vector::FromElementsOp::create(
|
|
rewriter, loc,
|
|
vecTy1.cloneWith(/*shape=*/std::nullopt, results.front().getType()),
|
|
results);
|
|
rewriter.replaceOp(op, fromElements);
|
|
return success();
|
|
}
|
|
|
|
Value repl = scalarFMA(op.getA(), op.getB(), op.getC());
|
|
rewriter.replaceOp(op, repl);
|
|
return success();
|
|
}
|
|
|
|
SymbolOpInterface symTable;
|
|
};
|
|
|
|
namespace {
|
|
struct MathToAPFloatConversionPass final
|
|
: impl::MathToAPFloatConversionPassBase<MathToAPFloatConversionPass> {
|
|
using Base::Base;
|
|
|
|
void runOnOperation() override;
|
|
};
|
|
|
|
void MathToAPFloatConversionPass::runOnOperation() {
|
|
MLIRContext *context = &getContext();
|
|
RewritePatternSet patterns(context);
|
|
|
|
patterns.add<AbsFOpToAPFloatConversion>(context, getOperation());
|
|
patterns.add<IsOpToAPFloatConversion<math::IsFiniteOp>>(context, "finite",
|
|
getOperation());
|
|
patterns.add<IsOpToAPFloatConversion<math::IsInfOp>>(context, "infinite",
|
|
getOperation());
|
|
patterns.add<IsOpToAPFloatConversion<math::IsNaNOp>>(context, "nan",
|
|
getOperation());
|
|
patterns.add<IsOpToAPFloatConversion<math::IsNormalOp>>(context, "normal",
|
|
getOperation());
|
|
patterns.add<FmaOpToAPFloatConversion>(context, getOperation());
|
|
|
|
LogicalResult result = success();
|
|
ScopedDiagnosticHandler scopedHandler(context, [&result](Diagnostic &diag) {
|
|
if (diag.getSeverity() == DiagnosticSeverity::Error) {
|
|
result = failure();
|
|
}
|
|
// NB: if you don't return failure, no other diag handlers will fire (see
|
|
// mlir/lib/IR/Diagnostics.cpp:DiagnosticEngineImpl::emit).
|
|
return failure();
|
|
});
|
|
walkAndApplyPatterns(getOperation(), std::move(patterns));
|
|
if (failed(result))
|
|
return signalPassFailure();
|
|
}
|
|
} // namespace
|