Merge branch 'develop' into add_op_test

PaddlePaddle · Nov 18, 2021 · 2bce9c8 · 2bce9c8
2 parents 3221e6a + 456c125
commit 2bce9c8
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 12 deletions.
diff --git a/cinn/hlir/pe/pe_elementwise_test.cc b/cinn/hlir/pe/pe_elementwise_test.cc
@@ -111,7 +111,7 @@ void TestElementwisePE(const std::string &fn_name,
 bool isnan(float e) { return std::isnan(e); }
 bool isfinite(float e) { return std::isfinite(e); }
 bool isinf(float e) { return std::isinf(e); }
-bool rsqrt(float e) { return 1.0f / std::sqrt(e); }
+float rsqrt(float e) { return 1.0f / std::sqrt(e); }
 
 #define TEST_ELEMENTWISE_PE_FP32(test_name__, PE__) \
   TEST(elementwise_pe, test_name__) { TestElementwisePE("PE_Elementwise_" #test_name__ "_fp32", PE__, test_name__); }

diff --git a/cinn/lang/builtin.cc b/cinn/lang/builtin.cc
@@ -57,6 +57,7 @@ Expr logic_or(const std::vector<Expr>& conds) {
 EXTERN_CALL_IMP(Exp, exp);
 EXTERN_CALL_IMP_NO_VEC(Erf, erf);
 EXTERN_CALL_IMP(Sqrt, sqrt);
+EXTERN_CALL_IMP(Rsqrt, rsqrt);
 EXTERN_CALL_IMP(Log, log);
 EXTERN_CALL_IMP(Log2, log2);
 EXTERN_CALL_IMP(Log10, log10);
@@ -77,16 +78,6 @@ EXTERN_CALL_IMP_NO_VEC(Asinh, asinh);
 EXTERN_CALL_IMP_NO_VEC(Atan, atan);
 EXTERN_CALL_IMP_NO_VEC(Atanh, atanh);
 
-#ifdef CINN_WITH_CUDA
-EXTERN_CALL_IMP(Rsqrt, rsqrt);
-#else
-Expr Rsqrt(Expr e) {
-  // TODO(Xreki): add the rsqrt's intrinsic of X86 and remove the ifdef.
-  auto one = make_const(e->type(), 1);
-  return one / Sqrt(e);
-}
-#endif
-
 Expr min_value(const Type& type) {
   CHECK_EQ(type.lanes(), 1);
 #define FOR_CASE(type__)                                \

diff --git a/cinn/optim/lower_intrin.h b/cinn/optim/lower_intrin.h
@@ -26,7 +26,7 @@ static const std::set<std::string> kIntrinsicCalls{
     {"exp",         "exp2",       "sqrt",        "log",         "log2",        "log10", "floor",
      "ceil",        "round",      "trunc",       "cos",         "cosh",        "tan",   "tanh",
      "sin",         "sinh",       "fabs",        "isnan",       "isfinite",    "isinf", "left_shift",
-     "right_shift", "bitwise_or", "bitwise_and", "bitwise_xor", "bitwise_not", "fma"}};
+     "right_shift", "bitwise_or", "bitwise_and", "bitwise_xor", "bitwise_not", "fma",   "rsqrt"}};
 
 /**
  * Map the Call nodes to llvm intrinsic.