From 04b3725e0dd5f7c2822a5164476318094d76aa17 Mon Sep 17 00:00:00 2001
From: Arpit Jaiswal <arjaiswal@nvidia.com>
Date: Fri, 29 Dec 2023 00:02:30 +0000
Subject: [PATCH] Fix for identifying builtin sqrt

---
 tools/cgeist/Lib/CGCall.cc     | 16 ++++++++++++++++
 tools/cgeist/Lib/clang-mlir.cc | 14 +-------------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/tools/cgeist/Lib/CGCall.cc b/tools/cgeist/Lib/CGCall.cc
index c5e4d40ca2d9..695d00b56edb 100644
--- a/tools/cgeist/Lib/CGCall.cc
+++ b/tools/cgeist/Lib/CGCall.cc
@@ -541,6 +541,22 @@ MLIRScanner::EmitClangBuiltinCallExpr(clang::CallExpr *expr) {
         ValueCategory(builder.create<arith::ConstantIntOp>(loc, 0, resultType),
                       /*isRef*/ false));
   }
+case Builtin::BIsqrt:
+  case Builtin::BIsqrtf:
+  case Builtin::BIsqrtl:
+  case Builtin::BI__builtin_sqrt:
+  case Builtin::BI__builtin_sqrtf:
+  case Builtin::BI__builtin_sqrtf16:
+  case Builtin::BI__builtin_sqrtl:
+  case Builtin::BI__builtin_sqrtf128:
+  case Builtin::BI__builtin_elementwise_sqrt: {
+    auto v = Visit(expr->getArg(0));
+    assert(!v.isReference);
+    Value res = builder.create<math::SqrtOp>(loc, v.val);
+    auto postTy = getMLIRType(expr->getType());
+    return success(
+        ValueCategory(res, /*isRef*/ false));
+  }
   case Builtin::BI__builtin_clzs:
   case Builtin::BI__builtin_clz:
   case Builtin::BI__builtin_clzl:
diff --git a/tools/cgeist/Lib/clang-mlir.cc b/tools/cgeist/Lib/clang-mlir.cc
index a3f14408b2c4..806810d31196 100644
--- a/tools/cgeist/Lib/clang-mlir.cc
+++ b/tools/cgeist/Lib/clang-mlir.cc
@@ -1824,19 +1824,7 @@ MLIRScanner::EmitBuiltinOps(clang::CallExpr *expr) {
                           /*isReference*/ false),
             true);
       }
-      if (sr->getDecl()->getIdentifier() &&
-          (sr->getDecl()->getName() == "sqrtf" ||
-           sr->getDecl()->getName() == "sqrt")) {
-        std::vector<mlir::Value> args;
-        for (auto a : expr->arguments()) {
-          args.push_back(Visit(a).getValue(loc, builder));
-        }
-        return make_pair(
-            ValueCategory(builder.create<mlir::math::SqrtOp>(loc, args[0]),
-                          /*isReference*/ false),
-            true);
-      }
-      if (sr->getDecl()->getIdentifier() &&
+            if (sr->getDecl()->getIdentifier() &&
           (sr->getDecl()->getName() == "expf" ||
            sr->getDecl()->getName() == "exp")) {
         std::vector<mlir::Value> args;