Matthias Springer 78994706d8
[mlir][arith] Add support for extf, truncf to ArithToAPFloat (#169275)
Add support for `arith.extf` and `arith.truncf`. No support for custom
rounding modes yet.
2025-11-25 10:09:26 +09:00

226 lines
9.1 KiB
C++

//===- ArithToAPFloat.cpp - Arithmetic 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 "mlir/Conversion/ArithToAPFloat/ArithToAPFloat.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Utils/Utils.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"
namespace mlir {
#define GEN_PASS_DEF_ARITHTOAPFLOATCONVERSIONPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
using namespace mlir;
using namespace mlir::func;
static FuncOp createFnDecl(OpBuilder &b, SymbolOpInterface symTable,
StringRef name, FunctionType funcT, bool setPrivate,
SymbolTableCollection *symbolTables = nullptr) {
OpBuilder::InsertionGuard g(b);
assert(!symTable->getRegion(0).empty() && "expected non-empty region");
b.setInsertionPointToStart(&symTable->getRegion(0).front());
FuncOp funcOp = FuncOp::create(b, symTable->getLoc(), name, funcT);
if (setPrivate)
funcOp.setPrivate();
if (symbolTables) {
SymbolTable &symbolTable = symbolTables->getSymbolTable(symTable);
symbolTable.insert(funcOp, symTable->getRegion(0).front().begin());
}
return funcOp;
}
/// Helper function to look up or create the symbol for a runtime library
/// function with the given parameter types. Always returns an int64_t.
static FailureOr<FuncOp>
lookupOrCreateApFloatFn(OpBuilder &b, SymbolOpInterface symTable,
StringRef name, TypeRange paramTypes,
SymbolTableCollection *symbolTables = nullptr) {
auto i64Type = IntegerType::get(symTable->getContext(), 64);
std::string funcName = (llvm::Twine("_mlir_apfloat_") + name).str();
auto funcT = FunctionType::get(b.getContext(), paramTypes, {i64Type});
FailureOr<FuncOp> func =
lookupFnDecl(symTable, funcName, funcT, symbolTables);
// Failed due to type mismatch.
if (failed(func))
return func;
// Successfully matched existing decl.
if (*func)
return *func;
return createFnDecl(b, symTable, funcName, funcT,
/*setPrivate=*/true, symbolTables);
}
/// Helper function to look up or create the symbol for a runtime library
/// function for a binary arithmetic operation.
///
/// Parameter 1: APFloat semantics
/// Parameter 2: Left-hand side operand
/// Parameter 3: Right-hand side operand
///
/// This function will return a failure if the function is found but has an
/// unexpected signature.
///
static FailureOr<FuncOp>
lookupOrCreateBinaryFn(OpBuilder &b, SymbolOpInterface symTable, StringRef name,
SymbolTableCollection *symbolTables = nullptr) {
auto i32Type = IntegerType::get(symTable->getContext(), 32);
auto i64Type = IntegerType::get(symTable->getContext(), 64);
return lookupOrCreateApFloatFn(b, symTable, name, {i32Type, i64Type, i64Type},
symbolTables);
}
static Value getSemanticsValue(OpBuilder &b, Location loc, FloatType floatTy) {
int32_t sem = llvm::APFloatBase::SemanticsToEnum(floatTy.getFloatSemantics());
return arith::ConstantOp::create(b, loc, b.getI32Type(),
b.getIntegerAttr(b.getI32Type(), sem));
}
/// Rewrite a binary arithmetic operation to an APFloat function call.
template <typename OpTy>
struct BinaryArithOpToAPFloatConversion final : OpRewritePattern<OpTy> {
BinaryArithOpToAPFloatConversion(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 {
// Get APFloat function from runtime library.
FailureOr<FuncOp> fn =
lookupOrCreateBinaryFn(rewriter, symTable, APFloatName);
if (failed(fn))
return fn;
rewriter.setInsertionPoint(op);
// Cast operands to 64-bit integers.
Location loc = op.getLoc();
auto floatTy = cast<FloatType>(op.getType());
auto intWType = rewriter.getIntegerType(floatTy.getWidth());
auto int64Type = rewriter.getI64Type();
Value lhsBits = arith::ExtUIOp::create(
rewriter, loc, int64Type,
arith::BitcastOp::create(rewriter, loc, intWType, op.getLhs()));
Value rhsBits = arith::ExtUIOp::create(
rewriter, loc, int64Type,
arith::BitcastOp::create(rewriter, loc, intWType, op.getRhs()));
// Call APFloat function.
Value semValue = getSemanticsValue(rewriter, loc, floatTy);
SmallVector<Value> params = {semValue, lhsBits, rhsBits};
auto resultOp =
func::CallOp::create(rewriter, loc, TypeRange(rewriter.getI64Type()),
SymbolRefAttr::get(*fn), params);
// Truncate result to the original width.
Value truncatedBits = arith::TruncIOp::create(rewriter, loc, intWType,
resultOp->getResult(0));
rewriter.replaceOp(
op, arith::BitcastOp::create(rewriter, loc, floatTy, truncatedBits));
return success();
}
SymbolOpInterface symTable;
const char *APFloatName;
};
template <typename OpTy>
struct FpToFpConversion final : OpRewritePattern<OpTy> {
FpToFpConversion(MLIRContext *context, SymbolOpInterface symTable,
PatternBenefit benefit = 1)
: OpRewritePattern<OpTy>(context, benefit), symTable(symTable) {}
LogicalResult matchAndRewrite(OpTy op,
PatternRewriter &rewriter) const override {
// Get APFloat function from runtime library.
auto i32Type = IntegerType::get(symTable->getContext(), 32);
auto i64Type = IntegerType::get(symTable->getContext(), 64);
FailureOr<FuncOp> fn = lookupOrCreateApFloatFn(
rewriter, symTable, "convert", {i32Type, i32Type, i64Type});
if (failed(fn))
return fn;
rewriter.setInsertionPoint(op);
// Cast operands to 64-bit integers.
Location loc = op.getLoc();
auto inFloatTy = cast<FloatType>(op.getOperand().getType());
auto inIntWType = rewriter.getIntegerType(inFloatTy.getWidth());
auto int64Type = rewriter.getI64Type();
Value operandBits = arith::ExtUIOp::create(
rewriter, loc, int64Type,
arith::BitcastOp::create(rewriter, loc, inIntWType, op.getOperand()));
// Call APFloat function.
Value inSemValue = getSemanticsValue(rewriter, loc, inFloatTy);
auto outFloatTy = cast<FloatType>(op.getType());
Value outSemValue = getSemanticsValue(rewriter, loc, outFloatTy);
std::array<Value, 3> params = {inSemValue, outSemValue, operandBits};
auto resultOp =
func::CallOp::create(rewriter, loc, TypeRange(rewriter.getI64Type()),
SymbolRefAttr::get(*fn), params);
// Truncate result to the original width.
auto outIntWType = rewriter.getIntegerType(outFloatTy.getWidth());
Value truncatedBits = arith::TruncIOp::create(rewriter, loc, outIntWType,
resultOp->getResult(0));
rewriter.replaceOp(
op, arith::BitcastOp::create(rewriter, loc, outFloatTy, truncatedBits));
return success();
}
SymbolOpInterface symTable;
};
namespace {
struct ArithToAPFloatConversionPass final
: impl::ArithToAPFloatConversionPassBase<ArithToAPFloatConversionPass> {
using Base::Base;
void runOnOperation() override;
};
void ArithToAPFloatConversionPass::runOnOperation() {
MLIRContext *context = &getContext();
RewritePatternSet patterns(context);
patterns.add<BinaryArithOpToAPFloatConversion<arith::AddFOp>>(context, "add",
getOperation());
patterns.add<BinaryArithOpToAPFloatConversion<arith::SubFOp>>(
context, "subtract", getOperation());
patterns.add<BinaryArithOpToAPFloatConversion<arith::MulFOp>>(
context, "multiply", getOperation());
patterns.add<BinaryArithOpToAPFloatConversion<arith::DivFOp>>(
context, "divide", getOperation());
patterns.add<BinaryArithOpToAPFloatConversion<arith::RemFOp>>(
context, "remainder", getOperation());
patterns
.add<FpToFpConversion<arith::ExtFOp>, FpToFpConversion<arith::TruncFOp>>(
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