Generate native sqrt for fast llvm sqrt operation and match reciprocal sqrt

Author: azabazno

IGC/VectorCompiler/lib/GenXCodeGen/GenXLowering.cpp

@@ -2494,8 +2494,14 @@ bool GenXLowering::lowerSqrt(CallInst *CI) {
   IGC_ASSERT_MESSAGE(GenXIntrinsic::getAnyIntrinsicID(CI) == Intrinsic::sqrt,
     "llvm.sqrt expected");
   auto *ResTy = CI->getType();
-  auto *SqrtDecl = GenXIntrinsic::getGenXDeclaration(
-      CI->getModule(), GenXIntrinsic::genx_ieee_sqrt, {ResTy, ResTy});
+
+  GenXIntrinsic::ID SqrtID = (CI->getType()->getScalarType()->isDoubleTy() ||
+                              !CI->getFastMathFlags().isFast())
+                                 ? GenXIntrinsic::genx_ieee_sqrt
+                                 : GenXIntrinsic::genx_sqrt;
+
+  auto *SqrtDecl =
+      GenXIntrinsic::getGenXDeclaration(CI->getModule(), SqrtID, {ResTy});
   Value *Result = IRBuilder<>(CI).CreateCall(SqrtDecl, {CI->getArgOperand(0)},
                                              CI->getName());
   CI->replaceAllUsesWith(Result);

IGC/VectorCompiler/lib/GenXCodeGen/GenXPatternMatch.cpp

@@ -166,6 +166,7 @@ class GenXPatternMatch : public FunctionPass,
   // flipBoolNot : flip a (vector) bool not instruction if beneficial
   bool flipBoolNot(Instruction *Inst);
   // foldBoolAnd : fold a (vector) bool and into sel/wrregion if beneficial
+  bool matchInverseSqrt(Instruction *I);
   bool foldBoolAnd(Instruction *Inst);
   bool simplifyPredRegion(CallInst *Inst);
   bool simplifyWrRegion(CallInst *Inst);
@@ -398,6 +399,9 @@ void GenXPatternMatch::visitCallInst(CallInst &I) {
   switch (unsigned ID = GenXIntrinsic::getGenXIntrinsicID(&I)) {
   default:
     break;
+  case GenXIntrinsic::genx_inv:
+    Changed |= matchInverseSqrt(&I);
+    break;
   case GenXIntrinsic::genx_ssadd_sat:
   case GenXIntrinsic::genx_suadd_sat:
   case GenXIntrinsic::genx_usadd_sat:
@@ -768,6 +772,42 @@ bool GenXPatternMatch::flipBoolNot(Instruction *Inst) {
   return true;
 }
 
+/// (inv (sqrt x)) -> (rsqrt x)
+bool GenXPatternMatch::matchInverseSqrt(Instruction *I) {
+  IGC_ASSERT(I);
+
+  auto *OpInst = dyn_cast<CallInst>(I->getOperand(0));
+  if (!OpInst)
+    return false;
+
+  // Leave as it is for double types
+  if (OpInst->getType()->getScalarType()->isDoubleTy())
+    return false;
+
+  // Generate inverse sqrt only if fast flag for llvm intrinsic is used or
+  // genx sqrt intrinsics is specified
+  auto IID = GenXIntrinsic::getAnyIntrinsicID(OpInst);
+  if (!(IID == GenXIntrinsic::genx_sqrt ||
+        (IID == Intrinsic::sqrt && OpInst->getFastMathFlags().isFast())))
+    return false;
+
+  // Leave as it if sqrt has multiple uses:
+  // generating rsqrt operation is not beneficial
+  if (OpInst->getNumUses() > 1)
+    return false;
+
+  Function *Decl = GenXIntrinsic::getGenXDeclaration(
+      I->getModule(), GenXIntrinsic::genx_rsqrt, {I->getType()});
+  auto NewInst =
+      CallInst::Create(Decl, {OpInst->getOperand(0)},
+                       OpInst->getName() + "inversed", I->getNextNode());
+  I->replaceAllUsesWith(NewInst);
+  I->eraseFromParent();
+
+  OpInst->eraseFromParent();
+  return true;
+}
+
 /***********************************************************************
  * foldBoolAnd : fold a (vector) bool and into sel/wrregion if beneficial
  *