AMDGPU: Handle multiple uses when matching sincos

Match how the generic implementation handles this. We now will leave
behind the dead other user for later passes to deal with.

https://reviews.llvm.org/D156707

Author: arsenm

llvm/lib/Target/AMDGPU/AMDGPULibCalls.cpp

@@ -82,8 +82,12 @@ class AMDGPULibCalls {
   // sqrt
   bool fold_sqrt(FPMathOperator *FPOp, IRBuilder<> &B, const FuncInfo &FInfo);
 
-  bool insertSinCos(CallInst *Sin, CallInst *Cos, IRBuilder<> &B,
-                    const FuncInfo &FInfo);
+  /// Insert a value to sincos function \p Fsincos. Returns (value of sin, value
+  /// of cos, sincos call).
+  std::tuple<Value *, Value *, Value *> insertSinCos(Value *Arg,
+                                                     FastMathFlags FMF,
+                                                     IRBuilder<> &B,
+                                                     FunctionCallee Fsincos);
 
   // sin/cos
   bool fold_sincos(FPMathOperator *FPOp, IRBuilder<> &B, const FuncInfo &FInfo);
@@ -1041,40 +1045,24 @@ bool AMDGPULibCalls::fold_sqrt(FPMathOperator *FPOp, IRBuilder<> &B,
   return false;
 }
 
-bool AMDGPULibCalls::insertSinCos(CallInst *Sin, CallInst *Cos, IRBuilder<> &B,
-                                  const FuncInfo &fInfo) {
-  Value *Arg = Sin->getOperand(0);
-  assert(Arg == Cos->getOperand(0));
-
+std::tuple<Value *, Value *, Value *>
+AMDGPULibCalls::insertSinCos(Value *Arg, FastMathFlags FMF, IRBuilder<> &B,
+                             FunctionCallee Fsincos) {
+  DebugLoc DL = B.getCurrentDebugLocation();
   Function *F = B.GetInsertBlock()->getParent();
-  Module *M = F->getParent();
-  // Merge the sin and cos.
-
-  // for OpenCL 2.0 we have only generic implementation of sincos
-  // function.
-  // FIXME: This is not true anymore
-  AMDGPULibFunc nf(AMDGPULibFunc::EI_SINCOS, fInfo);
-  nf.getLeads()[0].PtrKind =
-      AMDGPULibFunc::getEPtrKindFromAddrSpace(AMDGPUAS::FLAT_ADDRESS);
-  FunctionCallee Fsincos = getFunction(M, nf);
-  if (!Fsincos)
-    return false;
-
   B.SetInsertPointPastAllocas(F);
 
-  DILocation *MergedDebugLoc =
-      DILocation::getMergedLocation(Sin->getDebugLoc(), Cos->getDebugLoc());
-  B.SetCurrentDebugLocation(MergedDebugLoc);
-
-  AllocaInst *Alloc = B.CreateAlloca(Sin->getType(), nullptr, "__sincos_");
+  AllocaInst *Alloc = B.CreateAlloca(Arg->getType(), nullptr, "__sincos_");
 
   if (Instruction *ArgInst = dyn_cast<Instruction>(Arg)) {
     // If the argument is an instruction, it must dominate all uses so put our
     // sincos call there. Otherwise, right after the allocas works well enough
     // if it's an argument or constant.
 
     B.SetInsertPoint(ArgInst->getParent(), ++ArgInst->getIterator());
-    B.SetCurrentDebugLocation(MergedDebugLoc);
+
+    // SetInsertPoint unwelcomely always tries to set the debug loc.
+    B.SetCurrentDebugLocation(DL);
   }
 
   Value *P = Alloc;
@@ -1085,25 +1073,12 @@ bool AMDGPULibCalls::insertSinCos(CallInst *Sin, CallInst *Cos, IRBuilder<> &B,
   if (PTy->getPointerAddressSpace() != AMDGPUAS::PRIVATE_ADDRESS)
     P = B.CreateAddrSpaceCast(Alloc, PTy);
 
-  // Intersect the two sets of flags.
-  FastMathFlags FMF = cast<FPMathOperator>(Sin)->getFastMathFlags();
-  FMF &= cast<FPMathOperator>(Cos)->getFastMathFlags();
-  B.setFastMathFlags(FMF);
-
-  CallInst *Call = CreateCallEx2(B, Fsincos, Arg, P);
-  LoadInst *Reload = B.CreateLoad(Alloc->getAllocatedType(), Alloc);
-  Reload->setDebugLoc(Cos->getDebugLoc());
-
-  LLVM_DEBUG(errs() << "AMDIC: fold_sincos (" << *Sin << ", " << *Cos
-                    << ") with " << *Call << '\n');
-
-  Sin->replaceAllUsesWith(Call);
-  Sin->eraseFromParent();
-
-  Cos->replaceAllUsesWith(Reload);
-  Cos->eraseFromParent();
+  CallInst *SinCos = CreateCallEx2(B, Fsincos, Arg, P);
 
-  return true;
+  // TODO: Is it worth trying to preserve the location for the cos calls for the
+  // load?
+  LoadInst *LoadCos = B.CreateLoad(Alloc->getAllocatedType(), Alloc);
+  return {SinCos, LoadCos, SinCos};
 }
 
 // fold sin, cos -> sincos.
@@ -1121,33 +1096,92 @@ bool AMDGPULibCalls::fold_sincos(FPMathOperator *FPOp, IRBuilder<> &B,
 
   Value *CArgVal = FPOp->getOperand(0);
   CallInst *CI = cast<CallInst>(FPOp);
-  bool Changed = false;
 
+  Function *F = B.GetInsertBlock()->getParent();
+  Module *M = F->getParent();
+
+  // Merge the sin and cos.
+
+  // for OpenCL 2.0 we have only generic implementation of sincos
+  // function.
+  // FIXME: This is not true anymore
+  AMDGPULibFunc SinCosLibFunc(AMDGPULibFunc::EI_SINCOS, fInfo);
+  SinCosLibFunc.getLeads()[0].PtrKind =
+      AMDGPULibFunc::getEPtrKindFromAddrSpace(AMDGPUAS::FLAT_ADDRESS);
+  FunctionCallee FSinCos = getFunction(M, SinCosLibFunc);
+  if (!FSinCos)
+    return false;
+
+  SmallVector<CallInst *> SinCalls;
+  SmallVector<CallInst *> CosCalls;
+  SmallVector<CallInst *> SinCosCalls;
   FuncInfo PartnerInfo(isSin ? AMDGPULibFunc::EI_COS : AMDGPULibFunc::EI_SIN,
                        fInfo);
   const std::string PairName = PartnerInfo.mangle();
 
-  CallInst *UI = nullptr;
+  StringRef SinName = isSin ? CI->getCalledFunction()->getName() : PairName;
+  StringRef CosName = isSin ? PairName : CI->getCalledFunction()->getName();
+  const std::string SinCosName = SinCosLibFunc.mangle();
+
+  // Intersect the two sets of flags.
+  FastMathFlags FMF = FPOp->getFastMathFlags();
+  MDNode *FPMath = CI->getMetadata(LLVMContext::MD_fpmath);
+
+  SmallVector<DILocation *> MergeDbgLocs = {CI->getDebugLoc()};
 
-  // TODO: Handle repeated uses, the generic implementation does.
   for (User* U : CArgVal->users()) {
     CallInst *XI = dyn_cast<CallInst>(U);
-    if (!XI || XI->isNoBuiltin())
+    if (!XI || XI->getFunction() != F || XI->isNoBuiltin())
       continue;
 
     Function *UCallee = XI->getCalledFunction();
-    if (UCallee && UCallee->getName().equals(PairName))
-      UI = XI;
-    else if (UI)
-      return Changed;
+    if (!UCallee)
+      continue;
+
+    bool Handled = true;
+
+    if (UCallee->getName() == SinName)
+      SinCalls.push_back(XI);
+    else if (UCallee->getName() == CosName)
+      CosCalls.push_back(XI);
+    else if (UCallee->getName() == SinCosName)
+      SinCosCalls.push_back(XI);
+    else
+      Handled = false;
+
+    if (Handled) {
+      MergeDbgLocs.push_back(XI->getDebugLoc());
+      auto *OtherOp = cast<FPMathOperator>(XI);
+      FMF &= OtherOp->getFastMathFlags();
+      FPMath = MDNode::getMostGenericFPMath(
+          FPMath, XI->getMetadata(LLVMContext::MD_fpmath));
+    }
   }
 
-  if (!UI)
-    return Changed;
+  if (SinCalls.empty() || CosCalls.empty())
+    return false;
+
+  B.setFastMathFlags(FMF);
+  B.setDefaultFPMathTag(FPMath);
+  DILocation *DbgLoc = DILocation::getMergedLocations(MergeDbgLocs);
+  B.SetCurrentDebugLocation(DbgLoc);
+
+  auto [Sin, Cos, SinCos] = insertSinCos(CArgVal, FMF, B, FSinCos);
+
+  auto replaceTrigInsts = [](ArrayRef<CallInst *> Calls, Value *Res) {
+    for (CallInst *C : Calls)
+      C->replaceAllUsesWith(Res);
+
+    // Leave the other dead instructions to avoid clobbering iterators.
+  };
 
-  CallInst *Sin = isSin ? CI : UI;
-  CallInst *Cos = isSin ? UI : CI;
-  return insertSinCos(Sin, Cos, B, fInfo) || Changed;
+  replaceTrigInsts(SinCalls, Sin);
+  replaceTrigInsts(CosCalls, Cos);
+  replaceTrigInsts(SinCosCalls, SinCos);
+
+  // It's safe to delete the original now.
+  CI->eraseFromParent();
+  return true;
 }
 
 bool AMDGPULibCalls::evaluateScalarMathFunc(const FuncInfo &FInfo,

llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-sincos.defined.ll

@@ -110,6 +110,7 @@ define void @sincos_f32(float %x, ptr addrspace(1) nocapture writeonly %sin_out,
 ; CHECK-NEXT:    [[TMP1:%.*]] = call contract float @_Z6sincosfPU3AS0f(float [[X]], ptr [[TMP0]])
 ; CHECK-NEXT:    [[TMP2:%.*]] = load float, ptr addrspace(5) [[__SINCOS_]], align 4
 ; CHECK-NEXT:    store float [[TMP1]], ptr addrspace(1) [[SIN_OUT]], align 4
+; CHECK-NEXT:    [[CALL1:%.*]] = tail call contract float @_Z3cosf(float [[X]])
 ; CHECK-NEXT:    store float [[TMP2]], ptr addrspace(1) [[COS_OUT]], align 4
 ; CHECK-NEXT:    ret void
 ;
@@ -130,6 +131,7 @@ define void @sincos_f32_value_is_same_constantfp(ptr addrspace(1) nocapture writ
 ; CHECK-NEXT:    [[TMP1:%.*]] = call contract float @_Z6sincosfPU3AS0f(float 4.200000e+01, ptr [[TMP0]])
 ; CHECK-NEXT:    [[TMP2:%.*]] = load float, ptr addrspace(5) [[__SINCOS_]], align 4
 ; CHECK-NEXT:    store float [[TMP1]], ptr addrspace(1) [[SIN_OUT]], align 4
+; CHECK-NEXT:    [[CALL1:%.*]] = tail call contract float @_Z3cosf(float 4.200000e+01)
 ; CHECK-NEXT:    store float [[TMP2]], ptr addrspace(1) [[COS_OUT]], align 4
 ; CHECK-NEXT:    ret void
 ;
@@ -159,6 +161,7 @@ define void @sincos_v2f32(<2 x float> %x, ptr addrspace(1) nocapture writeonly %
 ; CHECK-NEXT:    [[TMP1:%.*]] = call contract <2 x float> @_Z6sincosDv2_fPU3AS0S_(<2 x float> [[X]], ptr [[TMP0]])
 ; CHECK-NEXT:    [[TMP2:%.*]] = load <2 x float>, ptr addrspace(5) [[__SINCOS_]], align 8
 ; CHECK-NEXT:    store <2 x float> [[TMP1]], ptr addrspace(1) [[SIN_OUT]], align 8
+; CHECK-NEXT:    [[CALL1:%.*]] = tail call contract <2 x float> @_Z3cosDv2_f(<2 x float> [[X]])
 ; CHECK-NEXT:    store <2 x float> [[TMP2]], ptr addrspace(1) [[COS_OUT]], align 8
 ; CHECK-NEXT:    ret void
 ;

llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-sincos.defined.nobuiltin.ll

@@ -60,6 +60,7 @@ define void @sincos_f32(float %x, ptr addrspace(1) nocapture writeonly %sin_out,
 ; CHECK-NEXT:    [[TMP0:%.*]] = call contract float @_Z6sincosfPU3AS0f(float [[X]], ptr [[__SINCOS_]])
 ; CHECK-NEXT:    [[TMP1:%.*]] = load float, ptr [[__SINCOS_]], align 4
 ; CHECK-NEXT:    store float [[TMP0]], ptr addrspace(1) [[SIN_OUT]], align 4
+; CHECK-NEXT:    [[CALL1:%.*]] = tail call contract float @_Z3cosf(float [[X]])
 ; CHECK-NEXT:    store float [[TMP1]], ptr addrspace(1) [[COS_OUT]], align 4
 ; CHECK-NEXT:    ret void
 ;