Fix floating point argmax

loopy/target/c/__init__.py

@@ -201,6 +201,12 @@ def map_nan(self, expr: p.NaN):
         self.saw_inf_or_nan = True
         return super().map_nan(expr)
 
+    @override
+    def map_variable(self, expr: p.Variable):
+        if expr.name in ("HUGE_VAL", "INFINITY"):
+            self.saw_inf_or_nan = True
+        return super().map_variable(expr)
+
 
 def c99_preamble_generator(preamble_info: PreambleInfo) -> Iterator[tuple[str, str]]:
     if any(dtype.is_integral() for dtype in preamble_info.seen_dtypes):
@@ -561,6 +567,11 @@ def c_symbol_mangler(kernel, name):
     if name in ["INT_MAX", "INT_MIN"]:
         return NumpyType(np.dtype(np.int32)), name
 
+    if name == "INFINITY":
+        return NumpyType(np.dtype(np.float32)), name
+    if name == "HUGE_VAL":
+        return NumpyType(np.dtype(np.float64)), name
+
     return None
 
 # }}}

test/test_target.py

@@ -875,6 +875,22 @@ def test_float3():
     assert "float3" in device_code
 
 
+def test_argmax_ctarget_floating_point():
+    for dtype in (np.float32, np.float64):
+        knl = lp.make_kernel(
+                "{[i]: 0<=i<n}",
+                """
+                max_val[0], max_ind[0] = argmax(i, a[i], i)
+                """,
+                target=lp.ExecutableCTarget())
+
+        knl = lp.set_options(knl, return_dict=True)
+        _evt, out_dict = knl(a=np.array([1, 3, 70, 5, 4], dtype=dtype))
+
+        assert out_dict["max_val"][0] == 70
+        assert out_dict["max_ind"][0] == 2
+
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) > 1: