102 lines · cpp
1//===- OptimizeForNVVM.cpp - Optimize LLVM IR for NVVM ---------===//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-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/LLVMIR/Transforms/OptimizeForNVVM.h"10 11#include "mlir/Dialect/LLVMIR/NVVMDialect.h"12#include "mlir/IR/Builders.h"13#include "mlir/IR/PatternMatch.h"14#include "mlir/Pass/Pass.h"15#include "mlir/Transforms/GreedyPatternRewriteDriver.h"16 17namespace mlir {18namespace LLVM {19#define GEN_PASS_DEF_NVVMOPTIMIZEFORTARGETPASS20#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc"21} // namespace LLVM22} // namespace mlir23 24using namespace mlir;25 26namespace {27// Replaces fdiv on fp16 with fp32 multiplication with reciprocal plus one28// (conditional) Newton iteration.29//30// This as accurate as promoting the division to fp32 in the NVPTX backend, but31// faster because it performs less Newton iterations, avoids the slow path32// for e.g. denormals, and allows reuse of the reciprocal for multiple divisions33// by the same divisor.34struct ExpandDivF16 : public OpRewritePattern<LLVM::FDivOp> {35 using OpRewritePattern<LLVM::FDivOp>::OpRewritePattern;36 37private:38 LogicalResult matchAndRewrite(LLVM::FDivOp op,39 PatternRewriter &rewriter) const override;40};41 42struct NVVMOptimizeForTarget43 : public LLVM::impl::NVVMOptimizeForTargetPassBase<NVVMOptimizeForTarget> {44 void runOnOperation() override;45 46 void getDependentDialects(DialectRegistry ®istry) const override {47 registry.insert<NVVM::NVVMDialect>();48 }49};50} // namespace51 52LogicalResult ExpandDivF16::matchAndRewrite(LLVM::FDivOp op,53 PatternRewriter &rewriter) const {54 if (!op.getType().isF16())55 return rewriter.notifyMatchFailure(op, "not f16");56 Location loc = op.getLoc();57 58 Type f32Type = rewriter.getF32Type();59 Type i32Type = rewriter.getI32Type();60 61 // Extend lhs and rhs to fp32.62 Value lhs = LLVM::FPExtOp::create(rewriter, loc, f32Type, op.getLhs());63 Value rhs = LLVM::FPExtOp::create(rewriter, loc, f32Type, op.getRhs());64 65 // float rcp = rcp.approx.ftz.f32(rhs), approx = lhs * rcp.66 Value rcp = NVVM::RcpApproxFtzF32Op::create(rewriter, loc, f32Type, rhs);67 Value approx = LLVM::FMulOp::create(rewriter, loc, lhs, rcp);68 69 // Refine the approximation with one Newton iteration:70 // float refined = approx + (lhs - approx * rhs) * rcp;71 Value err = LLVM::FMAOp::create(72 rewriter, loc, approx, LLVM::FNegOp::create(rewriter, loc, rhs), lhs);73 Value refined = LLVM::FMAOp::create(rewriter, loc, err, rcp, approx);74 75 // Use refined value if approx is normal (exponent neither all 0 or all 1).76 Value mask = LLVM::ConstantOp::create(77 rewriter, loc, i32Type, rewriter.getUI32IntegerAttr(0x7f800000));78 Value cast = LLVM::BitcastOp::create(rewriter, loc, i32Type, approx);79 Value exp = LLVM::AndOp::create(rewriter, loc, i32Type, cast, mask);80 Value zero = LLVM::ConstantOp::create(rewriter, loc, i32Type,81 rewriter.getUI32IntegerAttr(0));82 Value pred = LLVM::OrOp::create(83 rewriter, loc,84 LLVM::ICmpOp::create(rewriter, loc, LLVM::ICmpPredicate::eq, exp, zero),85 LLVM::ICmpOp::create(rewriter, loc, LLVM::ICmpPredicate::eq, exp, mask));86 Value result =87 LLVM::SelectOp::create(rewriter, loc, f32Type, pred, approx, refined);88 89 // Replace with trucation back to fp16.90 rewriter.replaceOpWithNewOp<LLVM::FPTruncOp>(op, op.getType(), result);91 92 return success();93}94 95void NVVMOptimizeForTarget::runOnOperation() {96 MLIRContext *ctx = getOperation()->getContext();97 RewritePatternSet patterns(ctx);98 patterns.add<ExpandDivF16>(ctx);99 if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))100 return signalPassFailure();101}102