Fix OpenCL positive and negative INF constants. (#8266)

halide · Jun 20, 2024 · b921710 · b921710
1 parent ea775cc
commit b921710
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/src/CodeGen_OpenCL_Dev.cpp b/src/CodeGen_OpenCL_Dev.cpp
@@ -12,6 +12,7 @@
 #include "EliminateBoolVectors.h"
 #include "EmulateFloat16Math.h"
 #include "ExprUsesVar.h"
+#include "Float16.h"
 #include "IRMutator.h"
 #include "IROperator.h"
 #include "Simplify.h"
@@ -1183,11 +1184,15 @@ void CodeGen_OpenCL_Dev::init_module() {
     }
 
     if (target.has_feature(Target::CLHalf)) {
+        const uint16_t nan_f16 = float16_t::make_nan().to_bits();
+        const uint16_t neg_inf_f16 = float16_t::make_negative_infinity().to_bits();
+        const uint16_t inf_f16 = float16_t::make_infinity().to_bits();
+
         src_stream << "#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n"
                    << "inline half half_from_bits(unsigned short x) {return __builtin_astype(x, half);}\n"
-                   << "inline half nan_f16() { return half_from_bits(32767); }\n"
-                   << "inline half neg_inf_f16() { return half_from_bits(31744); }\n"
-                   << "inline half inf_f16() { return half_from_bits(64512); }\n"
+                   << "inline half nan_f16() { return half_from_bits(" << nan_f16 << "); }\n"
+                   << "inline half neg_inf_f16() { return half_from_bits(" << neg_inf_f16 << "); }\n"
+                   << "inline half inf_f16() { return half_from_bits(" << inf_f16 << "); }\n"
                    << "inline bool is_nan_f16(half x) {return isnan(x); }\n"
                    << "inline bool is_inf_f16(half x) {return isinf(x); }\n"
                    << "inline bool is_finite_f16(half x) {return isfinite(x); }\n"

diff --git a/test/correctness/gpu_f16_intrinsics.cpp b/test/correctness/gpu_f16_intrinsics.cpp
@@ -4,8 +4,9 @@ using namespace Halide;
 int main(int argc, char *argv[]) {
 
     auto target = get_jit_target_from_environment();
-    if (!target.has_feature(Target::Metal)) {
-        printf("[SKIP] No metal target enabled.\n");
+    if (!target.has_feature(Target::Metal) &&
+        !target.features_all_of({Target::OpenCL, Target::CLHalf})) {
+        printf("[SKIP] Test only applies to Metal and OpenCL+CLHalf.\n");
         return 0;
     }