From 358f862bf7023470f65fb567d6655f1bed8667b6 Mon Sep 17 00:00:00 2001
From: Andrew Adams <andrew.b.adams@gmail.com>
Date: Tue, 11 Jun 2024 13:35:21 -0700
Subject: [PATCH] Don't try to codegen predicated atomic stores

By disabling predication if an Atomic node is found.

Fixes #8280.
---
 src/VectorizeLoops.cpp                     |  7 +++++++
 test/correctness/predicated_store_load.cpp | 23 +++++++++++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/VectorizeLoops.cpp b/src/VectorizeLoops.cpp
index 0745a34a9d39..1481e37b0345 100644
--- a/src/VectorizeLoops.cpp
+++ b/src/VectorizeLoops.cpp
@@ -380,6 +380,13 @@ class PredicateLoadStore : public IRMutator {
         return pred;
     }
 
+    Stmt visit(const Atomic *op) override {
+        // We don't support codegen for vectorized predicated atomic stores, so
+        // just bail out.
+        valid = false;
+        return op;
+    }
+
     Expr visit(const Load *op) override {
         valid = valid && ((op->predicate.type().lanes() == lanes) || (op->predicate.type().is_scalar() && !expr_uses_var(op->index, var)));
         if (!valid) {
diff --git a/test/correctness/predicated_store_load.cpp b/test/correctness/predicated_store_load.cpp
index 023a3c2f92f9..4d29ae630760 100644
--- a/test/correctness/predicated_store_load.cpp
+++ b/test/correctness/predicated_store_load.cpp
@@ -464,6 +464,22 @@ int vectorized_predicated_load_lut_test(const Target &t) {
     return 0;
 }
 
+int predicated_atomic_store_test(const Target &t) {
+    // We don't support atomic predicated stores, so ensure that we don't
+    // generate them. See https://github.com/halide/Halide/issues/8280
+    ImageParam in(Float(32), 1);
+    Func f;
+    Var x;
+    RDom r(0, 20);
+
+    f(x) = 0.f;
+    f(x) += in(r) + x;
+    f.update().vectorize(x, 8, TailStrategy::GuardWithIf).atomic().parallel(r);
+
+    f.compile_jit(t);
+    return 0;
+}
+
 }  // namespace
 
 int main(int argc, char **argv) {
@@ -529,6 +545,11 @@ int main(int argc, char **argv) {
         return 1;
     }
 
+    printf("predicated atomic store test\n");
+    if (predicated_atomic_store_test(t) != 0) {
+        return 1;
+    }
+
     printf("Success!\n");
     return 0;
-}
\ No newline at end of file
+}