Iterate over lets in the correct order in VectorizeLoops

src/VectorizeLoops.cpp

@@ -960,19 +960,27 @@ class VectorSubs : public IRMutator {
 
             vectorized_vars.push_back({op->name, min, (int)extent_int->value});
             update_replacements();
-            // Go over lets which were vectorized and update them according to the current
-            // loop level.
-            for (auto it = scope.cbegin(); it != scope.cend(); ++it) {
-                string vectorized_name = get_widened_var_name(it.name());
-                Expr vectorized_value = mutate(it.value());
+            // Go over lets which were vectorized in the order of their occurrence and update
+            // them according to the current loop level.
+            for (auto let = containing_lets.begin(); let != containing_lets.end(); let++) {
+                // Skip if this var wasn't vectorized.
+                if (!scope.contains(let->first)) {
+                    continue;
+                }
+                string vectorized_name = get_widened_var_name(let->first);
+                Expr vectorized_value = mutate(scope.get(let->first));
                 vector_scope.push(vectorized_name, vectorized_value);
             }
 
             body = mutate(body);
 
             // Append vectorized lets for this loop level.
-            for (auto it = scope.cbegin(); it != scope.cend(); ++it) {
-                string vectorized_name = get_widened_var_name(it.name());
+            for (auto let = containing_lets.rbegin(); let != containing_lets.rend(); let++) {
+                // Skip if this var wasn't vectorized.
+                if (!scope.contains(let->first)) {
+                    continue;
+                }
+                string vectorized_name = get_widened_var_name(let->first);
                 Expr vectorized_value = vector_scope.get(vectorized_name);
                 vector_scope.pop(vectorized_name);
                 InterleavedRamp ir;

test/correctness/vectorize_nested.cpp

@@ -192,6 +192,29 @@ int vectorize_all_d() {
     return 0;
 }
 
+int vectorize_lets_order() {
+    const int width = 128;
+    const int height = 128;
+
+    Var x("x"), y("y"), yo("yo"), yi("yi"), yoi("yoi"), yoio("yoio"), yoii("yoii");
+    Func f("f");
+    f(x, y) = x + y;
+    f.split(y, yo, yi, 8, TailStrategy::Auto)
+        .split(yo, yo, yoi, 4, TailStrategy::RoundUp)
+        .vectorize(yoi)
+        .vectorize(yi)
+        .split(yoi, yoio, yoii, 2, TailStrategy::Auto);
+    Buffer<int> result = f.realize({width, height});
+
+    auto cmp_func = [](int x, int y) {
+        return x + y;
+    };
+    if (check_image(result, cmp_func)) {
+        return 1;
+    }
+
+    return 0;
+}
 int vectorize_inner_of_scalarization() {
     ImageParam in(UInt(8), 2);
 
@@ -289,6 +312,11 @@ int main(int argc, char **argv) {
         return 1;
     }
 
+    if (vectorize_lets_order()) {
+        printf("vectorize_lets_order failed\n");
+        return 1;
+    }
+
     if (vectorize_inner_of_scalarization()) {
         printf("vectorize_inner_of_scalarization failed\n");
         return 1;