[Mlir-commits] [mlir] [mlir][math] `powf(a, b)` drop support when a < 0 (PR #126338)

[Benoit Jacob]

================
@@ -202,25 +202,85 @@ func.func @roundf_func(%a: f32) -> f32 {
 
 // CHECK-LABEL:   func @powf_func
 // CHECK-SAME:    ([[ARG0:%.+]]: f64, [[ARG1:%.+]]: f64)
-func.func @powf_func(%a: f64, %b: f64) ->f64 {
-  // CHECK-DAG: [[CST0:%.+]] = arith.constant 0.000000e+00
-  // CHECK-DAG: [[CST1:%.+]] = arith.constant 1.0
-  // CHECK-DAG: [[TWO:%.+]] = arith.constant 2.000000e+00
-  // CHECK-DAG: [[NEGONE:%.+]] = arith.constant -1.000000e+00
-  // CHECK-DAG: [[SQR:%.+]] = arith.mulf [[ARG0]], [[ARG0]]
-  // CHECK-DAG: [[HALF:%.+]] = arith.divf [[ARG1]], [[TWO]]
-  // CHECK-DAG: [[LOG:%.+]] = math.log [[SQR]]
-  // CHECK-DAG: [[MULT:%.+]] = arith.mulf [[HALF]], [[LOG]]
-  // CHECK-DAG: [[EXPR:%.+]] = math.exp [[MULT]]
-  // CHECK-DAG: [[NEGEXPR:%.+]] = arith.mulf [[EXPR]], [[NEGONE]]
-  // CHECK-DAG: [[REMF:%.+]] = arith.remf [[ARG1]], [[TWO]]
-  // CHECK-DAG: [[CMPNEG:%.+]] = arith.cmpf olt, [[ARG0]]
-  // CHECK-DAG: [[CMPZERO:%.+]] = arith.cmpf one, [[REMF]]
-  // CHECK-DAG: [[AND:%.+]] = arith.andi [[CMPZERO]], [[CMPNEG]]
-  // CHECK-DAG: [[CMPZERO:%.+]] = arith.cmpf oeq, [[ARG1]], [[CST0]]
-  // CHECK-DAG: [[SEL:%.+]] = arith.select [[AND]], [[NEGEXPR]], [[EXPR]]
-  // CHECK-DAG: [[SEL1:%.+]] = arith.select [[CMPZERO]], [[CST1]], [[SEL]]
-  // CHECK: return [[SEL1]]
+func.func @powf_func(%a: f64, %b: f64) -> f64 {
+  // CHECK: [[LOGA:%.+]] = math.log [[ARG0]] : f64
+  // CHECK: [[MUL:%.+]] = arith.mulf [[ARG1]], [[LOGA]] : f64
+  // CHECK: [[EXP:%.+]] = math.exp [[MUL]] : f64
+  // CHECK: return [[EXP]] : f64
+  %ret = math.powf %a, %b : f64
+  return %ret : f64
+}
+
+// CHECK-LABEL:   func @powf_func_zero
+// CHECK-SAME:    ([[ARG0:%.+]]: f64) -> f64
+func.func @powf_func_zero(%a: f64) -> f64{
+  // CHECK: [[ONE:%.+]] = arith.constant 1.000000e+00 : f64
+  // CHECK: return [[ONE]] : f64
+  %b = arith.constant 0.0 : f64
+  %ret = math.powf %a, %b : f64
+  return %ret : f64
+}
+
+// CHECK-LABEL:   func @powf_func_one
+// CHECK-SAME:    ([[ARG0:%.+]]: f64) -> f64
+func.func @powf_func_one(%a: f64) -> f64{
+  // CHECK: return [[ARG0]] : f64
+  %b = arith.constant 1.0 : f64
+  %ret = math.powf %a, %b : f64
+  return %ret : f64
+}
+
+// CHECK-LABEL:   func @powf_func_negone
+// CHECK-SAME:    ([[ARG0:%.+]]: f64) -> f64
+func.func @powf_func_negone(%a: f64) -> f64{
+  // CHECK: [[CSTONE:%.+]] = arith.constant 1.000000e+00 : f64
+  // CHECK: [[DIV:%.+]] = arith.divf [[CSTONE]], [[ARG0]] : f64
+  // CHECK: return [[DIV]] : f64
+  %b = arith.constant -1.0 : f64
+  %ret = math.powf %a, %b : f64
+  return %ret : f64
+}
+
+// CHECK-LABEL:   func @powf_func_half
+// CHECK-SAME:    ([[ARG0:%.+]]: f64) -> f64
+func.func @powf_func_half(%a: f64) -> f64{
+  // CHECK: [[SQRT:%.+]] = math.sqrt [[ARG0]] : f64
+  // CHECK: return [[SQRT]] : f64
+  %b = arith.constant 0.5 : f64
+  %ret = math.powf %a, %b : f64
+  return %ret : f64
+}
+
+// CHECK-LABEL:   func @powf_func_neghalf
+// CHECK-SAME:    ([[ARG0:%.+]]: f64) -> f64
+func.func @powf_func_neghalf(%a: f64) -> f64{
+  // CHECK: [[CSTONE:%.+]] = arith.constant 1.000000e+00 : f64
+  // CHECK: [[SQRT:%.+]] = math.sqrt [[ARG0]] : f64
+  // CHECK: [[DIV:%.+]] = arith.divf [[CSTONE]], [[SQRT]] : f64
----------------
bjacob wrote:

I don't understand how this test succeeds: The C++ code in the rewrite is correctly generating a `math.rsqrt`, which is indeed better than a `math.sqrt + arith.divf`. So how is this test, which requires the latter, passing?

https://github.com/llvm/llvm-project/pull/126338