feat: Add cbrt intrinsic support for f32 and f64

crates/dialect-mir/src/rust_intrinsics.rs

@@ -148,3 +148,7 @@ pub const CALLEE_ATAN2_F64: &str = placeholder!("atan2f64");
 pub const CALLEE_ATAN_F32: &str = placeholder!("atanf32");
 /// Placeholder call used for `f64::atan` / `std::sys::cmath::atan`.
 pub const CALLEE_ATAN_F64: &str = placeholder!("atanf64");
+/// Placeholder call used for `f32::cbrt` / `std::sys::cmath::cbrtf`.
+pub const CALLEE_CBRT_F32: &str = placeholder!("cbrtf32");
+/// Placeholder call used for `f64::cbrt` / `std::sys::cmath::cbrt`.
+pub const CALLEE_CBRT_F64: &str = placeholder!("cbrtf64");

crates/mir-importer/src/translator/terminator/intrinsics/float_math.rs

@@ -113,6 +113,10 @@ pub enum RustFloatMathIntrinsic {
     AtanF32,
     /// `f64::atan` / `std::sys::cmath::atan`.
     AtanF64,
+    /// `f32::cbrt` / `std::sys::cmath::cbrtf`.
+    CbrtF32,
+    /// `f64::cbrt` / `std::sys::cmath::cbrt`.
+    CbrtF64,
 }
 
 impl RustFloatMathIntrinsic {
@@ -182,6 +186,10 @@ impl RustFloatMathIntrinsic {
             "std::sys::cmath::atan2" => Some(Self::Atan2F64),
             "std::sys::cmath::atanf" => Some(Self::AtanF32),
             "std::sys::cmath::atan" => Some(Self::AtanF64),
+            "std::sys::cmath::cbrtf" => Some(Self::CbrtF32),
+            "std::sys::cmath::cbrt" => Some(Self::CbrtF64),
+            "core::num::imp::libm::cbrtf" => Some(Self::CbrtF32),
+            "core::num::imp::libm::cbrt" => Some(Self::CbrtF64),
             other => Self::from_libm_path(other),
         }
     }
@@ -245,6 +253,8 @@ impl RustFloatMathIntrinsic {
             "atan2" => Some(Self::Atan2F64),
             "atanf" => Some(Self::AtanF32),
             "atan" => Some(Self::AtanF64),
+            "cbrtf" => Some(Self::CbrtF32),
+            "cbrt" => Some(Self::CbrtF64),
             _ => None,
         }
     }
@@ -299,6 +309,8 @@ impl RustFloatMathIntrinsic {
             Self::Atan2F64 => rust_intrinsics::CALLEE_ATAN2_F64,
             Self::AtanF32 => rust_intrinsics::CALLEE_ATAN_F32,
             Self::AtanF64 => rust_intrinsics::CALLEE_ATAN_F64,
+            Self::CbrtF32 => rust_intrinsics::CALLEE_CBRT_F32,
+            Self::CbrtF64 => rust_intrinsics::CALLEE_CBRT_F64,
         }
     }
 }
@@ -556,6 +568,35 @@ mod tests {
         );
     }
 
+    /// `f{32,64}::cbrt` reaches device codegen as either the `std::sys::cmath`
+    /// C shim or the in-tree pure-Rust libm path, depending on toolchain.
+    /// Both must map to the libdevice-backed `Cbrt*` variants; pin them so a
+    /// rustc rename surfaces as a test failure rather than an undefined-symbol
+    /// PTX verification error.
+    #[test]
+    fn from_core_path_recognizes_cbrt_via_cmath_and_libm() {
+        for (path, expected) in [
+            ("std::sys::cmath::cbrtf", RustFloatMathIntrinsic::CbrtF32),
+            ("std::sys::cmath::cbrt", RustFloatMathIntrinsic::CbrtF64),
+            (
+                "core::num::imp::libm::cbrtf",
+                RustFloatMathIntrinsic::CbrtF32,
+            ),
+            (
+                "core::num::imp::libm::cbrt",
+                RustFloatMathIntrinsic::CbrtF64,
+            ),
+            ("libm::cbrtf", RustFloatMathIntrinsic::CbrtF32),
+            ("libm::math::cbrt::cbrt", RustFloatMathIntrinsic::CbrtF64),
+        ] {
+            assert_eq!(
+                RustFloatMathIntrinsic::from_core_path(path),
+                Some(expected),
+                "`{path}` did not map to the expected cbrt intrinsic"
+            );
+        }
+    }
+
     /// Libm interception must be anchored to the `libm` crate root. A user
     /// function that shares a libm function name, inside a path that merely
     /// mentions "libm", must stay a regular call: a bare `contains("libm")`

crates/mir-lower/src/convert/ops/call.rs

@@ -219,6 +219,8 @@ enum RustFloatMathIntrinsic {
     Atan2F64,
     AtanF32,
     AtanF64,
+    CbrtF32,
+    CbrtF64,
 }
 
 impl RustFloatMathIntrinsic {
@@ -272,6 +274,8 @@ impl RustFloatMathIntrinsic {
             rust_intrinsics::CALLEE_ATAN2_F64 => Some(Self::Atan2F64),
             rust_intrinsics::CALLEE_ATAN_F32 => Some(Self::AtanF32),
             rust_intrinsics::CALLEE_ATAN_F64 => Some(Self::AtanF64),
+            rust_intrinsics::CALLEE_CBRT_F32 => Some(Self::CbrtF32),
+            rust_intrinsics::CALLEE_CBRT_F64 => Some(Self::CbrtF64),
             _ => None,
         }
     }
@@ -336,6 +340,8 @@ impl RustFloatMathIntrinsic {
             Self::Atan2F64 => Ok("__nv_atan2"),
             Self::AtanF32 => Ok("__nv_atanf"),
             Self::AtanF64 => Ok("__nv_atan"),
+            Self::CbrtF32 => Ok("__nv_cbrtf"),
+            Self::CbrtF64 => Ok("__nv_cbrt"),
         }
     }
 
@@ -1388,6 +1394,14 @@ mod tests {
                 rust_intrinsics::CALLEE_ATAN_F64,
                 RustFloatMathIntrinsic::AtanF64,
             ),
+            (
+                rust_intrinsics::CALLEE_CBRT_F32,
+                RustFloatMathIntrinsic::CbrtF32,
+            ),
+            (
+                rust_intrinsics::CALLEE_CBRT_F64,
+                RustFloatMathIntrinsic::CbrtF64,
+            ),
         ];
 
         for (name, expected) in cases {
@@ -1416,6 +1430,8 @@ mod tests {
         assert_eq!(RustFloatMathIntrinsic::Atan2F64.arg_count(), 2);
         assert_eq!(RustFloatMathIntrinsic::AtanF32.arg_count(), 1);
         assert_eq!(RustFloatMathIntrinsic::AtanF64.arg_count(), 1);
+        assert_eq!(RustFloatMathIntrinsic::CbrtF32.arg_count(), 1);
+        assert_eq!(RustFloatMathIntrinsic::CbrtF64.arg_count(), 1);
         assert_eq!(RustFloatMathIntrinsic::FmaF32.arg_count(), 3);
         assert_eq!(RustFloatMathIntrinsic::FmuladdF64.arg_count(), 3);
         assert_eq!(RustFloatMathIntrinsic::MaxNumNszF32.arg_count(), 2);