diff --git a/slinky/builder/optimizations.cc b/slinky/builder/optimizations.cc index 209e2fcc..eb7a9212 100644 --- a/slinky/builder/optimizations.cc +++ b/slinky/builder/optimizations.cc @@ -528,6 +528,17 @@ class copy_aliaser : public stmt_mutator { std::optional& target_info = buffers[target_var]; assert(target_info); + if (target_info->producers > 1 && !target_info->is_output) { + // The target is an intermediate with more than one producer (e.g. a + // concatenate feeding a downstream consumer). Aliasing this buffer + // into one slice of the target is not correct in general: the other + // producers write the rest of the target via copies into its own + // allocation, and consumers expect the full buffer. (Symmetric to the + // `info.producers > 1` guard above, which only covers aliasing the + // multi-producer buffer itself.) + continue; + } + if (!alias_compatible(info, alias, target_var, *target_info)) { continue; } diff --git a/slinky/builder/test/cannot_alias.cc b/slinky/builder/test/cannot_alias.cc index 622753b4..1f5eec10 100644 --- a/slinky/builder/test/cannot_alias.cc +++ b/slinky/builder/test/cannot_alias.cc @@ -1102,4 +1102,118 @@ TEST_P(stride_propagation, cropped) { ASSERT_EQ(eval_ctx.heap.allocs.size(), constrain_outer ? 2 : 1); } +// Regression test for aliasing a producer of a *multi-producer* buffer. A +// concatenate has 3 producers (a, b, c) writing disjoint feature slices of +// `packed`, and a pass-2 that slices `packed` back into 3 views. With everything +// stored at its producing func's level (store_outputs_innermost), `packed` is a +// single buffer that all of the producers AND the slice-views alias into. It was +// tempting for the copy aliaser to alias a producer into `packed`; doing so +// corrupts the shared allocation (the other producers write the rest of it), and +// here it produces an out-of-bounds access. +index_t combine3( + const buffer& sa, const buffer& sb, const buffer& sc, const buffer& out) { + const buffer* s[3] = {&sa, &sb, &sc}; + for (index_t ii = out.dim(1).begin(); ii < out.dim(1).end(); ++ii) { + index_t base = 0; + for (int r = 0; r < 3; ++r) { + for (index_t ff = 0; ff < s[r]->dim(0).extent(); ++ff) { + int acc = 0; + for (index_t kk = s[r]->dim(1).begin(); kk < s[r]->dim(1).end(); ++kk) + acc += (*s[r])(s[r]->dim(0).min() + ff, kk, ii); + out(base + ff, ii) = acc; + } + base += s[r]->dim(0).extent(); + } + } + return 0; +} + +TEST(cannot_alias, concat_producer_batch_reduce) { + node_context ctx; + const int Wa = 1, Wb = 1, Wc = 4, K = 4, I = 3; + + auto in_a = buffer_expr::make(ctx, "in_a", 3, sizeof(int)); + auto in_b = buffer_expr::make(ctx, "in_b", 3, sizeof(int)); + auto in_c = buffer_expr::make(ctx, "in_c", 3, sizeof(int)); + auto a = buffer_expr::make(ctx, "a", 3, sizeof(int)); + auto b = buffer_expr::make(ctx, "b", 3, sizeof(int)); + auto c = buffer_expr::make(ctx, "c", 3, sizeof(int)); + auto packed = buffer_expr::make(ctx, "packed", 3, sizeof(int)); + auto sl_a = buffer_expr::make(ctx, "sl_a", 3, sizeof(int)); + auto sl_b = buffer_expr::make(ctx, "sl_b", 3, sizeof(int)); + auto sl_c = buffer_expr::make(ctx, "sl_c", 3, sizeof(int)); + auto out = buffer_expr::make(ctx, "out", 2, sizeof(int)); + + var f(ctx, "f"), k(ctx, "k"), i(ctx, "i"); + + func pa = func::make(add_1, {{in_a, {point(f), point(k), point(i)}}}, {{a, {f, k, i}}}); + func pb = func::make(add_1, {{in_b, {point(f), point(k), point(i)}}}, {{b, {f, k, i}}}); + func pc = func::make(add_1, {{in_c, {point(f), point(k), point(i)}}}, {{c, {f, k, i}}}); + + func concat = func::make_concat({a, b, c}, {packed, {f, k, i}}, 0, + {0, in_a->dim(0).extent(), in_a->dim(0).extent() + in_b->dim(0).extent(), + in_a->dim(0).extent() + in_b->dim(0).extent() + in_c->dim(0).extent()}); + + // Slice packed back into its 3 feature regions (views that ALSO alias + // into packed), then reduce each over k. + func slice_a = func::make_copy({packed, {point(f), point(k), point(i)}}, {sl_a, {f, k, i}}); + func slice_b = func::make_copy({packed, {point(f) + Wa, point(k), point(i)}}, {sl_b, {f, k, i}}); + func slice_c = func::make_copy({packed, {point(f) + (Wa + Wb), point(k), point(i)}}, {sl_c, {f, k, i}}); + + interval_expr all_k = in_a->dim(1).bounds; + func combine = func::make(combine3, + {{sl_a, {in_a->dim(0).bounds, all_k, point(i)}}, {sl_b, {in_b->dim(0).bounds, all_k, point(i)}}, + {sl_c, {in_c->dim(0).bounds, all_k, point(i)}}}, + {{out, {f, i}}}, call_stmt::attributes{.name = "combine"}); + + // One outer i-loop over everything; packed, producers and slice-views are all + // computed per-i and stored at the i level, so they alias into the SAME per-i + // packed symbol. + combine.loops({{i, 1}}); + concat.compute_at({&combine, i}); + slice_a.compute_at({&combine, i}); + slice_b.compute_at({&combine, i}); + slice_c.compute_at({&combine, i}); + pa.compute_at({&combine, i}); + pb.compute_at({&combine, i}); + pc.compute_at({&combine, i}); + // Store every buffer at its producing func's own level: packed at the concat's + // level, partials at theirs. + concat.store_outputs_innermost(); + pa.store_outputs_innermost(); + pb.store_outputs_innermost(); + pc.store_outputs_innermost(); + slice_a.store_outputs_innermost(); + slice_b.store_outputs_innermost(); + slice_c.store_outputs_innermost(); + + pipeline p = build_pipeline(ctx, {in_a, in_b, in_c}, {out}); + + buffer in_a_buf({Wa, K, I}), in_b_buf({Wb, K, I}), in_c_buf({Wc, K, I}); + init_random(in_a_buf); + init_random(in_b_buf); + init_random(in_c_buf); + buffer out_buf({Wa + Wb + Wc, I}); + out_buf.allocate(); + + const raw_buffer* inputs[] = {&in_a_buf, &in_b_buf, &in_c_buf}; + const raw_buffer* outputs[] = {&out_buf}; + test_context eval_ctx; + p.evaluate(inputs, outputs, eval_ctx); + + auto ref = [&](int ff, int ii) { + int s = 0; + for (int kk = 0; kk < K; ++kk) { + int v = ff < Wa ? in_a_buf(ff, kk, ii) + : ff < Wa + Wb ? in_b_buf(ff - Wa, kk, ii) + : in_c_buf(ff - Wa - Wb, kk, ii); + s += v + 1; + } + return s; + }; + for (int ii = 0; ii < I; ++ii) + for (int ff = 0; ff < Wa + Wb + Wc; ++ff) + ASSERT_EQ(out_buf(ff, ii), ref(ff, ii)) << "f=" << ff << " i=" << ii; +} + } // namespace slinky