Merge pull request #22262 from JuliaLang/yyc/codegen/muladd

Add a custom LLVM pass to replace fastmath multiple and add with muladd

Author: yuyichao

base/math.jl

@@ -941,8 +941,11 @@ mod2pi(x) = rem2pi(x,RoundDown)
 """
     muladd(x, y, z)
 
-Combined multiply-add, computes `x*y+z` in an efficient manner. This may on some systems be
-equivalent to `x*y+z`, or to `fma(x,y,z)`. `muladd` is used to improve performance.
+Combined multiply-add, computes `x*y+z` allowing the add and multiply to be contracted with
+each other or ones from other `muladd` and `@fastmath` to form `fma`
+if the transformation can improve performance.
+The result can be different on different machines and can also be different on the same machine
+due to constant propagation or other optimizations.
 See [`fma`](@ref).
 
 # Example

src/Makefile

@@ -50,7 +50,7 @@ endif
 LLVMLINK :=
 
 ifeq ($(JULIACODEGEN),LLVM)
-SRCS += codegen jitlayers disasm debuginfo llvm-simdloop llvm-ptls llvm-late-gc-lowering llvm-lower-handlers llvm-gc-invariant-verifier llvm-propagate-addrspaces cgmemmgr
+SRCS += codegen jitlayers disasm debuginfo llvm-simdloop llvm-ptls llvm-muladd llvm-late-gc-lowering llvm-lower-handlers llvm-gc-invariant-verifier llvm-propagate-addrspaces cgmemmgr
 FLAGS += -I$(shell $(LLVM_CONFIG_HOST) --includedir)
 LLVM_LIBS := all
 ifeq ($(USE_POLLY),1)

src/intrinsics.cpp

@@ -704,17 +704,23 @@ static Value *emit_checked_srem_int(jl_codectx_t &ctx, Value *x, Value *den)
 struct math_builder {
     IRBuilder<> &ctxbuilder;
     FastMathFlags old_fmf;
-    math_builder(jl_codectx_t &ctx, bool always_fast = false)
+    math_builder(jl_codectx_t &ctx, bool always_fast = false, bool contract = false)
       : ctxbuilder(ctx.builder),
         old_fmf(ctxbuilder.getFastMathFlags())
     {
+        FastMathFlags fmf;
         if (jl_options.fast_math != JL_OPTIONS_FAST_MATH_OFF &&
             (always_fast ||
              jl_options.fast_math == JL_OPTIONS_FAST_MATH_ON)) {
-            FastMathFlags fmf;
             fmf.setUnsafeAlgebra();
-            ctxbuilder.setFastMathFlags(fmf);
         }
+#if JL_LLVM_VERSION >= 50000
+        if (contract)
+            fmf.setAllowContract(true);
+#else
+        assert(!contract);
+#endif
+        ctxbuilder.setFastMathFlags(fmf);
     }
     IRBuilder<>& operator()() const { return ctxbuilder; }
     ~math_builder() {
@@ -913,10 +919,18 @@ static Value *emit_untyped_intrinsic(jl_codectx_t &ctx, intrinsic f, Value **arg
         return ctx.builder.CreateCall(fmaintr, {x, y, z});
     }
     case muladd_float: {
+#if JL_LLVM_VERSION >= 50000
+        // LLVM 5.0 can create FMA in the backend for contractable fmul and fadd
+        // Emitting fmul and fadd here since they are easier for other LLVM passes to
+        // optimize.
+        auto mathb = math_builder(ctx, false, true);
+        return mathb().CreateFAdd(mathb().CreateFMul(x, y), z);
+#else
         assert(y->getType() == x->getType());
         assert(z->getType() == y->getType());
         Value *muladdintr = Intrinsic::getDeclaration(jl_Module, Intrinsic::fmuladd, makeArrayRef(t));
         return ctx.builder.CreateCall(muladdintr, {x, y, z});
+#endif
     }
 
     case checked_sadd_int:

src/jitlayers.cpp

@@ -242,6 +242,7 @@ void addOptimizationPasses(legacy::PassManagerBase *PM, int opt_level)
     PM->add(createDeadCodeEliminationPass());
     PM->add(createLowerPTLSPass(imaging_mode));
 #endif
+    PM->add(createCombineMulAddPass());
 }
 
 extern "C" JL_DLLEXPORT

src/jitlayers.h

@@ -202,6 +202,7 @@ extern JuliaOJIT *jl_ExecutionEngine;
 JL_DLLEXPORT extern LLVMContext jl_LLVMContext;
 
 Pass *createLowerPTLSPass(bool imaging_mode);
+Pass *createCombineMulAddPass();
 Pass *createLateLowerGCFramePass();
 Pass *createLowerExcHandlersPass();
 Pass *createGCInvariantVerifierPass(bool Strong);

src/llvm-muladd.cpp

@@ -0,0 +1,123 @@
+// This file is a part of Julia. License is MIT: https://julialang.org/license
+
+#define DEBUG_TYPE "combine_muladd"
+#undef DEBUG
+#include "llvm-version.h"
+
+#include <llvm/IR/Value.h>
+#include <llvm/IR/Function.h>
+#include <llvm/IR/Instructions.h>
+#include <llvm/IR/IntrinsicInst.h>
+#include <llvm/IR/Module.h>
+#include <llvm/IR/Operator.h>
+#include <llvm/IR/IRBuilder.h>
+#include <llvm/Pass.h>
+#include <llvm/Support/Debug.h>
+#include "fix_llvm_assert.h"
+
+#include "julia.h"
+
+using namespace llvm;
+
+/**
+ * Combine
+ * ```
+ * %v0 = fmul ... %a, %b
+ * %v = fadd fast ... %v0, %c
+ * ```
+ * to
+ * `%v = call fast @llvm.fmuladd.<...>(... %a, ... %b, ... %c)`
+ * when `%v0` has no other use
+ */
+
+struct CombineMulAdd : public FunctionPass {
+    static char ID;
+    CombineMulAdd() : FunctionPass(ID)
+    {}
+
+private:
+    bool runOnFunction(Function &F) override;
+};
+
+// Return true if this function shouldn't be called again on the other operand
+// This will always return false on LLVM 5.0+
+static bool checkCombine(Module *m, Instruction *addOp, Value *maybeMul, Value *addend,
+                         bool negadd, bool negres)
+{
+    auto mulOp = dyn_cast<Instruction>(maybeMul);
+    if (!mulOp || mulOp->getOpcode() != Instruction::FMul)
+        return false;
+    if (!mulOp->hasOneUse())
+        return false;
+#if JL_LLVM_VERSION >= 50000
+    // On 5.0+ we only need to mark the mulOp as contract and the backend will do the work for us.
+    auto fmf = mulOp->getFastMathFlags();
+    fmf.setAllowContract(true);
+    mulOp->copyFastMathFlags(fmf);
+    return false;
+#else
+    IRBuilder<> builder(m->getContext());
+    builder.SetInsertPoint(addOp);
+    auto mul1 = mulOp->getOperand(0);
+    auto mul2 = mulOp->getOperand(1);
+    Value *muladdf = Intrinsic::getDeclaration(m, Intrinsic::fmuladd, addOp->getType());
+    if (negadd) {
+        auto newaddend = builder.CreateFNeg(addend);
+        // Might be a const
+        if (auto neginst = dyn_cast<Instruction>(newaddend))
+            neginst->setHasUnsafeAlgebra(true);
+        addend = newaddend;
+    }
+    Instruction *newv = builder.CreateCall(muladdf, {mul1, mul2, addend});
+    newv->setHasUnsafeAlgebra(true);
+    if (negres) {
+        // Shouldn't be a constant
+        newv = cast<Instruction>(builder.CreateFNeg(newv));
+        newv->setHasUnsafeAlgebra(true);
+    }
+    addOp->replaceAllUsesWith(newv);
+    addOp->eraseFromParent();
+    mulOp->eraseFromParent();
+    return true;
+#endif
+}
+
+bool CombineMulAdd::runOnFunction(Function &F)
+{
+    Module *m = F.getParent();
+    for (auto &BB: F) {
+        for (auto it = BB.begin(); it != BB.end();) {
+            auto &I = *it;
+            it++;
+            switch (I.getOpcode()) {
+            case Instruction::FAdd: {
+                if (!I.hasUnsafeAlgebra())
+                    continue;
+                checkCombine(m, &I, I.getOperand(0), I.getOperand(1), false, false) ||
+                    checkCombine(m, &I, I.getOperand(1), I.getOperand(0), false, false);
+                break;
+            }
+            case Instruction::FSub: {
+                if (!I.hasUnsafeAlgebra())
+                    continue;
+                checkCombine(m, &I, I.getOperand(0), I.getOperand(1), true, false) ||
+                    checkCombine(m, &I, I.getOperand(1), I.getOperand(0), true, true);
+                break;
+            }
+            default:
+                break;
+            }
+        }
+    }
+    return true;
+}
+
+char CombineMulAdd::ID = 0;
+static RegisterPass<CombineMulAdd> X("CombineMulAdd", "Combine mul and add to muladd",
+                                     false /* Only looks at CFG */,
+                                     false /* Analysis Pass */);
+
+Pass *createCombineMulAddPass()
+{
+    return new CombineMulAdd();
+}

test/llvmpasses/muladd.ll

@@ -0,0 +1,28 @@
+; RUN: opt -load libjulia.so -CombineMulAdd -S %s | FileCheck %s
+
+define double @fast_muladd1(double %a, double %b, double %c) {
+top:
+; CHECK: {{contract|fmuladd}}
+  %v1 = fmul double %a, %b
+  %v2 = fadd fast double %v1, %c
+; CHECK: ret double
+  ret double %v2
+}
+
+define double @fast_mulsub1(double %a, double %b, double %c) {
+top:
+; CHECK: {{contract|fmuladd}}
+  %v1 = fmul double %a, %b
+  %v2 = fsub fast double %v1, %c
+; CHECK: ret double
+  ret double %v2
+}
+
+define <2 x double> @fast_mulsub_vec1(<2 x double> %a, <2 x double> %b, <2 x double> %c) {
+top:
+; CHECK: {{contract|fmuladd}}
+  %v1 = fmul <2 x double> %a, %b
+  %v2 = fsub fast <2 x double> %c, %v1
+; CHECK: ret <2 x double>
+  ret <2 x double> %v2
+}