Add TailStrategy::Predicate (#5856)

* Add TailStrategy::Predicate

* Add some tests for TailStrategy::Predicate.

* Fix missing override.

* Fix comment.

* Tweak target behavior.

* Remove all heuristics

* clang-format.

* clang-tidy.

* TailStrategy::GuardWithIf isn't always faster than scalar code :(

* Use TailStrategy::Predicate in the predicated store/load test.

* What is this test

* Fix test bug.

* Revert x86 behavior.

* Move predicate to Internal namespace.

* Recursively strip tags.

* trigger buildbots

* strip_tags -> unwrap_tags

* Fix comment.

Co-authored-by: Steven Johnson <[email protected]>

Author: dsharletg

src/ApplySplit.cpp

@@ -49,7 +49,8 @@ vector<ApplySplitResult> apply_split(const Split &split, bool is_update, const s
         } else if (is_const_one(split.factor)) {
             // The split factor trivially divides the old extent,
             // but we know nothing new about the outer dimension.
-        } else if (tail == TailStrategy::GuardWithIf) {
+        } else if (tail == TailStrategy::GuardWithIf ||
+                   tail == TailStrategy::Predicate) {
             // It's an exact split but we failed to prove that the
             // extent divides the factor. Use predication to avoid
             // running off the end of the original loop.
@@ -70,6 +71,10 @@ vector<ApplySplitResult> apply_split(const Split &split, bool is_update, const s
             // Inject the if condition *after* doing the substitution
             // for the guarded version.
             Expr cond = likely(old_var <= old_max);
+            if (tail == TailStrategy::Predicate) {
+                // Add the hint for predication.
+                cond = predicate(cond);
+            }
             result.emplace_back(cond);
 
         } else if (tail == TailStrategy::ShiftInwards) {

src/Bounds.cpp

@@ -1232,9 +1232,7 @@ class Bounds : public IRVisitor {
 
             interval.min = Interval::make_max(interval.min, lower.min);
             interval.max = Interval::make_min(interval.max, upper.max);
-        } else if (op->is_intrinsic(Call::likely) ||
-                   op->is_intrinsic(Call::likely_if_innermost)) {
-            internal_assert(op->args.size() == 1);
+        } else if (Call::as_tag(op)) {
             op->args[0].accept(this);
         } else if (op->is_intrinsic(Call::return_second)) {
             internal_assert(op->args.size() == 2);
@@ -2477,10 +2475,8 @@ class BoxesTouched : public IRGraphVisitor {
             for (const auto &pair : cases) {
                 Expr c = pair.first;
                 Stmt body = pair.second;
-                const Call *call = c.as<Call>();
-                if (call && (call->is_intrinsic(Call::likely) ||
-                             call->is_intrinsic(Call::likely_if_innermost) ||
-                             call->is_intrinsic(Call::strict_float))) {
+                const Call *call = Call::as_tag(c);
+                if (call) {
                     c = call->args[0];
                 }
 

src/Derivative.cpp

@@ -1166,7 +1166,7 @@ void ReverseAccumulationVisitor::visit(const Call *op) {
             accumulate(op->args[0], adjoint * (make_one(op->type) - op->args[2]));
             accumulate(op->args[1], adjoint * op->args[2]);
             accumulate(op->args[2], adjoint * (op->args[1] - op->args[0]));
-        } else if (op->is_intrinsic(Call::likely)) {
+        } else if (Call::as_tag(op)) {
             accumulate(op->args[0], adjoint);
         } else if (op->is_intrinsic(Call::return_second)) {
             accumulate(op->args[0], make_const(op->type, 0.0));

src/Func.cpp

@@ -1057,9 +1057,9 @@ void Stage::split(const string &old, const string &outer, const string &inner, c
     }
 
     if (exact) {
-        user_assert(tail == TailStrategy::GuardWithIf)
+        user_assert(tail == TailStrategy::GuardWithIf || tail == TailStrategy::Predicate)
             << "When splitting Var " << old_name
-            << " the tail strategy must be GuardWithIf or Auto. "
+            << " the tail strategy must be GuardWithIf, Predicate or Auto. "
             << "Anything else may change the meaning of the algorithm\n";
     }
 

src/IR.cpp

@@ -617,6 +617,7 @@ const char *const intrinsic_op_names[] = {
     "mulhi_shr",
     "mux",
     "popcount",
+    "predicate",
     "prefetch",
     "promise_clamped",
     "random",

src/IR.h

@@ -529,6 +529,7 @@ struct Call : public ExprNode<Call> {
         mulhi_shr,  // Compute high_half(arg[0] * arg[1]) >> arg[3]. Note that this is a shift in addition to taking the upper half of multiply result. arg[3] must be an unsigned integer immediate.
         mux,
         popcount,
+        predicate,
         prefetch,
         promise_clamped,
         random,
@@ -661,6 +662,10 @@ struct Call : public ExprNode<Call> {
         return nullptr;
     }
 
+    static const Call *as_tag(const Expr &e) {
+        return as_intrinsic(e, {Call::likely, Call::likely_if_innermost, Call::predicate, Call::strict_float});
+    }
+
     bool is_extern() const {
         return (call_type == Extern ||
                 call_type == ExternCPlusPlus ||

src/IROperator.cpp

@@ -1035,6 +1035,19 @@ Stmt remove_likelies(const Stmt &s) {
     return RemoveLikelies().mutate(s);
 }
 
+Expr unwrap_tags(const Expr &e) {
+    if (const Call *tag = Call::as_tag(e)) {
+        return unwrap_tags(tag->args[0]);
+    }
+    return e;
+}
+
+Expr predicate(Expr e) {
+    Type t = e.type();
+    return Internal::Call::make(t, Internal::Call::predicate,
+                                {std::move(e)}, Internal::Call::PureIntrinsic);
+}
+
 Expr requirement_failed_error(Expr condition, const std::vector<Expr> &args) {
     return Internal::Call::make(Int(32),
                                 "halide_error_requirement_failed",

src/IROperator.h

@@ -307,6 +307,16 @@ Expr remove_likelies(const Expr &e);
  * all calls to likely() and likely_if_innermost() removed. */
 Stmt remove_likelies(const Stmt &s);
 
+/** If the expression is a tag helper call, remove it and return
+ * the tagged expression. If not, returns the expression. */
+Expr unwrap_tags(const Expr &e);
+
+/** Expressions tagged with this intrinsic are suggestions that
+ * vectorization of loops with guard ifs should be implemented with
+ * non-faulting predicated loads and stores, instead of scalarizing
+ * an if statement. */
+Expr predicate(Expr e);
+
 // Secondary args to print can be Exprs or const char *
 inline HALIDE_NO_USER_CODE_INLINE void collect_print_args(std::vector<Expr> &args) {
 }

src/IRPrinter.cpp

@@ -147,6 +147,9 @@ std::ostream &operator<<(std::ostream &out, const TailStrategy &t) {
     case TailStrategy::GuardWithIf:
         out << "GuardWithIf";
         break;
+    case TailStrategy::Predicate:
+        out << "Predicate";
+        break;
     case TailStrategy::ShiftInwards:
         out << "ShiftInwards";
         break;

src/LICM.cpp

@@ -97,9 +97,7 @@ class LiftLoopInvariants : public IRMutator {
             }
         }
         if (const Call *call = e.as<Call>()) {
-            if (call->is_intrinsic(Call::strict_float) ||
-                call->is_intrinsic(Call::likely) ||
-                call->is_intrinsic(Call::likely_if_innermost) ||
+            if (Call::as_tag(call) ||
                 call->is_intrinsic(Call::reinterpret)) {
                 // Don't lift these intrinsics. They're free.
                 return should_lift(call->args[0]);