[CINN] Fix kernel type errors (atan, pow, exp, bitwise_xxx) (#72025)

ggggxm · web-flow · commit 597fbc33911b · 2025-04-07T10:00:42.000+08:00
diff --git a/paddle/cinn/backends/llvm/llvm_intrin_rule.h b/paddle/cinn/backends/llvm/llvm_intrin_rule.h
@@ -72,7 +72,6 @@ void RegisterCpuIntrinRule() {
 #define __(intrin_name__, id)                                         \
   ir::Registry::Register("lower_cpu_intrinsic_" #intrin_name__, true) \
       .SetBody(MakeFloatIntrinOp<id, 1>);
-  __(exp, ::llvm::Intrinsic::exp)
   __(exp2, ::llvm::Intrinsic::exp2)
   __(sqrt, ::llvm::Intrinsic::sqrt)
   __(log, ::llvm::Intrinsic::log)
@@ -102,7 +101,53 @@ void RegisterCpuIntrinRule() {
       .SetBody(MakeFloatIntrinOp<::llvm::Intrinsic::fmuladd, 3, false>);
 
   ir::Registry::Register("lower_cpu_intrinsic_pow", true)
-      .SetBody(MakeFloatIntrinOp<::llvm::Intrinsic::pow, 2, false>);
+      .SetBody([](lang::Args args, lang::RetValue *rv) {
+        PADDLE_ENFORCE_GE(args.size(),
+                          1U,
+                          ::common::errors::InvalidArgument(
+                              "The number of args should be greater than 1."));
+        Expr arg0 = args[0];
+        ir::Call *node = arg0->as<ir::Call>();
+        PADDLE_ENFORCE_NOT_NULL(
+            node,
+            ::common::errors::InvalidArgument(
+                "The argument must be a valid call expression."));
+        PADDLE_ENFORCE_EQ(
+            node->read_args.size(),
+            2U,
+            ::common::errors::InvalidArgument(
+                "The number of read arguments of 'pow' should be 2."));
+        Type type = common::Float(64);
+        Expr casted_arg0 = ir::Cast::Make(type, node->read_args[0]);
+        Expr casted_arg1 = ir::Cast::Make(type, node->read_args[1]);
+        Expr log_base = lang::Log(casted_arg0);
+        Expr product = log_base * casted_arg1;
+        *rv = ir::Cast::Make(node->type(), lang::Exp(product));
+      });
+
+  ir::Registry::Register("lower_cpu_intrinsic_exp", true)
+      .SetBody([](lang::Args args, lang::RetValue *rv) {
+        PADDLE_ENFORCE_GE(args.size(),
+                          1U,
+                          ::common::errors::InvalidArgument(
+                              "The number of args should be greater than 1."));
+        Expr arg0 = args[0];
+        ir::Call *node = arg0->as<ir::Call>();
+        PADDLE_ENFORCE_NOT_NULL(
+            node,
+            ::common::errors::InvalidArgument(
+                "The argument must be a valid call expression."));
+        PADDLE_ENFORCE_EQ(
+            node->read_args.size(),
+            1,
+            ::common::errors::InvalidArgument(
+                "The number of read arguments of 'exp' should be 1."));
+        Type type = common::Float(64);
+        Expr casted_arg0 = ir::Cast::Make(type, node->read_args[0]);
+        Expr exp = ir::intrinsics::BuiltinIntrin::Make(
+            node->name, casted_arg0, ::llvm::Intrinsic::exp, 1, type);
+        *rv = ir::Cast::Make(node->type(), exp);
+      });
 
   ir::Registry::Register("lower_cpu_intrinsic_bitwise_not", true)
       .SetBody(MakeFloatIntrinOp<-1, 1, false>);
diff --git a/paddle/cinn/ir/op/ir_operators.cc b/paddle/cinn/ir/op/ir_operators.cc
@@ -102,7 +102,9 @@ Expr BitwiseOrCallImpl(common::UnknownArch,
 }
 
 Expr BitwiseOrCallImpl(common::X86Arch, const Target &target, Expr a, Expr b) {
-  return lang::CallExtern("bitwise_or", {a, b}, {{"vectorizable", false}});
+  Expr c = lang::CallExtern("bitwise_or", {a, b}, {{"vectorizable", false}});
+  c->set_type(a.type());
+  return c;
 }
 
 Expr BitwiseOrCallImpl(common::ARMArch, const Target &target, Expr a, Expr b) {
@@ -176,7 +178,9 @@ Expr BitwiseAndCallImpl(common::UnknownArch,
 }
 
 Expr BitwiseAndCallImpl(common::X86Arch, const Target &target, Expr a, Expr b) {
-  return lang::CallExtern("bitwise_and", {a, b}, {{"vectorizable", false}});
+  Expr c = lang::CallExtern("bitwise_and", {a, b}, {{"vectorizable", false}});
+  c->set_type(a.type());
+  return c;
 }
 
 Expr BitwiseAndCallImpl(common::ARMArch, const Target &target, Expr a, Expr b) {
@@ -250,7 +254,9 @@ Expr BitwiseXorCallImpl(common::UnknownArch,
 }
 
 Expr BitwiseXorCallImpl(common::X86Arch, const Target &target, Expr a, Expr b) {
-  return lang::CallExtern("bitwise_xor", {a, b}, {{"vectorizable", false}});
+  Expr c = lang::CallExtern("bitwise_xor", {a, b}, {{"vectorizable", false}});
+  c->set_type(a.type());
+  return c;
 }
 
 Expr BitwiseXorCallImpl(common::ARMArch, const Target &target, Expr a, Expr b) {
@@ -321,7 +327,9 @@ Expr BitwiseNotCallImpl(common::UnknownArch, const Target &target, Expr a) {
 }
 
 Expr BitwiseNotCallImpl(common::X86Arch, const Target &target, Expr a) {
-  return lang::CallExtern("bitwise_not", {a}, {{"vectorizable", false}});
+  Expr c = lang::CallExtern("bitwise_not", {a}, {{"vectorizable", false}});
+  c->set_type(a->type());
+  return c;
 }
 
 Expr BitwiseNotCallImpl(common::ARMArch, const Target &target, Expr a) {
diff --git a/paddle/cinn/optim/map_extern_call.cc b/paddle/cinn/optim/map_extern_call.cc
@@ -41,11 +41,14 @@ static const std::set<std::string> kExternInt32CallsGPU{{"left_shift",
                                                          "popc",
                                                          "mod"}};
 
+static const std::set<std::string> kExternFp64CallsCPU{"atan"};
+
 static const std::set<std::string> kExternFp32CallsCPU = {
     "erf", "acos", "acosh", "asin", "asinh", "atan", "atanh", "remainder"};
 
 void DealWithCpuIntrinsics(ir::Call *node, Expr *expr) {
-  if (kExternFp32CallsCPU.count(node->name)) {
+  if (kExternFp32CallsCPU.count(node->name) ||
+      kExternFp64CallsCPU.count(node->name)) {
     PADDLE_ENFORCE_GE(
         node->read_args.size(),
         1UL,
@@ -62,6 +65,11 @@ void DealWithCpuIntrinsics(ir::Call *node, Expr *expr) {
       auto out_type = node->type();
       *expr = lang::CallExtern(node->name + "f", node->read_args);
     }
+    if (node->read_args.front().type().is_float(64)) {
+      auto out_type = node->type();
+      *expr = lang::CallExtern("cinn_host_" + node->name + "_fp64",
+                               node->read_args);
+    }
   }
 }
 
diff --git a/paddle/cinn/runtime/cpu/host_intrinsics.cc b/paddle/cinn/runtime/cpu/host_intrinsics.cc
@@ -244,6 +244,8 @@ inline double FN_FP64(cbrt)(double x) { return cbrt(x); }
 
 inline double FN_FP64(pow)(double x, double y) { return pow(x, y); }
 
+inline double FN_FP64(atan)(double x) { return atan(x); }
+
 #undef FN_FP64
 
 #define FN_INT32(func) cinn_host_##func##_int32
@@ -302,6 +304,7 @@ CINN_REGISTER_HELPER(host_intrinsics) {
   REGISTER_EXTERN_FUNC_1_IN_1_OUT(func__, host_target, double, double);
 
   REGISTER_EXTERN_FUNC_1_IN_1_OUT_FP64(cinn_host_cbrt_fp64);
+  REGISTER_EXTERN_FUNC_1_IN_1_OUT_FP64(cinn_host_atan_fp64);
 
 #undef REGISTER_EXTERN_FUNC_1_IN_1_OUT_FP64
 
diff --git a/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh b/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh
@@ -228,6 +228,10 @@ __device__ inline int FN_INT32(mod)(int a, int b) {
   if ((res != 0) && ((b ^ res) < 0)) res += b;
   return res;
 }
+__device__ inline int FN_INT32(exp)(int a) {
+  float res = exp(__int2float_rd(a));
+  return __float2int_rn(res);
+}
 
 // *************************************************************** //
 
@@ -247,6 +251,10 @@ __device__ inline long long int FN_INT64(mod)(long long int a, long long int b)
   if ((res != 0) && ((b ^ res) < 0)) res += b;
   return res;
 }
+__device__ inline long long int FN_INT64(exp)(long long int a) {
+  double res = exp(__ll2double_rd(a));
+  return __double2ll_rn(res);
+}
 
 __device__ inline long long int FN_INT64(pow)(long long int a, long long int b) {
   double res = pow(__ll2double_rd(a), __ll2double_rd(b));
diff --git a/paddle/cinn/runtime/cuda/cuda_intrinsics.cc b/paddle/cinn/runtime/cuda/cuda_intrinsics.cc
@@ -222,6 +222,7 @@ CINN_REGISTER_HELPER(cuda_intrinsics) {
   REGISTER_EXTERN_FUNC_1_IN_1_INT32(clz)
   REGISTER_EXTERN_FUNC_1_IN_1_INT32(popc)
   REGISTER_EXTERN_FUNC_1_IN_1_INT32(trunc)
+  REGISTER_EXTERN_FUNC_1_IN_1_INT32(exp)
 
 #undef REGISTER_EXTERN_FUNC_1_IN_1_INT32
 
@@ -233,6 +234,7 @@ CINN_REGISTER_HELPER(cuda_intrinsics) {
   REGISTER_EXTERN_FUNC_1_IN_1_INT64(clz)
   REGISTER_EXTERN_FUNC_1_IN_1_INT64(popc)
   REGISTER_EXTERN_FUNC_1_IN_1_INT64(trunc)
+  REGISTER_EXTERN_FUNC_1_IN_1_INT64(exp)
 
 #undef REGISTER_EXTERN_FUNC_1_IN_1_INT64
 

Original file line number	Diff line number	Diff line change
`@@ -102,7 +102,9 @@ Expr BitwiseOrCallImpl(common::UnknownArch,`
`102`	`102`	`}`
`103`	`103`
`104`	`104`	`Expr BitwiseOrCallImpl(common::X86Arch, const Target &target, Expr a, Expr b) {`
`105`		`- return lang::CallExtern("bitwise_or", {a, b}, {{"vectorizable", false}});`
	`105`	`+ Expr c = lang::CallExtern("bitwise_or", {a, b}, {{"vectorizable", false}});`
	`106`	`+ c->set_type(a.type());`
	`107`	`+ return c;`
`106`	`108`	`}`
`107`	`109`
`108`	`110`	`Expr BitwiseOrCallImpl(common::ARMArch, const Target &target, Expr a, Expr b) {`
`@@ -176,7 +178,9 @@ Expr BitwiseAndCallImpl(common::UnknownArch,`
`176`	`178`	`}`
`177`	`179`
`178`	`180`	`Expr BitwiseAndCallImpl(common::X86Arch, const Target &target, Expr a, Expr b) {`
`179`		`- return lang::CallExtern("bitwise_and", {a, b}, {{"vectorizable", false}});`
	`181`	`+ Expr c = lang::CallExtern("bitwise_and", {a, b}, {{"vectorizable", false}});`
	`182`	`+ c->set_type(a.type());`
	`183`	`+ return c;`
`180`	`184`	`}`
`181`	`185`
`182`	`186`	`Expr BitwiseAndCallImpl(common::ARMArch, const Target &target, Expr a, Expr b) {`
`@@ -250,7 +254,9 @@ Expr BitwiseXorCallImpl(common::UnknownArch,`
`250`	`254`	`}`
`251`	`255`
`252`	`256`	`Expr BitwiseXorCallImpl(common::X86Arch, const Target &target, Expr a, Expr b) {`
`253`		`- return lang::CallExtern("bitwise_xor", {a, b}, {{"vectorizable", false}});`
	`257`	`+ Expr c = lang::CallExtern("bitwise_xor", {a, b}, {{"vectorizable", false}});`
	`258`	`+ c->set_type(a.type());`
	`259`	`+ return c;`
`254`	`260`	`}`
`255`	`261`
`256`	`262`	`Expr BitwiseXorCallImpl(common::ARMArch, const Target &target, Expr a, Expr b) {`
`@@ -321,7 +327,9 @@ Expr BitwiseNotCallImpl(common::UnknownArch, const Target &target, Expr a) {`
`321`	`327`	`}`
`322`	`328`
`323`	`329`	`Expr BitwiseNotCallImpl(common::X86Arch, const Target &target, Expr a) {`
`324`		`- return lang::CallExtern("bitwise_not", {a}, {{"vectorizable", false}});`
	`330`	`+ Expr c = lang::CallExtern("bitwise_not", {a}, {{"vectorizable", false}});`
	`331`	`+ c->set_type(a->type());`
	`332`	`+ return c;`
`325`	`333`	`}`
`326`	`334`
`327`	`335`	`Expr BitwiseNotCallImpl(common::ARMArch, const Target &target, Expr a) {`