Add a custom LLVM pass to allow contraction of fast fadd/fsub with a …

…normal fmul.

Fix #18654, Fix #22217

src/Makefile

@@ -53,7 +53,8 @@ endif
 LLVMLINK :=
 
 ifeq ($(JULIACODEGEN),LLVM)
-SRCS += codegen jitlayers disasm debuginfo llvm-simdloop llvm-ptls llvm-gcroot llvm-lower-handlers cgmemmgr
+SRCS += codegen jitlayers disasm debuginfo llvm-simdloop llvm-ptls llvm-gcroot \
+	llvm-lower-handlers llvm-muladd cgmemmgr
 FLAGS += -I$(shell $(LLVM_CONFIG_HOST) --includedir)
 LLVM_LIBS := all
 ifeq ($(USE_POLLY),1)

src/intrinsics.cpp

@@ -955,7 +955,17 @@ static Value *emit_untyped_intrinsic(intrinsic f, Value **argvalues, size_t narg
 #endif
     }
     case muladd_float: {
-#if JL_LLVM_VERSION >= 30400
+#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);
+        auto mul = mathb().CreateFMul(x, y);
+        auto add = mathb().CreateFAdd(mul, z);
+        mul->setHasAllowContract(true);
+        add->setHasAllowContract(true);
+        return add;
+#elif JL_LLVM_VERSION >= 30400
         assert(y->getType() == x->getType());
         assert(z->getType() == y->getType());
         Value *muladdintr = Intrinsic::getDeclaration(jl_Module, Intrinsic::fmuladd, makeArrayRef(t));

src/jitlayers.cpp

@@ -266,6 +266,7 @@ void addOptimizationPasses(PassManager *PM)
     PM->add(createLoopVectorizePass());         // Vectorize loops
     PM->add(createInstructionCombiningPass());  // Clean up after loop vectorizer
 #endif
+    PM->add(createCombineMulAddPass());
 }
 
 #ifdef USE_ORCJIT

src/jitlayers.h

@@ -249,6 +249,7 @@ JL_DLLEXPORT extern LLVMContext &jl_LLVMContext;
 Pass *createLowerPTLSPass(bool imaging_mode);
 Pass *createLowerGCFramePass();
 Pass *createLowerExcHandlersPass();
+Pass *createCombineMulAddPass();
 // Whether the Function is an llvm or julia intrinsic.
 static inline bool isIntrinsicFunction(Function *F)
 {

src/llvm-muladd.cpp

@@ -0,0 +1,121 @@
+// 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.
+    addOp->setHasAllowContract(true);
+    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
+}