Feedback from Andrew.

Author: mcourteaux

src/LegalizeVectors.cpp

@@ -8,6 +8,8 @@
 #include "Util.h"
 
 #include <optional>
+#include <unordered_set>
+#include <vector>
 
 namespace Halide {
 namespace Internal {
@@ -16,13 +18,19 @@ namespace {
 
 using namespace std;
 
-const char *legalization_error_guide = "\n(This issue can most likely be resolved by reducing lane count for vectorize() calls in the schedule, or disabling it.)";
+const char *legalization_error_guide = "\n"
+                                       "(This is an implemenation limitation in Halide right now. This issue can most likely be \n"
+                                       " worked around by reducing lane count for vectorize() calls in GPU schedules, or disabling it.)";
 
 int max_lanes_for_device(DeviceAPI api, int parent_max_lanes) {
+    // The environment variable below (HL_FORCE_VECTOR_LEGALIZATION) is here solely for testing purposes.
+    // It is useful to "stress-test" this lowering pass by forcing a shorter maximal vector size across
+    // all codegen across the entire test suite. This should not be used in real uses of Halide.
     std::string envvar = Halide::Internal::get_env_variable("HL_FORCE_VECTOR_LEGALIZATION");
     if (!envvar.empty()) {
         return std::atoi(envvar.c_str());
     }
+    // The remainder of this function correctly determines the number of lanes the device API supports.
     switch (api) {
     case DeviceAPI::Metal:
     case DeviceAPI::WebGPU:
@@ -53,13 +61,13 @@ std::string vec_name(const string &name, int lane_start, int lane_count) {
 class LiftLetToLetStmt : public IRMutator {
     using IRMutator::visit;
 
+    unordered_set<string> lifted_let_names;
     vector<const Let *> lets;
     Expr visit(const Let *op) override {
-        for (const Let *existing : lets) {
-            internal_assert(existing->name != op->name)
-                << "Let " << op->name << " = ...  cannot be lifted to LetStmt because the name is not unique.";
-        }
+        internal_assert(lifted_let_names.count(op->name) == 0)
+            << "Let " << op->name << " = ...  cannot be lifted to LetStmt because the name is not unique.";
         lets.push_back(op);
+        lifted_let_names.insert(op->name);
         return mutate(op->body);
     }
 
@@ -124,8 +132,7 @@ class ExtractLanes : public IRMutator {
             return result;
         }
 
-        internal_error << "Unhandled trace call in LegalizeVectors' ExtractLanes: " << *event << legalization_error_guide << "\n"
-                       << "Please report this error on GitHub." << legalization_error_guide;
+        internal_error << "Unhandled trace call in LegalizeVectors' ExtractLanes: " << *event << legalization_error_guide;
         return Expr(0);
     }
 
@@ -332,7 +339,7 @@ class LiftExceedingVectors : public IRMutator {
         just_in_let_definition = false;
         Stmt mutated = IRMutator::mutate(s);
         for (auto &let : reverse_view(lets)) {
-            // There is no recurse into let.second. This is handled by repeatedly calling this tranform.
+            // There is no recurse into let.second. This is handled by repeatedly calling this transform.
             mutated = LetStmt::make(let.first, let.second, mutated);
         }
         return mutated;
@@ -576,17 +583,12 @@ Stmt legalize_vectors_in_device_loop(const For *op) {
 }
 
 Stmt legalize_vectors(const Stmt &s) {
-    class LegalizeDeviceLoops : public IRMutator {
-        using IRMutator::visit;
-        Stmt visit(const For *op) override {
-            if (max_lanes_for_device(op->device_api, 0)) {
-                return legalize_vectors_in_device_loop(op);
-            } else {
-                return IRMutator::visit(op);
-            }
+    return mutate_with(s, [&](auto *self, const For *op) {
+        if (max_lanes_for_device(op->device_api, 0)) {
+            return legalize_vectors_in_device_loop(op);
         }
-    } mutator;
-    return mutator.mutate(s);
+        return self->visit_base(op);
+    });
 }
 }  // namespace Internal
 }  // namespace Halide

src/Simplify_Let.cpp

@@ -187,6 +187,11 @@ Body Simplify::simplify_let(const LetOrLetStmt *op, ExprInfo *info) {
                 // pure operations from _all_ arguments to the Shuffle, we will
                 // instead substitute all of the vars that go in the shuffle, and
                 // instead guard against side effects by checking with `is_pure()`.
+                //
+                // Also, it is safe to substitute in without combinatorial
+                // blow-up, because deeply nested concats implies a
+                // combinatorially-large number of vector lanes, which we can't
+                // express in the type system anyway.
                 replacement = substitute(f.new_name, shuffle, replacement);
                 f.new_value = Expr();
                 break;

src/Simplify_Shuffle.cpp

@@ -191,7 +191,9 @@ Expr Simplify::visit(const Shuffle *op, ExprInfo *info) {
                 }
             } else {
                 // We can't... Leave it as a Shuffle of Loads.
-                // Note: don't proceed down.
+                // Note: no mutate-recursion as we are dealing here with a
+                // Shuffle of Loads, which have already undergone mutation
+                // early in this function (new_vectors).
                 return result;
             }
         }
@@ -362,13 +364,15 @@ Expr Simplify::visit(const Shuffle *op, ExprInfo *info) {
                 }
             }
 
+#if 0 // Not sure what this was for. Disabling for now, and will run tests to see what's up.
             for (size_t i = 0; i < new_vectors.size() && can_collapse; i++) {
                 if (new_vectors[i].as<Load>()) {
                     // Don't create a Ramp of a Load, like:
                     // ramp(buf[x], buf[x + 1] - buf[x], ...)
                     can_collapse = false;
                 }
             }
+#endif
 
             if (can_collapse) {
                 return Ramp::make(new_vectors[0], stride, op->indices.size());