Addressed review comments

1. Removed extra arguments passed to tryFoldBinaryFMul.
2. Removed temporary storage to collect the binary instructions.
3. Made guarding condition little easier to read.
4. Added one more test scenario.

Author: rajatbajpai

llvm/lib/Target/NVPTX/CMakeLists.txt

@@ -17,7 +17,7 @@ set(NVPTXCodeGen_sources
   NVPTXAssignValidGlobalNames.cpp
   NVPTXAtomicLower.cpp
   NVPTXCtorDtorLowering.cpp
-  NVPTXFoldFMA.cpp
+  NVPTXIRPeephole.cpp
   NVPTXForwardParams.cpp
   NVPTXFrameLowering.cpp
   NVPTXGenericToNVVM.cpp

llvm/lib/Target/NVPTX/NVPTX.h

@@ -52,7 +52,7 @@ FunctionPass *createNVPTXLowerAllocaPass();
 FunctionPass *createNVPTXLowerUnreachablePass(bool TrapUnreachable,
                                               bool NoTrapAfterNoreturn);
 FunctionPass *createNVPTXTagInvariantLoadsPass();
-FunctionPass *createNVPTXFoldFMAPass();
+FunctionPass *createNVPTXIRPeepholePass();
 MachineFunctionPass *createNVPTXPeephole();
 MachineFunctionPass *createNVPTXProxyRegErasurePass();
 MachineFunctionPass *createNVPTXForwardParamsPass();
@@ -77,14 +77,14 @@ void initializeNVPTXAAWrapperPassPass(PassRegistry &);
 void initializeNVPTXExternalAAWrapperPass(PassRegistry &);
 void initializeNVPTXPeepholePass(PassRegistry &);
 void initializeNVPTXTagInvariantLoadLegacyPassPass(PassRegistry &);
-void initializeNVPTXFoldFMAPass(PassRegistry &);
+void initializeNVPTXIRPeepholePass(PassRegistry &);
 void initializeNVPTXPrologEpilogPassPass(PassRegistry &);
 
 struct NVVMIntrRangePass : PassInfoMixin<NVVMIntrRangePass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
-struct NVPTXFoldFMAPass : PassInfoMixin<NVPTXFoldFMAPass> {
+struct NVPTXIRPeepholePass : PassInfoMixin<NVPTXIRPeepholePass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 

llvm/lib/Target/NVPTX/NVPTXFoldFMA.cpp llvm/lib/Target/NVPTX/NVPTXIRPeephole.cppllvm/lib/Target/NVPTX/NVPTXFoldFMA.cpp renamed to llvm/lib/Target/NVPTX/NVPTXIRPeephole.cpp

@@ -1,4 +1,4 @@
-//===------ NVPTXFoldFMA.cpp - Fold FMA --------------===//
+//===------ NVPTXIRPeephole.cpp - NVPTX IR Peephole --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -22,18 +22,37 @@
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 
-#define DEBUG_TYPE "nvptx-fold-fma"
+#define DEBUG_TYPE "nvptx-ir-peephole"
 
 using namespace llvm;
 
-static bool tryFoldBinaryFMul(BinaryOperator *BI, Value *MulOperand,
-                              Value *OtherOperand, bool IsFirstOperand,
-                              bool IsFSub) {
-  auto *FMul = dyn_cast<BinaryOperator>(MulOperand);
-  if (!FMul || FMul->getOpcode() != Instruction::FMul || !FMul->hasOneUse() ||
-      !FMul->hasAllowContract())
+static bool tryFoldBinaryFMul(BinaryOperator *BI) {
+  Value *Op0 = BI->getOperand(0);
+  Value *Op1 = BI->getOperand(1);
+
+  auto *FMul0 = dyn_cast<BinaryOperator>(Op0);
+  auto *FMul1 = dyn_cast<BinaryOperator>(Op1);
+
+  BinaryOperator *FMul = nullptr;
+  Value *OtherOperand = nullptr;
+  bool IsFirstOperand = false;
+
+  // Either Op0 or Op1 should be a valid FMul
+  if (FMul0 && FMul0->getOpcode() == Instruction::FMul && FMul0->hasOneUse() &&
+      FMul0->hasAllowContract()) {
+    FMul = FMul0;
+    OtherOperand = Op1;
+    IsFirstOperand = true;
+  } else if (FMul1 && FMul1->getOpcode() == Instruction::FMul &&
+             FMul1->hasOneUse() && FMul1->hasAllowContract()) {
+    FMul = FMul1;
+    OtherOperand = Op0;
+    IsFirstOperand = false;
+  } else {
     return false;
+  }
 
+  bool IsFSub = BI->getOpcode() == Instruction::FSub;
   LLVM_DEBUG({
     const char *OpName = IsFSub ? "FSub" : "FAdd";
     dbgs() << "Found " << OpName << " with FMul (single use) as "
@@ -87,10 +106,9 @@ static bool tryFoldBinaryFMul(BinaryOperator *BI, Value *MulOperand,
 
 static bool foldFMA(Function &F) {
   bool Changed = false;
-  SmallVector<BinaryOperator *, 16> FAddFSubInsts;
 
-  // Collect all float/double FAdd/FSub instructions with allow-contract
-  for (auto &I : instructions(F)) {
+  // Iterate and process float/double FAdd/FSub instructions with allow-contract
+  for (auto &I : llvm::make_early_inc_range(instructions(F))) {
     if (auto *BI = dyn_cast<BinaryOperator>(&I)) {
       // Only FAdd and FSub are supported.
       if (BI->getOpcode() != Instruction::FAdd &&
@@ -105,42 +123,35 @@ static bool foldFMA(Function &F) {
       if (!BI->getType()->isFloatTy() && !BI->getType()->isDoubleTy())
         continue;
 
-      FAddFSubInsts.push_back(BI);
+      if (tryFoldBinaryFMul(BI))
+        Changed = true;
     }
   }
-
-  for (auto *BI : FAddFSubInsts) {
-    Value *Op0 = BI->getOperand(0);
-    Value *Op1 = BI->getOperand(1);
-    bool IsFSub = BI->getOpcode() == Instruction::FSub;
-
-    if (tryFoldBinaryFMul(BI, Op0, Op1, true /*IsFirstOperand*/, IsFSub) ||
-        tryFoldBinaryFMul(BI, Op1, Op0, false /*IsFirstOperand*/, IsFSub))
-      Changed = true;
-  }
-
   return Changed;
 }
 
 namespace {
 
-struct NVPTXFoldFMA : public FunctionPass {
+struct NVPTXIRPeephole : public FunctionPass {
   static char ID;
-  NVPTXFoldFMA() : FunctionPass(ID) {}
+  NVPTXIRPeephole() : FunctionPass(ID) {}
   bool runOnFunction(Function &F) override;
 };
 
 } // namespace
 
-char NVPTXFoldFMA::ID = 0;
-INITIALIZE_PASS(NVPTXFoldFMA, "nvptx-fold-fma", "NVPTX Fold FMA", false, false)
+char NVPTXIRPeephole::ID = 0;
+INITIALIZE_PASS(NVPTXIRPeephole, "nvptx-ir-peephole", "NVPTX IR Peephole",
+                false, false)
 
-bool NVPTXFoldFMA::runOnFunction(Function &F) { return foldFMA(F); }
+bool NVPTXIRPeephole::runOnFunction(Function &F) { return foldFMA(F); }
 
-FunctionPass *llvm::createNVPTXFoldFMAPass() { return new NVPTXFoldFMA(); }
+FunctionPass *llvm::createNVPTXIRPeepholePass() {
+  return new NVPTXIRPeephole();
+}
 
-PreservedAnalyses NVPTXFoldFMAPass::run(Function &F,
-                                        FunctionAnalysisManager &) {
+PreservedAnalyses NVPTXIRPeepholePass::run(Function &F,
+                                           FunctionAnalysisManager &) {
   if (!foldFMA(F))
     return PreservedAnalyses::all();
 

llvm/lib/Target/NVPTX/NVPTXPassRegistry.def

@@ -40,5 +40,5 @@ FUNCTION_PASS("nvvm-intr-range", NVVMIntrRangePass())
 FUNCTION_PASS("nvptx-copy-byval-args", NVPTXCopyByValArgsPass())
 FUNCTION_PASS("nvptx-lower-args", NVPTXLowerArgsPass(*this))
 FUNCTION_PASS("nvptx-tag-invariant-loads", NVPTXTagInvariantLoadsPass())
-FUNCTION_PASS("nvptx-fold-fma", NVPTXFoldFMAPass())
+FUNCTION_PASS("nvptx-ir-peephole", NVPTXIRPeepholePass())
 #undef FUNCTION_PASS

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

@@ -51,11 +51,12 @@ static cl::opt<bool>
                                cl::desc("Disable load/store vectorizer"),
                                cl::init(false), cl::Hidden);
 
-// FoldFMA is a new pass; this option will lets us turn it off in case we
-// encounter some issues.
-static cl::opt<bool> DisableFoldFMA("disable-nvptx-fold-fma",
-                                    cl::desc("Disable NVPTX Fold FMA"),
-                                    cl::init(false), cl::Hidden);
+// NVPTX IR Peephole is a new pass; this option will lets us turn it off in case
+// we encounter some issues.
+static cl::opt<bool>
+    DisableNVPTXIRPeephole("disable-nvptx-ir-peephole",
+                           cl::desc("Disable NVPTX IR Peephole"),
+                           cl::init(false), cl::Hidden);
 
 // TODO: Remove this flag when we are confident with no regressions.
 static cl::opt<bool> DisableRequireStructuredCFG(
@@ -121,7 +122,7 @@ extern "C" LLVM_ABI LLVM_EXTERNAL_VISIBILITY void LLVMInitializeNVPTXTarget() {
   initializeNVPTXExternalAAWrapperPass(PR);
   initializeNVPTXPeepholePass(PR);
   initializeNVPTXTagInvariantLoadLegacyPassPass(PR);
-  initializeNVPTXFoldFMAPass(PR);
+  initializeNVPTXIRPeepholePass(PR);
   initializeNVPTXPrologEpilogPassPass(PR);
 }
 
@@ -404,8 +405,8 @@ void NVPTXPassConfig::addIRPasses() {
       addPass(createLoadStoreVectorizerPass());
     addPass(createSROAPass());
     addPass(createNVPTXTagInvariantLoadsPass());
-    if (!DisableFoldFMA)
-      addPass(createNVPTXFoldFMAPass());
+    if (!DisableNVPTXIRPeephole)
+      addPass(createNVPTXIRPeepholePass());
   }
 
   if (ST.hasPTXASUnreachableBug()) {

llvm/test/CodeGen/NVPTX/nvptx-fold-fma.ll

@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt < %s -passes=nvptx-fold-fma -S | FileCheck %s
+; RUN: opt < %s -passes=nvptx-ir-peephole -S | FileCheck %s
 
 target triple = "nvptx64-nvidia-cuda"
 
@@ -47,6 +47,25 @@ define float @test_fsub_fmul_fmul(float %a, float %b, float %c, float %d) {
 }
 
 
+; fsub(fmul(a, b), fmul(c, d)) => fma(fneg(c), d, fmul(a, b)))
+; fmul(a, b) has multiple uses.
+define float @test_fsub_fmul_fmul_multiple_use(float %a, float %b, float %c, float %d) {
+; CHECK-LABEL: define float @test_fsub_fmul_fmul_multiple_use(
+; CHECK-SAME: float [[A:%.*]], float [[B:%.*]], float [[C:%.*]], float [[D:%.*]]) {
+; CHECK-NEXT:    [[MUL1:%.*]] = fmul contract float [[A]], [[B]]
+; CHECK-NEXT:    [[TMP1:%.*]] = fneg contract float [[C]]
+; CHECK-NEXT:    [[TMP2:%.*]] = call contract float @llvm.fma.f32(float [[TMP1]], float [[D]], float [[MUL1]])
+; CHECK-NEXT:    [[ADD:%.*]] = fadd float [[TMP2]], [[MUL1]]
+; CHECK-NEXT:    ret float [[ADD]]
+;
+  %mul1 = fmul contract float %a, %b
+  %mul2 = fmul contract float %c, %d
+  %sub = fsub contract float %mul1, %mul2
+  %add = fadd float %sub, %mul1
+  ret float %add
+}
+
+
 ; fsub(fmul(a, b), c) => fma(a, b, fneg(c)) where fsub and fmul are in different BBs
 define float @test_fsub_fmul_different_BB(float %a, float %b, float %c, i32 %n) {
 ; CHECK-LABEL: define float @test_fsub_fmul_different_BB(