diff --git a/src/ascend/transform/auto_schedule/buffer_alias_analysis.h b/src/ascend/transform/auto_schedule/buffer_alias_analysis.h index bb7e8c987d..df0a2e65ad 100644 --- a/src/ascend/transform/auto_schedule/buffer_alias_analysis.h +++ b/src/ascend/transform/auto_schedule/buffer_alias_analysis.h @@ -1245,11 +1245,17 @@ class BufferAliasAnalyzerImpl { IRStructure *child = access.task; while (child != nullptr && child->GetParent() != parent_scope) child = child->GetParent(); - if (child == nullptr || !child->IsControl()) + // An access outside this scope is a boundary of the child lifetime. + // Every access inside this scope, however, must be covered by the + // selected child. Dropping a direct or single-trip sibling access can + // hide the writer of a value that survives across parent iterations. + if (child == nullptr) continue; + if (!child->IsControl()) + return nullptr; auto *candidate = static_cast(child); if (IsGuaranteedSingleTrip(candidate)) - continue; + return nullptr; if (control == nullptr) { control = candidate; } else if (control != candidate) { @@ -1297,11 +1303,12 @@ class BufferAliasAnalyzerImpl { } bool parent_live_through = child_result.live_through; - if (child_result.promote_live_through_as_live_in) { + if (child_result.promote_live_through_as_live_in && + parent_scope == nullptr) { // The storage spans the whole child control, so it cannot alias inside - // that child. At the parent level the same lifetime is a finite prefix - // ending at the child's last read and may alias a later sibling - // generation. + // that child. Only a non-periodic parent can reduce it to a finite + // prefix ending at the child's last read. A loop parent needs the + // incoming value again on its next iteration. parent_live_through = false; child_result.live_out = false; } @@ -1368,7 +1375,7 @@ class BufferAliasAnalyzerImpl { } } child_result.live_through = parent_live_through; - child_result.promote_live_through_as_live_in = false; + child_result.promote_live_through_as_live_in &= parent_scope != nullptr; child_result.generations = std::move(generations); child_result.summary_control = control; return child_result; diff --git a/testing/ascend/auto_schedule/test_buffer_lifetimes.py b/testing/ascend/auto_schedule/test_buffer_lifetimes.py index 486a11cb32..1c0226fe93 100644 --- a/testing/ascend/auto_schedule/test_buffer_lifetimes.py +++ b/testing/ascend/auto_schedule/test_buffer_lifetimes.py @@ -967,5 +967,74 @@ def test_l0_address_reuse_includes_the_bound_sf_lifetime(sticky): assert sa.data not in aliases and sb.data not in aliases +@pytest.mark.parametrize("reader_depth", [1, 2]) +def test_preloaded_nested_reader_stays_live_across_parent_iterations(reader_depth): + @T.prim_func + def main( + x: T.Tensor((1,), "int32"), + y: T.Tensor((4,), "int32"), + out_x: T.Tensor((4, 2, 2), "int32"), + out_y: T.Tensor((4,), "int32"), + ): + with T.Kernel(1): + a = T.alloc_shared((8,), "int32") + b = T.alloc_shared((8,), "int32") + a[0] = x[0] + for k in T.Pipelined(4, num_stages=1): + for j in T.serial(2): + if reader_depth == 2: + for i in T.serial(2): + out_x[k, j, i] = a[0] + else: + out_x[k, j, 0] = a[0] + b[0] = y[k] + out_y[k] = b[0] + + # b(k) must not overwrite the preload read again by a child of k + 1. + assert ("a", "b") not in _insert_sync_alias_pairs(main) + + if __name__ == "__main__": tilelang.testing.main() + + +def _make_parent_writer_nested_reader_program(versions: int, inner_extent: int, single_trip_writer: bool): + @T.prim_func + def main( + x: T.Tensor((8,), "int32"), + y: T.Tensor((8,), "int32"), + out_x: T.Tensor((8, inner_extent), "int32"), + out_y: T.Tensor((8,), "int32"), + ): + with T.Kernel(1): + a = T.alloc_shared((8,), "int32") + b = T.alloc_shared((8,), "int32") + T.annotate_buffer_versions({a: versions}) + for k in T.Pipelined(8, num_stages=5, annotations={"enable_offset": True}): + with T.Stage(0): + if single_trip_writer: + for _ in T.serial(1): + with T.Stage(0): + a[0] = x[k] + else: + a[0] = x[k] + with T.Stage(1): + for j in T.serial(inner_extent): + with T.Stage(0): + out_x[k, j] = a[0] + with T.Stage(4): + b[0] = y[k] + out_y[k] = b[0] + + return main + + +@pytest.mark.parametrize("single_trip_writer", [False, True], ids=["direct-writer", "single-trip-writer"]) +@pytest.mark.parametrize("versions", [2, 3]) +@pytest.mark.parametrize("inner_extent", [1, 2], ids=["single-trip", "multi-trip"]) +def test_parent_writer_and_nested_reader_keep_periodic_lifetime(versions, inner_extent, single_trip_writer): + # At physical step 6, a(6) is waiting for its stage-1 reader when b(2) + # writes. Whole-allocation reuse is unsafe even though a(k)'s last read + # precedes b(k)'s first write in the same logical iteration. + program = _make_parent_writer_nested_reader_program(versions, inner_extent, single_trip_writer) + assert ("a", "b") not in _insert_sync_alias_pairs(program)