[CINN] Add support for mod and pow in x86 kernel (#70144)

lshpku · web-flow · commit 6fc942c7d800 · 2024-12-12T14:10:06.000+08:00
diff --git a/paddle/cinn/backends/llvm/llvm_intrin_rule.h b/paddle/cinn/backends/llvm/llvm_intrin_rule.h
@@ -91,12 +91,18 @@ void RegisterCpuIntrinRule() {
 #define RegisterBitwise(intrin_name__)                                \
   ir::Registry::Register("lower_cpu_intrinsic_" #intrin_name__, true) \
       .SetBody(MakeFloatIntrinOp<-1, 2, false>);
-  RegisterBitwise(bitwise_or) RegisterBitwise(bitwise_xor) RegisterBitwise(
-      bitwise_and) RegisterBitwise(left_shift) RegisterBitwise(right_shift)
+  RegisterBitwise(bitwise_or);
+  RegisterBitwise(bitwise_xor);
+  RegisterBitwise(bitwise_and);
+  RegisterBitwise(left_shift);
+  RegisterBitwise(right_shift);
 #undef RegisterBitwise
 
-      ir::Registry::Register("lower_cpu_intrinsic_fma", true)
-          .SetBody(MakeFloatIntrinOp<::llvm::Intrinsic::fmuladd, 3, false>);
+  ir::Registry::Register("lower_cpu_intrinsic_fma", true)
+      .SetBody(MakeFloatIntrinOp<::llvm::Intrinsic::fmuladd, 3, false>);
+
+  ir::Registry::Register("lower_cpu_intrinsic_pow", true)
+      .SetBody(MakeFloatIntrinOp<::llvm::Intrinsic::pow, 2, false>);
 
   ir::Registry::Register("lower_cpu_intrinsic_bitwise_not", true)
       .SetBody(MakeFloatIntrinOp<-1, 1, false>);
@@ -298,6 +304,28 @@ void RegisterCpuIntrinRule() {
         *rv = (lang::Exp(arg) - lang::Exp(arg * make_const(arg->type(), -1))) /
               make_const(arg->type(), 2);
       });
+
+  ir::Registry::Register("lower_cpu_intrinsic_mod", 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. Received null."));
+        PADDLE_ENFORCE_EQ(node->read_args.size(),
+                          2UL,
+                          ::common::errors::InvalidArgument(
+                              "The 'mod' op must have exactly 2 read_args."));
+
+        Expr lhs = node->read_args[0];
+        Expr rhs = node->read_args[1];
+        *rv = lhs % rhs;
+      });
 }
 }  // namespace codegen
 }  // namespace cinn
diff --git a/paddle/cinn/optim/lower_intrin.h b/paddle/cinn/optim/lower_intrin.h
@@ -28,7 +28,7 @@ static const std::set<std::string> kIntrinsicCalls{
      "cos",         "cosh",        "tan",        "tanh",        "sin",
      "sinh",        "fabs",        "isnan",      "isfinite",    "isinf",
      "left_shift",  "right_shift", "bitwise_or", "bitwise_and", "bitwise_xor",
-     "bitwise_not", "fma",         "rsqrt"}};
+     "bitwise_not", "fma",         "rsqrt",      "mod",         "pow"}};
 
 /**
  * Map the Call nodes to llvm intrinsic.
