Make IRComparer consider nans to be less than non-nans.

src/IREquality.cpp

@@ -106,6 +106,22 @@ IRComparer::CmpResult IRComparer::compare_scalar(T a, T b) {
         return result;
     }
 
+    if constexpr (std::is_floating_point_v<T>) {
+        // NaNs are equal to each other and less than non-nans
+        if (std::isnan(a) && std::isnan(b)) {
+            result = Equal;
+            return result;
+        }
+        if (std::isnan(a)) {
+            result = LessThan;
+            return result;
+        }
+        if (std::isnan(b)) {
+            result = GreaterThan;
+            return result;
+        }
+    }
+
     if (a < b) {
         result = LessThan;
     } else if (a > b) {
@@ -125,6 +141,7 @@ IRComparer::CmpResult IRComparer::compare_expr(const Expr &a, const Expr &b) {
         return result;
     }
 
+    // Undefined values are equal to each other and less than defined values
     if (!a.defined() && !b.defined()) {
         result = Equal;
         return result;

test/correctness/CMakeLists.txt

@@ -63,6 +63,7 @@ tests(GROUPS correctness
       convolution.cpp
       convolution_multiple_kernels.cpp
       cross_compilation.cpp
+      cse_nan.cpp
       cuda_8_bit_dot_product.cpp
       custom_allocator.cpp
       custom_auto_scheduler.cpp

test/correctness/cse_nan.cpp

@@ -0,0 +1,38 @@
+#include <cmath>
+#include <cstdio>
+
+#include "Halide.h"
+using namespace Halide;
+
+int main() {
+    ImageParam xyz{Float(32), 3, "xyz"};
+    Target t = get_jit_target_from_environment().with_feature(Target::StrictFloat);
+
+    Var col{"col"}, row{"row"};
+    Func nan_or_one{"nan_or_one"};
+    nan_or_one(col, row) = Halide::select(is_nan(xyz(col, row, 0)), NAN, 1.0f);
+
+    Buffer<float> true_buf{1, 1, 1};
+    true_buf(0, 0, 0) = NAN;
+
+    Buffer<float> false_buf{1, 1, 1};
+    false_buf(0, 0, 0) = 2.0f;
+
+    Buffer<float> true_result{1, 1};
+    Buffer<float> false_result{1, 1};
+
+    xyz.set(true_buf);
+    nan_or_one.realize({true_result}, t);
+
+    xyz.set(false_buf);
+    nan_or_one.realize({false_result}, t);
+
+    if (std::isnan(true_result(0, 0)) && false_result(0, 0) == 1.0f) {
+        printf("Success!\n");
+        return 0;
+    } else {
+        fprintf(stderr, "ERROR: T = %f ; TR = %f ; F = %f ; FR = %f\n",
+                true_buf(0, 0, 0), true_result(0, 0), false_buf(0, 0, 0), false_result(0, 0));
+        return -1;
+    }
+}