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

This currently does the replacement as long as one of the instruction allows unsafe arithmetic
mainly because LLVM vectorization pass does not always preserve the fastmath flags.

Fix #18654

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/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,110 @@
+// This file is a part of Julia. License is MIT: https://julialang.org/license
+
+#define DEBUG_TYPE "lower_handlers"
+#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/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 ... %v0, %c
+ * ```
+ * to
+ * `%v = call ... @llvm.fmuladd.<...>(... %a, ... %b, ... %c)`
+ * when `%v0` has no other use and at least one of them allows unsafe arithmetic
+ */
+
+struct CombineMulAdd : public FunctionPass {
+ static char ID;
+ CombineMulAdd() : FunctionPass(ID)
+ {}
+
+private:
+ bool runOnFunction(Function &F) override;
+};
+
+static bool checkCombine(Module *m, Instruction *addOp, Value *maybeMul, Value *addend,
+ bool negadd, bool negres, bool has_unsafe)
+{
+ auto mulOp = dyn_cast<Instruction>(maybeMul);
+ if (!mulOp || mulOp->getOpcode() != Instruction::FMul ||
+ (!has_unsafe && !mulOp->hasUnsafeAlgebra()) || !mulOp->hasOneUse())
+ return false;
+ auto mul1 = mulOp->getOperand(0);
+ auto mul2 = mulOp->getOperand(1);
+ Type *T = addOp->getType();
+ Value *muladdf = Intrinsic::getDeclaration(m, Intrinsic::fmuladd, T);
+ if (negadd) {
+ auto newadd = BinaryOperator::CreateFNeg(addend);
+ newadd->setHasUnsafeAlgebra(true);
+ newadd->insertBefore(addOp);
+ addend = newadd;
+ }
+ Instruction *newv = CallInst::Create(muladdf, {mul1, mul2, addend}, "", addOp);
+ if (negres) {
+ newv = BinaryOperator::CreateFNeg(newv);
+ newv->setHasUnsafeAlgebra(true);
+ newv->insertBefore(addOp);
+ }
+ addOp->replaceAllUsesWith(newv);
+ addOp->eraseFromParent();
+ mulOp->eraseFromParent();
+ return true;
+}
+
+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: {
+ bool has_unsafe = I.hasUnsafeAlgebra();
+ checkCombine(m, &I, I.getOperand(0), I.getOperand(1),
+ false, false, has_unsafe) ||
+ checkCombine(m, &I, I.getOperand(1), I.getOperand(0),
+ false, false, has_unsafe);
+ break;
+ }
+ case Instruction::FSub: {
+ bool has_unsafe = I.hasUnsafeAlgebra();
+ checkCombine(m, &I, I.getOperand(0), I.getOperand(1),
+ true, false, has_unsafe) ||
+ checkCombine(m, &I, I.getOperand(1), I.getOperand(0),
+ true, true, has_unsafe);
+ 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();
+}