|
| 1 | +//===- ArithToAPFloat.cpp - Arithmetic to APFloat Conversion --------------===// |
| 2 | +// |
| 3 | +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| 4 | +// See https://llvm.org/LICENSE.txt for license information. |
| 5 | +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| 6 | +// |
| 7 | +//===----------------------------------------------------------------------===// |
| 8 | + |
| 9 | +#include "mlir/Conversion/ArithToAPFloat/ArithToAPFloat.h" |
| 10 | + |
| 11 | +#include "mlir/Dialect/Arith/IR/Arith.h" |
| 12 | +#include "mlir/Dialect/Arith/Transforms/Passes.h" |
| 13 | +#include "mlir/Dialect/Func/IR/FuncOps.h" |
| 14 | +#include "mlir/Dialect/Func/Utils/Utils.h" |
| 15 | +#include "mlir/IR/PatternMatch.h" |
| 16 | +#include "mlir/IR/Verifier.h" |
| 17 | +#include "mlir/Transforms/WalkPatternRewriteDriver.h" |
| 18 | +#include "llvm/Support/Debug.h" |
| 19 | + |
| 20 | +#define DEBUG_TYPE "arith-to-apfloat" |
| 21 | + |
| 22 | +namespace mlir { |
| 23 | +#define GEN_PASS_DEF_ARITHTOAPFLOATCONVERSIONPASS |
| 24 | +#include "mlir/Conversion/Passes.h.inc" |
| 25 | +} // namespace mlir |
| 26 | + |
| 27 | +using namespace mlir; |
| 28 | +using namespace mlir::func; |
| 29 | + |
| 30 | +static FuncOp createFnDecl(OpBuilder &b, SymbolOpInterface symTable, |
| 31 | + StringRef name, FunctionType funcT, bool setPrivate, |
| 32 | + SymbolTableCollection *symbolTables = nullptr) { |
| 33 | + OpBuilder::InsertionGuard g(b); |
| 34 | + assert(!symTable->getRegion(0).empty() && "expected non-empty region"); |
| 35 | + b.setInsertionPointToStart(&symTable->getRegion(0).front()); |
| 36 | + FuncOp funcOp = FuncOp::create(b, symTable->getLoc(), name, funcT); |
| 37 | + if (setPrivate) |
| 38 | + funcOp.setPrivate(); |
| 39 | + if (symbolTables) { |
| 40 | + SymbolTable &symbolTable = symbolTables->getSymbolTable(symTable); |
| 41 | + symbolTable.insert(funcOp, symTable->getRegion(0).front().begin()); |
| 42 | + } |
| 43 | + return funcOp; |
| 44 | +} |
| 45 | + |
| 46 | +/// Helper function to look up or create the symbol for a runtime library |
| 47 | +/// function for a binary arithmetic operation. |
| 48 | +/// |
| 49 | +/// Parameter 1: APFloat semantics |
| 50 | +/// Parameter 2: Left-hand side operand |
| 51 | +/// Parameter 3: Right-hand side operand |
| 52 | +/// |
| 53 | +/// This function will return a failure if the function is found but has an |
| 54 | +/// unexpected signature. |
| 55 | +/// |
| 56 | +static FailureOr<FuncOp> |
| 57 | +lookupOrCreateBinaryFn(OpBuilder &b, SymbolOpInterface symTable, StringRef name, |
| 58 | + SymbolTableCollection *symbolTables = nullptr) { |
| 59 | + auto i32Type = IntegerType::get(symTable->getContext(), 32); |
| 60 | + auto i64Type = IntegerType::get(symTable->getContext(), 64); |
| 61 | + |
| 62 | + std::string funcName = (llvm::Twine("__mlir_apfloat_") + name).str(); |
| 63 | + FunctionType funcT = |
| 64 | + FunctionType::get(b.getContext(), {i32Type, i64Type, i64Type}, {i64Type}); |
| 65 | + FailureOr<FuncOp> func = |
| 66 | + lookupFnDecl(symTable, funcName, funcT, symbolTables); |
| 67 | + // Failed due to type mismatch. |
| 68 | + if (failed(func)) |
| 69 | + return func; |
| 70 | + // Successfully matched existing decl. |
| 71 | + if (*func) |
| 72 | + return *func; |
| 73 | + |
| 74 | + return createFnDecl(b, symTable, funcName, funcT, |
| 75 | + /*setPrivate=*/true, symbolTables); |
| 76 | +} |
| 77 | + |
| 78 | +/// Rewrite a binary arithmetic operation to an APFloat function call. |
| 79 | +template <typename OpTy, const char *APFloatName> |
| 80 | +struct BinaryArithOpToAPFloatConversion final : OpRewritePattern<OpTy> { |
| 81 | + BinaryArithOpToAPFloatConversion(MLIRContext *context, PatternBenefit benefit, |
| 82 | + SymbolOpInterface symTable) |
| 83 | + : OpRewritePattern<OpTy>(context, benefit), symTable(symTable) {}; |
| 84 | + |
| 85 | + LogicalResult matchAndRewrite(OpTy op, |
| 86 | + PatternRewriter &rewriter) const override { |
| 87 | + // Get APFloat function from runtime library. |
| 88 | + FailureOr<FuncOp> fn = |
| 89 | + lookupOrCreateBinaryFn(rewriter, symTable, APFloatName); |
| 90 | + if (failed(fn)) |
| 91 | + return fn; |
| 92 | + |
| 93 | + rewriter.setInsertionPoint(op); |
| 94 | + // Cast operands to 64-bit integers. |
| 95 | + Location loc = op.getLoc(); |
| 96 | + auto floatTy = cast<FloatType>(op.getType()); |
| 97 | + auto intWType = rewriter.getIntegerType(floatTy.getWidth()); |
| 98 | + auto int64Type = rewriter.getI64Type(); |
| 99 | + Value lhsBits = arith::ExtUIOp::create( |
| 100 | + rewriter, loc, int64Type, |
| 101 | + arith::BitcastOp::create(rewriter, loc, intWType, op.getLhs())); |
| 102 | + Value rhsBits = arith::ExtUIOp::create( |
| 103 | + rewriter, loc, int64Type, |
| 104 | + arith::BitcastOp::create(rewriter, loc, intWType, op.getRhs())); |
| 105 | + |
| 106 | + // Call APFloat function. |
| 107 | + int32_t sem = |
| 108 | + llvm::APFloatBase::SemanticsToEnum(floatTy.getFloatSemantics()); |
| 109 | + Value semValue = arith::ConstantOp::create( |
| 110 | + rewriter, loc, rewriter.getI32Type(), |
| 111 | + rewriter.getIntegerAttr(rewriter.getI32Type(), sem)); |
| 112 | + SmallVector<Value> params = {semValue, lhsBits, rhsBits}; |
| 113 | + auto resultOp = |
| 114 | + func::CallOp::create(rewriter, loc, TypeRange(rewriter.getI64Type()), |
| 115 | + SymbolRefAttr::get(*fn), params); |
| 116 | + |
| 117 | + // Truncate result to the original width. |
| 118 | + Value truncatedBits = arith::TruncIOp::create(rewriter, loc, intWType, |
| 119 | + resultOp->getResult(0)); |
| 120 | + rewriter.replaceOp( |
| 121 | + op, arith::BitcastOp::create(rewriter, loc, floatTy, truncatedBits)); |
| 122 | + return success(); |
| 123 | + } |
| 124 | + |
| 125 | + SymbolOpInterface symTable; |
| 126 | +}; |
| 127 | + |
| 128 | +namespace { |
| 129 | +struct ArithToAPFloatConversionPass final |
| 130 | + : impl::ArithToAPFloatConversionPassBase<ArithToAPFloatConversionPass> { |
| 131 | + using Base::Base; |
| 132 | + |
| 133 | + void runOnOperation() override { |
| 134 | + MLIRContext *context = &getContext(); |
| 135 | + RewritePatternSet patterns(context); |
| 136 | + static const char add[] = "add"; |
| 137 | + static const char subtract[] = "subtract"; |
| 138 | + static const char multiply[] = "multiply"; |
| 139 | + static const char divide[] = "divide"; |
| 140 | + static const char remainder[] = "remainder"; |
| 141 | + patterns. add<BinaryArithOpToAPFloatConversion<arith::AddFOp, add>, |
| 142 | + BinaryArithOpToAPFloatConversion<arith::SubFOp, subtract>, |
| 143 | + BinaryArithOpToAPFloatConversion<arith::MulFOp, multiply>, |
| 144 | + BinaryArithOpToAPFloatConversion<arith::DivFOp, divide>, |
| 145 | + BinaryArithOpToAPFloatConversion<arith::RemFOp, remainder>>( |
| 146 | + context, 1, getOperation()); |
| 147 | + LogicalResult result = success(); |
| 148 | + ScopedDiagnosticHandler scopedHandler(context, [&result](Diagnostic &diag) { |
| 149 | + if (diag.getSeverity() == DiagnosticSeverity::Error) { |
| 150 | + result = failure(); |
| 151 | + } |
| 152 | + // NB: if you don't return failure, no other diag handlers will fire (see |
| 153 | + // mlir/lib/IR/Diagnostics.cpp:DiagnosticEngineImpl::emit). |
| 154 | + return failure(); |
| 155 | + }); |
| 156 | + walkAndApplyPatterns(getOperation(), std::move(patterns)); |
| 157 | + if (failed(result)) |
| 158 | + return signalPassFailure(); |
| 159 | + } |
| 160 | +}; |
| 161 | +} // namespace |
0 commit comments