From d3a35a768b9e0c2f826b5d196b286a0cd0c65328 Mon Sep 17 00:00:00 2001
From: Snehal <snehalreddy@google.com>
Date: Sat, 8 Aug 2026 15:29:07 +0000
Subject: [PATCH] fix(warp): fix float casting in warp_match_any/warp_match_all

This commit separates the f32 and f64 types from the impl_match! macro and correctly handles float equality matching using .to_bits(), preventing catastrophic truncation precision loss during implicit integer casts.

TAG=agy
---
 crates/cuda_std/src/warp.rs | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/crates/cuda_std/src/warp.rs b/crates/cuda_std/src/warp.rs
index 6f3e13ab..58c1b3f9 100644
--- a/crates/cuda_std/src/warp.rs
+++ b/crates/cuda_std/src/warp.rs
@@ -272,7 +272,7 @@ macro_rules! impl_match {
                     }
                     unsafe fn match_all(mask: u32, value: Self) -> Option<u32> {
                         let (val, pred) = unsafe { [<match_all_ $width>](mask, value as [<u $width>]) };
-                        pred.then(|| val)
+                        pred.then_some(val)
                     }
                 }
             }
@@ -285,8 +285,26 @@ impl_match! {
     i64, 64,
     u32, 32,
     u64, 64,
-    f32, 32,
-    f64, 64,
+}
+
+impl WarpMatchValue for f32 {
+    unsafe fn match_any(mask: u32, value: Self) -> u32 {
+        unsafe { match_any_32(mask, value.to_bits()) }
+    }
+    unsafe fn match_all(mask: u32, value: Self) -> Option<u32> {
+        let (val, pred) = unsafe { match_all_32(mask, value.to_bits()) };
+        pred.then_some(val)
+    }
+}
+
+impl WarpMatchValue for f64 {
+    unsafe fn match_any(mask: u32, value: Self) -> u32 {
+        unsafe { match_any_64(mask, value.to_bits()) }
+    }
+    unsafe fn match_all(mask: u32, value: Self) -> Option<u32> {
+        let (val, pred) = unsafe { match_all_64(mask, value.to_bits()) };
+        pred.then_some(val)
+    }
 }
 
 #[gpu_only]