diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 3fe14d22e5..7342a0dc84 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -1,5 +1,6 @@ """Special-register dependencies participate in scheduling and core placement.""" +import inspect from unittest.mock import patch import pytest @@ -160,21 +161,56 @@ def test_cube_mode_schedule_preserves_iteration_order(mode, enable_offset, start tvm.ir.assert_structural_equal(child.body.body, operations[index]) -def test_manual_stages_break_scalar_pipe_ties_by_iteration(): +@pytest.mark.parametrize("stages", [(0, 0), (0, 1), (6, 10), (1, 0)], ids=["same", "ascending", "nonadjacent", "descending"]) +def test_manual_scalar_issue_order_survives_stage_ties(stages): operations = [tirx.Evaluate(T.set_hf32_mode("nearest_even")), tirx.Evaluate(T.set_mmad_direction("n"))] - body = seq(*(unit(op, core=0, stage=stage, cost=(1, 1)) for stage, op in enumerate(operations))) + body = seq(*(unit(op, core=0, stage=stage, cost=(1, 1)) for stage, op in zip(stages, operations))) loop = tirx.For(tirx.Var("i", "int32"), 0, 3, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) before = kernel(unit(loop, core=None, stage=-1)) # These independent setters share PIPE_Scalar, which has no resource - # exclusion. The real solver chooses II=1, forcing equal issue phases. - # The earlier logical iteration takes precedence over per-pipe source order. + # exclusion. Ascending stages need distinct phases so the reconstruction + # tie-break cannot reverse manual source order; other stage orders may tie. after = transform.AutoSchedule()(before) children = after["main"].body.block.body.body.body.seq - assert [int(child.node["stage"]) for child in children] == [1, 0] - for child, op in zip(children, reversed(operations)): + assert [int(child.node["stage"]) for child in children] == list(stages) + for child, op in zip(children, operations): tvm.ir.assert_structural_equal(child.body.body, op) +@pytest.mark.parametrize("cores", [(1, 4), (1, 2), (2, 4)], ids=["aiv0-aiv1", "aiv0-aic", "aic-aiv1"]) +@pytest.mark.parametrize("in_loop", [False, True], ids=["sequence", "loop"]) +def test_manual_pipe_order_does_not_order_distinct_physical_cores(cores, in_loop): + sources = [tirx.decl_buffer((64,), "float32", name=f"A{i}") for i in range(2)] + buffers = [ + tirx.decl_buffer((64,), "float32", name=f"tile{i}", scope="shared.l1" if core == 2 else "shared.dyn") + for i, core in enumerate(cores) + ] + stages = (0, 1) if in_loop else (0, 0) + body = seq( + *( + unit(copy(source, buffer), core=core, stage=stage, cost=(1, 1)) + for source, buffer, core, stage in zip(sources, buffers, cores, stages) + ) + ) + if in_loop: + loop = tirx.For(tirx.Var("i", "int32"), 0, 3, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) + body = unit(loop, core=None, stage=-1) + name = "z3_schedule_loop_python" if in_loop else "z3_schedule_python" + original = getattr(z3_scheduler, name) + inputs = [] + + def capture(*args, **kwargs): + inputs.append(dict(inspect.signature(original).bind(*args, **kwargs).arguments)) + return original(*args, **kwargs) + + with patch.object(z3_scheduler, name, side_effect=capture): + transform.AutoSchedule()(kernel(body, buffers=buffers, params=sources)) + (actual,) = inputs + # Equal pipe names on separate cores do not impose manual issue order. + # Shared HBM resources may independently require resource exclusion. + assert actual["pipe_order_deps"] == [] + + @pytest.mark.parametrize("mode", ["hf32", "mmad-direction"]) def test_cube_mode_writers_follow_nested_gemm_core(mode): before, setter = _cube_mode_program(mode) diff --git a/tilelang/ascend/language/schedule_hint.py b/tilelang/ascend/language/schedule_hint.py index 51e0fdb371..4c15dfe662 100644 --- a/tilelang/ascend/language/schedule_hint.py +++ b/tilelang/ascend/language/schedule_hint.py @@ -122,10 +122,12 @@ def Stage(stage: int): ``T.Stage`` annotation, AutoSchedule treats that child list as manual. One scope may contain several consecutive statements; every scheduler task materialized from those statements receives the same stage. Unannotated - siblings default to stage 0. Within each hardware pipe, source order is the - manual issue order; tasks on different pipes remain freely reorderable, so - no separate order annotation is required. Nested child lists select - automatic or manual scheduling independently. + siblings default to stage 0. On each shared core and hardware pipe, source + order is the manual issue order. The scheduler separates issue phases when + a stage tie-break would reverse this order. Data and register dependencies + remain in force; an incompatible manual order fails scheduling. Tasks on + different pipes remain freely reorderable subject to those dependencies. + Nested child lists select automatic or manual scheduling independently. ``T.Stage`` must wrap complete scheduler tasks. It must not appear inside ``T.Task``, ``T.SimtVF``, ``T.SimdVF``, an SBlock, or a parallel loop. In diff --git a/tilelang/ascend/transform/z3_scheduler.py b/tilelang/ascend/transform/z3_scheduler.py index 51e9f02bf2..4f771a922b 100644 --- a/tilelang/ascend/transform/z3_scheduler.py +++ b/tilelang/ascend/transform/z3_scheduler.py @@ -427,7 +427,11 @@ def z3_schedule_loop_python( for i, stage in enumerate(manual_stages): solver.add(k_vars[i] == stage) for previous, current in pipe_order_deps: - solver.add(r_vars[previous] <= r_vars[current]) + # Reconstruction places larger stages first at equal physical + # phases. Separate ascending stages to preserve manual source + # order; this gap expresses ordering, not instruction latency. + gap = int(manual_schedule and enable_offset and manual_stages[previous] < manual_stages[current]) + solver.add(r_vars[previous] + gap <= r_vars[current]) buffer_vars = [z3.Int(f"buf_{i}", ctx) for i in range(len(buffer_sizes))] for i in range(len(buffer_sizes)):