[CINN] Add type support for Sigmoid and LogicalNot (#72307)

ggggxm · web-flow · commit 3a7b2b0137e0 · 2025-04-21T11:31:38.000+08:00
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
@@ -1254,19 +1254,21 @@ class SigmoidOpPattern
             .dyn_cast<paddle::dialect::DenseTensorType>()
             .dtype());
 
+    auto one_type = input_dtype;
     auto in = op->operand_source(0);
     bool need_cast = (input_dtype == phi::DataType::FLOAT16 ||
                       input_dtype == phi::DataType::BFLOAT16 ||
                       input_dtype == phi::DataType::UINT16);
     if (need_cast) {
       in = rewriter.Build<paddle::dialect::CastOp>(in, phi::DataType::FLOAT32)
                .result(0);
+      one_type = phi::DataType::FLOAT32;
     }
 
     // 1 / ( 1 + exp(-x))
     auto one = rewriter
                    .Build<paddle::dialect::FullOp>(
-                       std::vector<int64_t>({1}), 1.0, phi::DataType::FLOAT32)
+                       std::vector<int64_t>({1}), 1.0, one_type)
                    .result(0);
     auto minus_x =
         rewriter.Build<paddle::dialect::ScaleOp>(in, -1.0, 0.0).result(0);
diff --git a/paddle/cinn/lang/builtin.h b/paddle/cinn/lang/builtin.h
@@ -96,7 +96,7 @@ Expr Abs(Expr e);
 
 inline Expr Negative(Expr e) { return -e; }
 inline Expr Identity(Expr e) { return e; }
-inline Expr LogicalNot(Expr e) { return !e; }
+inline Expr LogicalNot(Expr e) { return !(ir::Cast::Make(common::Bool(), e)); }
 inline Expr BitwiseNot(Expr e) { return ~e; }
 inline Expr BitwiseAnd(Expr a, Expr b) { return a & b; }
 inline Expr BitwiseOr(Expr a, Expr b) { return a | b; }
