From 3c264f6f983787af9ba2f8336030cbbdc223b458 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 11:52:50 +0800 Subject: [PATCH 1/4] [Ascend] Preserve manual pipe order across physical phase ties --- src/ascend/transform/auto_schedule.cc | 2 + .../auto_schedule/test_special_registers.py | 44 ++++++++++++++++--- tilelang/ascend/language/schedule_hint.py | 10 +++-- tilelang/ascend/transform/z3_scheduler.py | 6 ++- 4 files changed, 51 insertions(+), 11 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 1687d6b936..868f8ca40a 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -611,6 +611,7 @@ std::vector> ScheduleBuilder::Z3SchedulePython( resource_deps.emplace_back(i, j); } if (manual_schedule && + (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) != 0 && (NodeResourceMask(nodes[i]) & NodeResourceMask(nodes[j])) != 0) { pipe_order_deps.emplace_back(i, j); } @@ -1035,6 +1036,7 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, resource_deps.emplace_back(i, j); } if (manual_schedule && + (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) != 0 && (NodeResourceMask(nodes[i]) & NodeResourceMask(nodes[j])) != 0) { pipe_order_deps.emplace_back(i, j); } diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 3fe14d22e5..01142a226c 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,52 @@ 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_seq_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)): From ca0b4ef54929fe565622ef8835b3277469814374 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 11:58:47 +0800 Subject: [PATCH 2/4] [Ascend] Format manual pipe-order regression --- testing/ascend/auto_schedule/test_special_registers.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 01142a226c..5795eeb820 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -186,8 +186,12 @@ def test_manual_pipe_order_does_not_order_distinct_physical_cores(cores, in_loop 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))) + 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) From da7716998c36761c2845377c1de2e3e6b20eac8a Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 12:04:11 +0800 Subject: [PATCH 3/4] [Ascend] Reuse the existing physical core and pipe masks --- src/ascend/transform/auto_schedule.cc | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 868f8ca40a..1687d6b936 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -611,7 +611,6 @@ std::vector> ScheduleBuilder::Z3SchedulePython( resource_deps.emplace_back(i, j); } if (manual_schedule && - (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) != 0 && (NodeResourceMask(nodes[i]) & NodeResourceMask(nodes[j])) != 0) { pipe_order_deps.emplace_back(i, j); } @@ -1036,7 +1035,6 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, resource_deps.emplace_back(i, j); } if (manual_schedule && - (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) != 0 && (NodeResourceMask(nodes[i]) & NodeResourceMask(nodes[j])) != 0) { pipe_order_deps.emplace_back(i, j); } From 9ec3fc08a8b5d3884621e71ee4ef3885c7b2236f Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 12:05:47 +0800 Subject: [PATCH 4/4] [Ascend] Capture the sequence solver in core-order coverage --- testing/ascend/auto_schedule/test_special_registers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 5795eeb820..7342a0dc84 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -195,7 +195,7 @@ def test_manual_pipe_order_does_not_order_distinct_physical_cores(cores, in_loop 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_seq_python" + name = "z3_schedule_loop_python" if in_loop else "z3_schedule_python" original = getattr(z3_scheduler, name) inputs = []