Add a math.cbrt instruction and lowering to libm.

[jreiffers]

Summary

There's currently no way to get accurate cube roots in the math dialect.
powf(x, 1/3.0) is too inaccurate in some cases.

mlir/include/mlir/Dialect/Math/IR/MathOps.td

let description = [{
     // Scalar variant.
     %a = math.atan2 %b, %c : f32
     ```
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// CbrtOp
+//===----------------------------------------------------------------------===//
+
+def Math_CbrtOp : Math_FloatUnaryOp<"cbrt"> {
+  let summary = "cube root of the specified value";
+  let description = [{
+    The `cbrt` operation computes the cube root. It takes one operand of
+    floating point type (i.e., scalar, tensor or vector) and returns one result
+    of the same type. It has no standard attributes.
+
+    Example:
+
+    ```mlir
+    // Scalar cube root value.
+    %a = math.cbrt %b : f64
+    ```
+
+    Note: This op is not equivalent to powf(..., 1/3.0).
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // CeilOp
 //===----------------------------------------------------------------------===//
 
 def Math_CeilOp : Math_FloatUnaryOp<"ceil"> {
   let summary = "ceiling of the specified value";
   let description = [{
     Syntax:
 

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

patterns.add<PromoteOpToF32<math::Atan2Op>, PromoteOpToF32<math::ExpM1Op>,
                PromoteOpToF32<math::SinOp>, PromoteOpToF32<math::ErfOp>,
                PromoteOpToF32<math::RoundEvenOp>, PromoteOpToF32<math::RoundOp>,
                PromoteOpToF32<math::AtanOp>, PromoteOpToF32<math::TanOp>,
                PromoteOpToF32<math::TruncOp>>(patterns.getContext(), benefit);
   patterns.add<ScalarOpToLibmCall<math::AtanOp>>(patterns.getContext(), "atanf",
                                                  "atan", benefit);
   patterns.add<ScalarOpToLibmCall<math::Atan2Op>>(patterns.getContext(),
                                                   "atan2f", "atan2", benefit);
+  patterns.add<ScalarOpToLibmCall<math::CbrtOp>>(patterns.getContext(), "cbrtf",
+                                                 "cbrt", benefit);
   patterns.add<ScalarOpToLibmCall<math::ErfOp>>(patterns.getContext(), "erff",
                                                 "erf", benefit);
   patterns.add<ScalarOpToLibmCall<math::ExpM1Op>>(patterns.getContext(),
                                                   "expm1f", "expm1", benefit);
   patterns.add<ScalarOpToLibmCall<math::TanOp>>(patterns.getContext(), "tanf",
                                                 "tan", benefit);
   patterns.add<ScalarOpToLibmCall<math::TanhOp>>(patterns.getContext(), "tanhf",
                                                  "tanh", benefit);

mlir/test/Conversion/MathToLibm/convert-to-libm.mlir

 // RUN: mlir-opt %s -convert-math-to-libm -canonicalize | FileCheck %s
 
 // CHECK-DAG: @atan(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @atanf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @erf(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @erff(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @expm1(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @expm1f(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @atan2(f64, f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @atan2f(f32, f32) -> f32 attributes {llvm.readnone}
+// CHECK-DAG: @cbrt(f64) -> f64 attributes {llvm.readnone}
+// CHECK-DAG: @cbrtf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @tan(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @tanf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @tanh(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @tanhf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @round(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @roundf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @roundeven(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @roundevenf(f32) -> f32 attributes {llvm.readnone}
func.func @trunc_caller(%float: f32, %double: f64) -> (f32, f64) {
   // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @truncf(%[[FLOAT]]) : (f32) -> f32
   %float_result = math.trunc %float : f32
   // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @trunc(%[[DOUBLE]]) : (f64) -> f64
   %double_result = math.trunc %double : f64
   // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
   return %float_result, %double_result : f32, f64
 }
 
+// CHECK-LABEL: func @cbrt_caller
+// CHECK-SAME: %[[FLOAT:.*]]: f32
+// CHECK-SAME: %[[DOUBLE:.*]]: f64
+func.func @cbrt_caller(%float: f32, %double: f64) -> (f32, f64)  {
+  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @cbrtf(%[[FLOAT]]) : (f32) -> f32
+  %float_result = math.cbrt %float : f32
+  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @cbrt(%[[DOUBLE]]) : (f64) -> f64
+  %double_result = math.cbrt %double : f64
+  // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
+  return %float_result, %double_result : f32, f64
+}
+
 // CHECK-LABEL: func @cos_caller
 // CHECK-SAME: %[[FLOAT:.*]]: f32
 // CHECK-SAME: %[[DOUBLE:.*]]: f64
 func.func @cos_caller(%float: f32, %double: f64) -> (f32, f64)  {
   // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @cosf(%[[FLOAT]]) : (f32) -> f32
   %float_result = math.cos %float : f32
   // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @cos(%[[DOUBLE]]) : (f64) -> f64
   %double_result = math.cos %double : f64

mlir/test/Dialect/Math/ops.mlir

func.func @atan2(%f: f32, %v: vector<4xf32>, %t: tensor<4x4x?xf32>) {
   %0 = math.atan2 %f, %f : f32
   // CHECK: %{{.*}} = math.atan2 %[[V]], %[[V]] : vector<4xf32>
   %1 = math.atan2 %v, %v : vector<4xf32>
   // CHECK: %{{.*}} = math.atan2 %[[T]], %[[T]] : tensor<4x4x?xf32>
   %2 = math.atan2 %t, %t : tensor<4x4x?xf32>
   return
 }
 
+// CHECK-LABEL: func @cbrt(
+// CHECK-SAME:             %[[F:.*]]: f32, %[[V:.*]]: vector<4xf32>, %[[T:.*]]: tensor<4x4x?xf32>)
+func.func @cbrt(%f: f32, %v: vector<4xf32>, %t: tensor<4x4x?xf32>) {
+  // CHECK: %{{.*}} = math.cbrt %[[F]] : f32
+  %0 = math.cbrt %f : f32
+  // CHECK: %{{.*}} = math.cbrt %[[V]] : vector<4xf32>
+  %1 = math.cbrt %v : vector<4xf32>
+  // CHECK: %{{.*}} = math.cbrt %[[T]] : tensor<4x4x?xf32>
+  %2 = math.cbrt %t : tensor<4x4x?xf32>
+  return
+}
+
 // CHECK-LABEL: func @cos(
 // CHECK-SAME:            %[[F:.*]]: f32, %[[V:.*]]: vector<4xf32>, %[[T:.*]]: tensor<4x4x?xf32>)
 func.func @cos(%f: f32, %v: vector<4xf32>, %t: tensor<4x4x?xf32>) {
   // CHECK: %{{.*}} = math.cos %[[F]] : f32
   %0 = math.cos %f : f32
   // CHECK: %{{.*}} = math.cos %[[V]] : vector<4xf32>
   %1 = math.cos %v : vector<4xf32>
   // CHECK: %{{.*}} = math.cos %[[T]] : tensor<4x4x?xf32>