From 2d483166ad6805077f856b8e9e7215aa2c97895e Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Sun, 18 Oct 2026 03:51:10 +0800 Subject: [PATCH 01/10] [Ascend] Order adjacent manual Cube configurations by physical phase --- src/ascend/transform/auto_schedule.cc | 130 ++++++++++++++++-- .../auto_schedule/test_special_registers.py | 88 ++++++++++++ tilelang/ascend/language/schedule_hint.py | 6 + tilelang/ascend/transform/z3_scheduler.py | 23 ++++ 4 files changed, 232 insertions(+), 15 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 1687d6b936..8bcea5b156 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -61,6 +61,7 @@ #include "./auto_schedule/scheduled_tir.h" #include "./auto_schedule/task_analysis.h" #include "./auto_schedule/task_annotations.h" +#include "ascend/op/builtin.h" #include "ascend/op/copy.h" #include "ascend/op/utils.h" #include "ascend/transform/buffer_version.h" @@ -80,6 +81,87 @@ using namespace ascend; namespace { +bool IsCubeConfiguration(const Stmt &stmt) { + if (const auto *attr = stmt.as()) + return IsCubeConfiguration(attr->body); + if (const auto *seq = stmt.as()) { + return std::all_of(seq->seq.begin(), seq->seq.end(), [](const Stmt &child) { + return IsCubeConfiguration(child); + }); + } + const auto *evaluate = stmt.as(); + if (!evaluate) + return false; + const auto *call = evaluate->value.as(); + return call && (call->op.same_as(tl::ascend_set_hf32_mode()) || + call->op.same_as(tl::ascend_set_mmad_direction())); +} + +// A physical mode has no versions. Only move complete, unconditional config +// regions from stage s to s+1; their source order also preserves the exit mode. +uint16_t ManualRegisterPhaseMask(const std::vector &nodes, + const std::vector &stages) { + for (const IRStructure *node : nodes) { + if (!node->IsTask() || node->ContainsLoopBreak()) + return 0; + const auto *task = static_cast(node); + if (task->IsPerCoreTask()) + return 0; + bool opaque = false; + PostOrderVisit(task->stmt, [&](const ObjectRef &object) { + if (const auto *call = object.as()) { + opaque |= !call->op.as() || + call->op.same_as(builtin::call_extern()) || + call->op.same_as(builtin::tvm_call_packed()) || + call->op.same_as(builtin::tvm_call_packed_lowered()) || + call->op.same_as(builtin::tvm_call_cpacked()) || + call->op.same_as(builtin::tvm_call_cpacked_lowered()) || + call->op.same_as(builtin::call_llvm_intrin()); + } + }); + if (opaque) + return 0; + } + + uint16_t mask = 0; + for (SpecialRegister reg : + {SpecialRegister::kHf32Mode, SpecialRegister::kMmadDirection}) { + uint16_t bit = static_cast(reg); + int64_t first_stage = -1; + int64_t writer_stage = -1; + bool valid = true; + for (size_t i = 0; i < nodes.size(); ++i) { + const IRStructure *node = nodes[i]; + bool writes = (node->GetSpecialWriteMask() & bit) != 0; + bool reads = (node->GetSpecialReadMask() & bit) != 0; + if (!writes && !reads) + continue; + if (node->HasConditions()) { + valid = false; + break; + } + if (writes) { + if (!IsCubeConfiguration(static_cast(node)->stmt) || + (writer_stage >= 0 && + (stages[i] < writer_stage || stages[i] > first_stage + 1))) { + valid = false; + break; + } + if (first_stage < 0) + first_stage = stages[i]; + writer_stage = stages[i]; + } + if (reads && (writer_stage < 0 || stages[i] != writer_stage)) { + valid = false; + break; + } + } + if (valid && first_stage >= 0 && writer_stage == first_stage + 1) + mask |= bit; + } + return mask; +} + // AutoSchedule target resource limits. // Check whether a kernel body contains any SIMT_VF blocks. SIMT_VF @@ -786,6 +868,8 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, std::vector resource_flags; std::vector> data_deps; // (i, j, distance, latency) + std::vector> + register_phase_deps; // (i, j, physical wrap, latency) std::vector> resource_deps; std::vector> owner_exclusion_deps; @@ -952,35 +1036,48 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, } } - // Special-register ordering: a node writing register R must be ordered - // against every node reading R (RAW forward, WAR backward) and every other - // writer of R (WAW), preserving program order. loop_break writes - // kLoopControl, which every node reads, so it fences against all siblings; - // hf32/pad match only their consumers. The distance bit encodes intra-iter - // (0) vs next-iter (1) direction. + // Config regions in adjacent manual stages share one physical register. + // Keep logical dependencies for all other special state, including fences. + uint16_t phase_mask = manual_schedule && enable_offset + ? ManualRegisterPhaseMask(nodes, manual_stages) + : 0; + auto add_register_dep = [&](size_t u, size_t v, int wrap, int64_t latency, + uint16_t registers) { + if (registers & ~phase_mask) + data_deps.emplace_back(u, v, wrap, latency); + if (registers & phase_mask) { + // Reconstruction issues larger stages first at equal physical times. + // A forward cross-stage WAR must therefore be strictly ordered. + if (wrap == 0 && manual_stages[u] < manual_stages[v]) + latency = std::max(latency, 1); + register_phase_deps.emplace_back(u, v, wrap, latency); + } + }; for (size_t i = 0; i < n; ++i) { uint16_t wmask = nodes[i]->GetSpecialWriteMask(); if (wmask == 0) continue; for (size_t j = 0; j < n; ++j) { - if (j == i || (wmask & nodes[j]->GetSpecialReadMask()) == 0 || + uint16_t overlap = wmask & nodes[j]->GetSpecialReadMask(); + if (j == i || overlap == 0 || (!(wmask & static_cast(SpecialRegister::kLoopControl)) && (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) == 0)) continue; if (i < j) { - data_deps.emplace_back(i, j, 0, latencies[i]); - data_deps.emplace_back(j, i, 1, 0); + add_register_dep(i, j, 0, latencies[i], overlap); + add_register_dep(j, i, 1, 0, overlap); } else { - data_deps.emplace_back(i, j, 1, latencies[i]); - data_deps.emplace_back(j, i, 0, 0); + add_register_dep(i, j, 1, latencies[i], overlap); + add_register_dep(j, i, 0, 0, overlap); } } for (size_t j = i + 1; j < n; ++j) { - if ((wmask & nodes[j]->GetSpecialWriteMask()) == 0 || + uint16_t overlap = wmask & nodes[j]->GetSpecialWriteMask(); + if (overlap == 0 || (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) == 0) continue; - data_deps.emplace_back(i, j, 0, latencies[i]); - data_deps.emplace_back(j, i, 1, latencies[j]); + add_register_dep(i, j, 0, latencies[i], overlap); + add_register_dep(j, i, 1, latencies[j], overlap); } } @@ -1092,6 +1189,9 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, ffi::Array> tvm_pipe_order_deps; ffi::Array tvm_storage_sizes; ffi::Array tvm_manual_stages; + ffi::Array> tvm_register_phase_deps; + for (const auto &[u, v, wrap, latency] : register_phase_deps) + tvm_register_phase_deps.push_back({u, v, wrap, latency}); for (auto val : latencies) { tvm_latencies.push_back(val); @@ -1165,7 +1265,7 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, tvm_data_deps, tvm_resource_deps, tvm_owner_exclusion_deps, tvm_pipe_order_deps, tvm_storage_sizes, tvm_memory_groups, tvm_stage_order_deps, enable_offset, tvm_manual_stages, - manual_schedule) + manual_schedule, tvm_register_phase_deps) .cast, ffi::Array, int64_t>>(); ffi::Array start_times = return_val.get<0>(); diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 3fe14d22e5..f038f3a0d7 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -181,3 +181,91 @@ def test_cube_mode_writers_follow_nested_gemm_core(mode): after = transform.ResolveCore()(before) assert len(calls(after, setter)) == 2 assert task_core_masks(after, setter) == {2} + + +def _manual_cube_mode_loop(stages, mode, *, guarded_writer=False, mixed_modes=False): + a = tirx.decl_buffer((16, 16), "float32", name="a", scope="shared.l0a") + b = tirx.decl_buffer((16, 16), "float32", name="b", scope="shared.l0b") + outputs = [tirx.decl_buffer((16, 16), "float32", name=f"c{i}", scope="shared.l0c") for i in range(2)] + setters = ( + [T.set_hf32_mode("nearest_even"), T.set_hf32_mode(None)] + if mode == "hf32" + else [T.set_mmad_direction("m"), T.set_mmad_direction("n")] + ) + i = tirx.Var("i", "int32") + operations = [ + tirx.Evaluate(setters[0]), + tirx.Evaluate(T.gemm(a, b, outputs[0], transpose_B=True, clear_accum=True)), + tirx.Evaluate(setters[1]), + tirx.Evaluate(T.gemm(a, b, outputs[1], transpose_B=True, clear_accum=True)), + ] + if mixed_modes: + # HF32 has only its stage-4 writer, but its second reader is stage 5. + # MMAD's two complete regions can use phases; HF32 must stay logical. + operations[0] = seq(tirx.Evaluate(T.set_hf32_mode(None)), operations[0]) + body = seq( + *( + unit( + op, + core=0, + stage=stage, + cost=(1, 1) if index % 2 == 0 else (64, 32), + guard=i == 0 if guarded_writer and index == 2 else None, + ) + for index, (op, stage) in enumerate(zip(operations, stages)) + ) + ) + loop = tirx.For(i, 0, 3, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) + return kernel(unit(loop, core=None, stage=-1), buffers=[a, b, *outputs]), operations + + +@pytest.mark.parametrize("mode", ["hf32", "mmad-direction"]) +@pytest.mark.parametrize("base_stage", [0, 4, 12]) +def test_manual_cube_mode_regions_use_physical_order(mode, base_stage): + stages = [base_stage, base_stage, base_stage + 1, base_stage + 1] + before, operations = _manual_cube_mode_loop(stages, mode) + # Use the real solver: fixed logical 0/1 WAW distances make this infeasible. + after = transform.AutoSchedule()(before) + children = after["main"].body.block.body.body.body.seq + assert [int(child.node["stage"]) for child in children] == stages + for child, operation in zip(children, operations): + tvm.ir.assert_structural_equal(child.body.body, operation) + + +@pytest.mark.parametrize( + "stages, guarded_writer", + [ + pytest.param([5, 5, 4, 4], False, id="exit-mode"), + pytest.param([4, 5, 5, 5], False, id="reader-crosses-stage"), + pytest.param([4, 4, 5, 5], True, id="conditional-writer"), + pytest.param([4, 4, 6, 6], False, id="nonadjacent-stages"), + ], +) +def test_manual_cube_mode_unsafe_regions_keep_logical_dependencies(stages, guarded_writer): + before, _ = _manual_cube_mode_loop(stages, "mmad-direction", guarded_writer=guarded_writer) + with pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"): + transform.AutoSchedule()(before) + + +def test_manual_cube_modes_preserve_other_register_dependencies(): + before, _ = _manual_cube_mode_loop([4, 4, 5, 5], "mmad-direction", mixed_modes=True) + inputs = [] + original = z3_scheduler.z3_schedule_loop_python + + def capture(*args, **kwargs): + inputs.append((args[4], kwargs["register_phase_deps"])) + return original(*args, **kwargs) + + with ( + patch.object(z3_scheduler, "z3_schedule_loop_python", side_effect=capture), + pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"), + ): + transform.AutoSchedule()(before) + assert len(inputs) == 1 + logical, physical = inputs[0] + # The same writer/reader pair must retain HF32's logical dependency as + # well as MMAD's physical dependency; classifying whole pairs loses one. + assert (0, 3, 0, 1) in logical + assert (0, 3, 0, 1) in physical + # At equal phases the stage-5 setter would precede the stage-4 GEMM. + assert (1, 2, 0, 1) in physical diff --git a/tilelang/ascend/language/schedule_hint.py b/tilelang/ascend/language/schedule_hint.py index 51e0fdb371..a1a95e6df8 100644 --- a/tilelang/ascend/language/schedule_hint.py +++ b/tilelang/ascend/language/schedule_hint.py @@ -145,6 +145,12 @@ def Stage(stage: int): ``T.Pipelined(..., num_stages=...)`` controls the automatic buffer-version search bound; it does not limit the values accepted by ``T.Stage``. + Unconditional HF32/MMAD configuration regions in adjacent stages may share + the physical register in source issue order. Each consumer must use its + preceding setter's stage, and writer stages must not decrease. This keeps + both consumer modes and the loop's outgoing mode unchanged. Nested or + conditional configuration regions retain logical iteration dependencies. + Example ------- >>> for i in T.Pipelined(8, annotations={"enable_offset": True}): diff --git a/tilelang/ascend/transform/z3_scheduler.py b/tilelang/ascend/transform/z3_scheduler.py index 51e9f02bf2..19c2675bdf 100644 --- a/tilelang/ascend/transform/z3_scheduler.py +++ b/tilelang/ascend/transform/z3_scheduler.py @@ -283,6 +283,7 @@ def z3_schedule_loop_python( manual_schedule: bool = False, seed: int | None = 42, verbose: bool = False, + register_phase_deps: list[tuple[int, int, int, int]] | None = None, ) -> tuple[list[int], list[int], int]: """Schedule a loop body with modulo, stage, and storage constraints. @@ -340,6 +341,11 @@ def z3_schedule_loop_python( verbose : bool Print solver inputs, search progress, and the selected schedule. + register_phase_deps : list[tuple[int, int, int, int]] | None + Manual configuration ordering (u, v, physical_wrap, gap), requiring + r[v] + physical_wrap * II >= r[u] + gap. These do not carry buffer + versions or logical iteration distances. + Returns ------- tuple[list[int], list[int], int] @@ -369,12 +375,22 @@ def z3_schedule_loop_python( if not enable_offset and any(stage != 0 for stage in manual_stages): raise ValueError("Non-zero manual stages require enable_offset=True") + if register_phase_deps is None: + register_phase_deps = [] + if register_phase_deps: + if not manual_schedule or not enable_offset: + raise ValueError("Register phase dependencies require a manual offset schedule") + for u, v, wrap, gap in register_phase_deps: + if not (0 <= u < n and 0 <= v < n and wrap in (0, 1) and gap >= 0): + raise ValueError(f"Invalid register phase dependency: {(u, v, wrap, gap)}") + if verbose: print(f"[Python Z3 Loop] Starting scheduling for {n} tasks") print(f"[Python Z3 Loop] Latencies: {latencies}") print(f"[Python Z3 Loop] IIs: {iis}") print(f"[Python Z3 Loop] Resource flags: {resource_flags}") print(f"[Python Z3 Loop] Data dependencies with distances: {data_deps}") + print(f"[Python Z3 Loop] Register phase dependencies: {register_phase_deps}") print(f"[Python Z3 Loop] Resource dependencies: {resource_deps}") print(f"[Python Z3 Loop] Owner exclusions: {owner_exclusion_deps}") print(f"[Python Z3 Loop] Stage-order dependencies: {stage_order_deps}") @@ -439,6 +455,9 @@ def z3_schedule_loop_python( if idxs: solver.add(z3.Sum([buffer_vars[i] * buffer_sizes[i] for i in idxs]) <= cap) + for u, v, wrap, gap in register_phase_deps: + solver.add(r_vars[v] + wrap * ii_mid >= r_vars[u] + gap) + # Add data dependency constraints with distance for u, v, distance, latency in data_deps: if distance >= 0: @@ -614,6 +633,7 @@ def z3_schedule_loop_python( "iis": iis, "resource_flags": resource_flags, "data_dependencies": data_deps, + "register_phase_dependencies": register_phase_deps, "resource_dependencies": resource_deps, "owner_exclusions": owner_exclusion_deps, "manual_schedule": manual_schedule, @@ -797,6 +817,7 @@ def z3_schedule_loop_ffi( enable_offset=False, manual_stages=None, manual_schedule=False, + register_phase_deps=None, ): """FFI wrapper for z3_schedule_loop_python. @@ -818,6 +839,7 @@ def z3_schedule_loop_ffi( if memory_groups is not None: for i in range(len(memory_groups)): memory_groups_list.append(list(memory_groups[i])) + register_phase_deps_list = _convert_int_tuples(register_phase_deps, 4) data_deps_list = _convert_int_tuples(data_deps, 4) resource_deps_list = _convert_int_tuples(resource_deps, 2) stage_order_deps_list = _convert_int_tuples(stage_order_deps, 2) @@ -846,6 +868,7 @@ def z3_schedule_loop_ffi( enable_offset=bool(enable_offset), manual_stages=manual_stages_list, manual_schedule=bool(manual_schedule), + register_phase_deps=register_phase_deps_list, seed=42, ) From 617b3f3e7ec3182355bb75066293642b7d8a58d2 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Sun, 18 Oct 2026 03:55:35 +0800 Subject: [PATCH 02/10] [Ascend] Keep unrelated nested work transparent to Cube mode phases --- src/ascend/transform/auto_schedule.cc | 50 +++++++++++-------- .../auto_schedule/test_special_registers.py | 21 +++++++- 2 files changed, 50 insertions(+), 21 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 8bcea5b156..aed0fd84c3 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -101,26 +101,35 @@ bool IsCubeConfiguration(const Stmt &stmt) { // regions from stage s to s+1; their source order also preserves the exit mode. uint16_t ManualRegisterPhaseMask(const std::vector &nodes, const std::vector &stages) { - for (const IRStructure *node : nodes) { - if (!node->IsTask() || node->ContainsLoopBreak()) - return 0; - const auto *task = static_cast(node); - if (task->IsPerCoreTask()) - return 0; - bool opaque = false; - PostOrderVisit(task->stmt, [&](const ObjectRef &object) { - if (const auto *call = object.as()) { - opaque |= !call->op.as() || - call->op.same_as(builtin::call_extern()) || - call->op.same_as(builtin::tvm_call_packed()) || - call->op.same_as(builtin::tvm_call_packed_lowered()) || - call->op.same_as(builtin::tvm_call_cpacked()) || - call->op.same_as(builtin::tvm_call_cpacked_lowered()) || - call->op.same_as(builtin::call_llvm_intrin()); - } - }); - if (opaque) + constexpr uint16_t kCubeModes = + static_cast(SpecialRegister::kHf32Mode) | + static_cast(SpecialRegister::kMmadDirection); + for (IRStructure *node : nodes) { + if (node->ContainsLoopBreak()) return 0; + std::vector tasks; + CollectAllTaskNodes(node, tasks); + for (const TaskNode *task : tasks) { + // PerCoreTask's aggregate write mask does not include configuration + // writes. Inspect its body before treating unrelated work as transparent. + if (task->IsPerCoreTask() && + (AnalyzeSpecialRegisterWrites(task->stmt) & kCubeModes)) + return 0; + bool opaque = false; + PostOrderVisit(task->stmt, [&](const ObjectRef &object) { + if (const auto *call = object.as()) { + opaque |= !call->op.as() || + call->op.same_as(builtin::call_extern()) || + call->op.same_as(builtin::tvm_call_packed()) || + call->op.same_as(builtin::tvm_call_packed_lowered()) || + call->op.same_as(builtin::tvm_call_cpacked()) || + call->op.same_as(builtin::tvm_call_cpacked_lowered()) || + call->op.same_as(builtin::call_llvm_intrin()); + } + }); + if (opaque) + return 0; + } } uint16_t mask = 0; @@ -136,7 +145,8 @@ uint16_t ManualRegisterPhaseMask(const std::vector &nodes, bool reads = (node->GetSpecialReadMask() & bit) != 0; if (!writes && !reads) continue; - if (node->HasConditions()) { + if (!node->IsTask() || node->HasConditions() || + static_cast(node)->IsPerCoreTask()) { valid = false; break; } diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index f038f3a0d7..803077b97d 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -183,7 +183,7 @@ def test_cube_mode_writers_follow_nested_gemm_core(mode): assert task_core_masks(after, setter) == {2} -def _manual_cube_mode_loop(stages, mode, *, guarded_writer=False, mixed_modes=False): +def _manual_cube_mode_loop(stages, mode, *, guarded_writer=False, mixed_modes=False, nested=None): a = tirx.decl_buffer((16, 16), "float32", name="a", scope="shared.l0a") b = tirx.decl_buffer((16, 16), "float32", name="b", scope="shared.l0b") outputs = [tirx.decl_buffer((16, 16), "float32", name=f"c{i}", scope="shared.l0c") for i in range(2)] @@ -215,6 +215,9 @@ def _manual_cube_mode_loop(stages, mode, *, guarded_writer=False, mixed_modes=Fa for index, (op, stage) in enumerate(zip(operations, stages)) ) ) + if nested is not None: + inner = tirx.For(tirx.Var("j", "int32"), 0, 2, tirx.ForKind.SERIAL, unit(tirx.Evaluate(nested), core=0, stage=0, cost=(1, 1))) + body = seq(body, unit(inner, core=None, stage=stages[-1])) loop = tirx.For(i, 0, 3, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) return kernel(unit(loop, core=None, stage=-1), buffers=[a, b, *outputs]), operations @@ -269,3 +272,19 @@ def capture(*args, **kwargs): assert (0, 3, 0, 1) in physical # At equal phases the stage-5 setter would precede the stage-4 GEMM. assert (1, 2, 0, 1) in physical + + +@pytest.mark.parametrize("nested", ["unrelated", "mode-writer", "opaque"]) +def test_manual_cube_mode_regions_inspect_nested_effects(nested): + operation = { + "unrelated": T.set_atomic_none(), + "mode-writer": T.set_mmad_direction("m"), + "opaque": tirx.call_extern("int32", "custom_stateful_call"), + }[nested] + before, _ = _manual_cube_mode_loop([4, 4, 5, 5], "mmad-direction", nested=operation) + if nested == "unrelated": + after = transform.AutoSchedule()(before) + assert len(calls(after, "tl.ascend_set_mmad_direction")) == 2 + else: + with pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"): + transform.AutoSchedule()(before) From 950811018176abc7bc7160f8c1a19d8e469325ac Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Sun, 18 Oct 2026 04:06:45 +0800 Subject: [PATCH 03/10] [Ascend] Preserve Cube mode phases across nested consumers --- src/ascend/transform/auto_schedule.cc | 83 +++++++++++++------ .../auto_schedule/test_special_registers.py | 62 ++++++++++++-- 2 files changed, 114 insertions(+), 31 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index aed0fd84c3..2a6b3aca37 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -104,32 +104,59 @@ uint16_t ManualRegisterPhaseMask(const std::vector &nodes, constexpr uint16_t kCubeModes = static_cast(SpecialRegister::kHf32Mode) | static_cast(SpecialRegister::kMmadDirection); - for (IRStructure *node : nodes) { + auto has_opaque_call = [](const ObjectRef &ir) { + bool opaque = false; + PostOrderVisit(ir, [&](const ObjectRef &object) { + if (const auto *call = object.as()) { + opaque |= !call->op.as() || + call->op.same_as(builtin::call_extern()) || + call->op.same_as(builtin::tvm_call_packed()) || + call->op.same_as(builtin::tvm_call_packed_lowered()) || + call->op.same_as(builtin::tvm_call_cpacked()) || + call->op.same_as(builtin::tvm_call_cpacked_lowered()) || + call->op.same_as(builtin::call_llvm_intrin()); + } + }); + return opaque; + }; + auto has_unsafe_effect = [&](auto &&self, const IRStructure *node) -> bool { if (node->ContainsLoopBreak()) - return 0; - std::vector tasks; - CollectAllTaskNodes(node, tasks); - for (const TaskNode *task : tasks) { - // PerCoreTask's aggregate write mask does not include configuration - // writes. Inspect its body before treating unrelated work as transparent. - if (task->IsPerCoreTask() && - (AnalyzeSpecialRegisterWrites(task->stmt) & kCubeModes)) - return 0; - bool opaque = false; - PostOrderVisit(task->stmt, [&](const ObjectRef &object) { - if (const auto *call = object.as()) { - opaque |= !call->op.as() || - call->op.same_as(builtin::call_extern()) || - call->op.same_as(builtin::tvm_call_packed()) || - call->op.same_as(builtin::tvm_call_packed_lowered()) || - call->op.same_as(builtin::tvm_call_cpacked()) || - call->op.same_as(builtin::tvm_call_cpacked_lowered()) || - call->op.same_as(builtin::call_llvm_intrin()); + return true; + for (const auto &guard : node->GetGuards()) { + if (guard->IsCondition()) { + if (has_opaque_call( + static_cast(guard.get())->condition)) + return true; + } else { + const auto *attribute = + static_cast(guard.get()); + if (has_opaque_call(attribute->value)) + return true; + if (const auto *expr = attribute->node.as()) { + if (has_opaque_call(GetRef(expr))) + return true; } - }); - if (opaque) - return 0; + } } + if (node->IsControl()) { + const auto *control = static_cast(node); + // The header and guards can change modes even when the body only reads. + if (has_opaque_call(control->task->stmt)) + return true; + return std::any_of(control->children.begin(), control->children.end(), + [&](const std::shared_ptr &child) { + return self(self, child.get()); + }); + } + const auto *task = static_cast(node); + // PerCoreTask's aggregate write mask does not include configuration writes. + return has_opaque_call(task->stmt) || + (task->IsPerCoreTask() && + (AnalyzeSpecialRegisterWrites(task->stmt) & kCubeModes)); + }; + for (const IRStructure *node : nodes) { + if (has_unsafe_effect(has_unsafe_effect, node)) + return 0; } uint16_t mask = 0; @@ -145,13 +172,15 @@ uint16_t ManualRegisterPhaseMask(const std::vector &nodes, bool reads = (node->GetSpecialReadMask() & bit) != 0; if (!writes && !reads) continue; - if (!node->IsTask() || node->HasConditions() || - static_cast(node)->IsPerCoreTask()) { + if (node->HasConditions() || + (node->IsTask() && + static_cast(node)->IsPerCoreTask())) { valid = false; break; } if (writes) { - if (!IsCubeConfiguration(static_cast(node)->stmt) || + if (!node->IsTask() || + !IsCubeConfiguration(static_cast(node)->stmt) || (writer_stage >= 0 && (stages[i] < writer_stage || stages[i] > first_stage + 1))) { valid = false; @@ -161,6 +190,8 @@ uint16_t ManualRegisterPhaseMask(const std::vector &nodes, first_stage = stages[i]; writer_stage = stages[i]; } + // A loop that only reads this mode stays in its writer's region as one + // outer unit. Its inner stages and buffer versions do not change the mode. if (reads && (writer_stage < 0 || stages[i] != writer_stage)) { valid = false; break; diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 803077b97d..049c6ebd57 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -183,7 +183,7 @@ def test_cube_mode_writers_follow_nested_gemm_core(mode): assert task_core_masks(after, setter) == {2} -def _manual_cube_mode_loop(stages, mode, *, guarded_writer=False, mixed_modes=False, nested=None): +def _manual_cube_mode_loop(stages, mode, *, guarded_writer=False, mixed_modes=False, nested=None, nested_readers=False, nested_effect=None): a = tirx.decl_buffer((16, 16), "float32", name="a", scope="shared.l0a") b = tirx.decl_buffer((16, 16), "float32", name="b", scope="shared.l0b") outputs = [tirx.decl_buffer((16, 16), "float32", name=f"c{i}", scope="shared.l0c") for i in range(2)] @@ -203,8 +203,29 @@ def _manual_cube_mode_loop(stages, mode, *, guarded_writer=False, mixed_modes=Fa # HF32 has only its stage-4 writer, but its second reader is stage 5. # MMAD's two complete regions can use phases; HF32 must stay logical. operations[0] = seq(tirx.Evaluate(T.set_hf32_mode(None)), operations[0]) - body = seq( - *( + units = [] + for index, (op, stage) in enumerate(zip(operations, stages)): + if nested_readers and index % 2: + extent = 4 + guard = None + if index == 1 and nested_effect == "opaque-header": + extent = tirx.call_extern("int32", "change_mode_and_return_count") + if index == 1 and nested_effect == "opaque-guard": + guard = tirx.call_extern("bool", "change_mode_and_return_condition") + inner_body = unit(op, core=0, stage=0, cost=(64, 32), guard=guard) + if index == 1 and nested_effect == "mode-writer": + inner_body = seq(unit(tirx.Evaluate(setters[0]), core=0, stage=0, cost=(1, 1)), inner_body) + inner = tirx.For( + tirx.Var(f"j{index}", "int32"), + 0, + extent, + tirx.ForKind.SERIAL, + inner_body, + annotations={"num_stages": 2}, + ) + units.append(unit(inner, core=None, stage=stage)) + continue + units.append( unit( op, core=0, @@ -212,9 +233,8 @@ def _manual_cube_mode_loop(stages, mode, *, guarded_writer=False, mixed_modes=Fa cost=(1, 1) if index % 2 == 0 else (64, 32), guard=i == 0 if guarded_writer and index == 2 else None, ) - for index, (op, stage) in enumerate(zip(operations, stages)) ) - ) + body = seq(*units) if nested is not None: inner = tirx.For(tirx.Var("j", "int32"), 0, 2, tirx.ForKind.SERIAL, unit(tirx.Evaluate(nested), core=0, stage=0, cost=(1, 1))) body = seq(body, unit(inner, core=None, stage=stages[-1])) @@ -235,6 +255,38 @@ def test_manual_cube_mode_regions_use_physical_order(mode, base_stage): tvm.ir.assert_structural_equal(child.body.body, operation) +@pytest.mark.parametrize("mode", ["hf32", "mmad-direction"]) +def test_manual_cube_mode_regions_keep_nested_consumers(mode): + stages = [4, 4, 5, 5] + before, _ = _manual_cube_mode_loop(stages, mode, nested_readers=True) + phases = [] + original = z3_scheduler.z3_schedule_loop_python + + def capture(*args, **kwargs): + phases.append(kwargs["register_phase_deps"]) + return original(*args, **kwargs) + + with patch.object(z3_scheduler, "z3_schedule_loop_python", side_effect=capture): + after = transform.AutoSchedule()(before) + children = after["main"].body.block.body.body.body.seq + assert [int(child.node["stage"]) for child in children] == stages + for index in (1, 3): + inner = children[index].body + assert isinstance(inner, tirx.For) + assert int(inner.extent) == 4 + assert len(calls(inner, "tl.tileop.gemm")) == 1 + # The whole first loop must precede the next stage's configuration. At a + # tie reconstruction would emit the stage-5 setter before that loop. + assert (1, 2, 0, 1) in phases[-1] + + +@pytest.mark.parametrize("effect", ["mode-writer", "opaque-header", "opaque-guard"]) +def test_manual_cube_mode_nested_consumers_reject_mode_changes(effect): + before, _ = _manual_cube_mode_loop([4, 4, 5, 5], "mmad-direction", nested_readers=True, nested_effect=effect) + with pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"): + transform.AutoSchedule()(before) + + @pytest.mark.parametrize( "stages, guarded_writer", [ From 2837a7837d7570279570b66cd5a430e3309c5ae1 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 11:50:01 +0800 Subject: [PATCH 04/10] [Ascend] Scope register overlap checks to the matched register --- src/ascend/transform/auto_schedule.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 2a6b3aca37..223e6c37ed 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -1101,7 +1101,7 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, for (size_t j = 0; j < n; ++j) { uint16_t overlap = wmask & nodes[j]->GetSpecialReadMask(); if (j == i || overlap == 0 || - (!(wmask & static_cast(SpecialRegister::kLoopControl)) && + (!(overlap & static_cast(SpecialRegister::kLoopControl)) && (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) == 0)) continue; if (i < j) { From d425fe8274b646f44aee0cdb418270c713487958 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 11:51:16 +0800 Subject: [PATCH 05/10] [Ascend] Format Cube phase comment --- src/ascend/transform/auto_schedule.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 223e6c37ed..0657d8294f 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -191,7 +191,8 @@ uint16_t ManualRegisterPhaseMask(const std::vector &nodes, writer_stage = stages[i]; } // A loop that only reads this mode stays in its writer's region as one - // outer unit. Its inner stages and buffer versions do not change the mode. + // outer unit. Its inner stages and buffer versions do not change the + // mode. if (reads && (writer_stage < 0 || stages[i] != writer_stage)) { valid = false; break; From 34cb234d20425cf914a8b482f713a151ad869c78 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 16:06:12 +0800 Subject: [PATCH 06/10] [Ascend] Unify single-state pipeline dependencies --- src/ascend/transform/auto_schedule.cc | 460 +++++++++++++----- .../test_schedule_integration.py | 44 ++ .../auto_schedule/test_special_registers.py | 150 ++++-- tilelang/ascend/language/schedule_hint.py | 13 +- tilelang/ascend/transform/z3_scheduler.py | 35 +- 5 files changed, 526 insertions(+), 176 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 0657d8294f..8109b0bd27 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -81,127 +81,277 @@ using namespace ascend; namespace { -bool IsCubeConfiguration(const Stmt &stmt) { - if (const auto *attr = stmt.as()) - return IsCubeConfiguration(attr->body); - if (const auto *seq = stmt.as()) { - return std::all_of(seq->seq.begin(), seq->seq.end(), [](const Stmt &child) { - return IsCubeConfiguration(child); - }); +// Configuration values are state; loop-control and pipe registers are fences. +constexpr uint16_t kConfigurationRegisters = + static_cast(SpecialRegister::kHf32Mode) | + static_cast(SpecialRegister::kMmadDirection) | + static_cast(SpecialRegister::kPadValue) | + static_cast(SpecialRegister::kAtomicMode); + +using PhaseDependency = std::tuple; + +struct SingleState { + uint16_t reg{0}; + Optional storage; + + bool Reads(const IRStructure *node) const { + return storage.defined() ? node->ReadsStorage(storage.value()) + : (node->GetSpecialReadMask() & reg) != 0; + } + bool Writes(const IRStructure *node) const { + return storage.defined() ? node->WritesStorage(storage.value()) + : (node->GetSpecialWriteMask() & reg) != 0; + } + bool MustWrite(const Stmt &stmt) const { + if (const auto *attr = stmt.as()) + return MustWrite(attr->body); + if (const auto *seq = stmt.as()) + return std::any_of(seq->seq.begin(), seq->seq.end(), + [&](const Stmt &child) { return MustWrite(child); }); + if (const auto *branch = stmt.as()) + return branch->else_case.defined() && MustWrite(branch->then_case) && + MustWrite(branch->else_case.value()); + const auto *evaluate = stmt.as(); + const auto *call = evaluate ? evaluate->value.as() : nullptr; + if (!call || !(AnalyzeSpecialRegisterWrites(stmt) & reg)) + return false; + // A write nested in an argument (for example if_then_else) need not + // execute. Only the call's own configuration effect is a definite kill. + for (const PrimExpr &arg : call->args) { + if (AnalyzeSpecialRegisterWrites(Evaluate(arg)) & reg) + return false; + } + return true; } - const auto *evaluate = stmt.as(); - if (!evaluate) - return false; - const auto *call = evaluate->value.as(); - return call && (call->op.same_as(tl::ascend_set_hf32_mode()) || - call->op.same_as(tl::ascend_set_mmad_direction())); +}; + +struct StateAccess { + // Reads of the incoming value, before a guaranteed write in this unit. + bool reads{false}; + bool writes{false}; + bool must_write{false}; +}; + +// Task-local scalar assignments can kill earlier values before a later read. +// Summarize those effects in execution order, including branch intersections. +class ScalarStateAccessAnalyzer : public StmtExprVisitor { +public: + explicit ScalarStateAccessAnalyzer(Var storage) + : storage_(std::move(storage)) {} + + StateAccess Analyze(const Stmt &stmt) { + VisitStmt(stmt); + return access_; + } + +private: + Var storage_; + StateAccess access_; + + void VisitExpr_(const BufferLoadNode *op) final { + if (op->buffer->data.same_as(storage_)) + access_.reads |= !access_.must_write; + StmtExprVisitor::VisitExpr_(op); + } + void VisitExpr_(const VarNode *op) final { + // An address passed to an intrinsic can read the incoming scalar too. + if (GetRef(op).same_as(storage_)) + access_.reads |= !access_.must_write; + } + void VisitStmt_(const BufferStoreNode *op) final { + VisitExpr(op->value); + for (const PrimExpr &index : op->indices) + VisitExpr(index); + if (op->predicate.defined()) + VisitExpr(op->predicate.value()); + if (op->buffer->data.same_as(storage_)) { + access_.writes = true; + access_.must_write |= + !op->predicate.defined() || is_one(op->predicate.value()); + } + } + void VisitStmt_(const IfThenElseNode *op) final { + VisitExpr(op->condition); + bool incoming = access_.must_write; + VisitStmt(op->then_case); + bool then_written = access_.must_write; + access_.must_write = incoming; + if (op->else_case.defined()) + VisitStmt(op->else_case.value()); + access_.must_write &= then_written; + } + void VisitStmt_(const ForNode *op) final { + VisitExpr(op->min); + VisitExpr(op->extent); + bool incoming = access_.must_write; + VisitStmt(op->body); + const auto *extent = op->extent.as(); + access_.must_write = + incoming || (extent && extent->value > 0 && access_.must_write); + } + void VisitStmt_(const SBlockRealizeNode *op) final { + for (const PrimExpr &value : op->iter_values) + VisitExpr(value); + VisitExpr(op->predicate); + bool incoming = access_.must_write; + VisitStmt(op->block); + access_.must_write = + incoming || (is_one(op->predicate) && access_.must_write); + } +}; + +CoreMask NodeCoreMask(const IRStructure *node); + +StateAccess SummarizeState(const IRStructure *node, const SingleState &state, + CoreMask core) { + StateAccess result; + if (!(NodeCoreMask(node) & core)) + return result; + if (node->IsTask()) { + const auto *task = static_cast(node); + if (state.storage.defined()) { + result = + ScalarStateAccessAnalyzer(state.storage.value()).Analyze(task->stmt); + result.writes |= state.Writes(node); + } else { + result = {state.Reads(node), state.Writes(node), + state.MustWrite(task->stmt)}; + } + } else { + const auto *control = static_cast(node); + result = SummarizeState(control->task.get(), state, core); + std::vector children; + for (const auto &child : control->children) + children.push_back(child.get()); + // Match buffer dependency analysis: stages establish logical iteration + // order; the sibling list establishes physical per-pipe issue order. + std::stable_sort(children.begin(), children.end(), + [](const IRStructure *a, const IRStructure *b) { + return a->GetStage() < b->GetStage(); + }); + for (const IRStructure *child : children) { + StateAccess access = SummarizeState(child, state, core); + result.reads |= access.reads && !result.must_write; + result.writes |= access.writes; + result.must_write |= access.must_write; + } + const auto *extent = control->control->extent.as(); + result.must_write &= extent && extent->value > 0; + } + if (state.storage.defined()) + result.reads |= node->GuardsTouchStorage(state.storage.value()); + result.must_write &= !node->HasConditions(); + return result; } -// A physical mode has no versions. Only move complete, unconditional config -// regions from stage s to s+1; their source order also preserves the exit mode. -uint16_t ManualRegisterPhaseMask(const std::vector &nodes, - const std::vector &stages) { - constexpr uint16_t kCubeModes = - static_cast(SpecialRegister::kHf32Mode) | - static_cast(SpecialRegister::kMmadDirection); - auto has_opaque_call = [](const ObjectRef &ir) { - bool opaque = false; +bool HasUnknownStateEffects(const IRStructure *node) { + auto opaque = [](const ObjectRef &ir) { + bool found = false; PostOrderVisit(ir, [&](const ObjectRef &object) { if (const auto *call = object.as()) { - opaque |= !call->op.as() || - call->op.same_as(builtin::call_extern()) || - call->op.same_as(builtin::tvm_call_packed()) || - call->op.same_as(builtin::tvm_call_packed_lowered()) || - call->op.same_as(builtin::tvm_call_cpacked()) || - call->op.same_as(builtin::tvm_call_cpacked_lowered()) || - call->op.same_as(builtin::call_llvm_intrin()); + found |= !call->op.as() || + call->op.same_as(builtin::call_extern()) || + call->op.same_as(builtin::tvm_call_packed()) || + call->op.same_as(builtin::tvm_call_packed_lowered()) || + call->op.same_as(builtin::tvm_call_cpacked()) || + call->op.same_as(builtin::tvm_call_cpacked_lowered()) || + call->op.same_as(builtin::call_llvm_intrin()); } }); - return opaque; + return found; }; - auto has_unsafe_effect = [&](auto &&self, const IRStructure *node) -> bool { - if (node->ContainsLoopBreak()) - return true; - for (const auto &guard : node->GetGuards()) { - if (guard->IsCondition()) { - if (has_opaque_call( - static_cast(guard.get())->condition)) - return true; - } else { - const auto *attribute = - static_cast(guard.get()); - if (has_opaque_call(attribute->value)) + if (node->ContainsLoopBreak()) + return true; + for (const auto &guard : node->GetGuards()) { + if (guard->IsCondition()) { + if (opaque(static_cast(guard.get())->condition)) + return true; + } else { + const auto *attr = static_cast(guard.get()); + if (opaque(attr->value)) + return true; + if (const auto *expr = attr->node.as()) { + if (opaque(GetRef(expr))) return true; - if (const auto *expr = attribute->node.as()) { - if (has_opaque_call(GetRef(expr))) - return true; - } } } - if (node->IsControl()) { - const auto *control = static_cast(node); - // The header and guards can change modes even when the body only reads. - if (has_opaque_call(control->task->stmt)) - return true; - return std::any_of(control->children.begin(), control->children.end(), - [&](const std::shared_ptr &child) { - return self(self, child.get()); - }); - } - const auto *task = static_cast(node); - // PerCoreTask's aggregate write mask does not include configuration writes. - return has_opaque_call(task->stmt) || - (task->IsPerCoreTask() && - (AnalyzeSpecialRegisterWrites(task->stmt) & kCubeModes)); - }; - for (const IRStructure *node : nodes) { - if (has_unsafe_effect(has_unsafe_effect, node)) - return 0; } + if (node->IsTask()) { + const auto *task = static_cast(node); + // PerCoreTask does not expose all internal configuration effects. + return task->IsPerCoreTask() || opaque(task->stmt); + } + const auto *control = static_cast(node); + return opaque(control->task->stmt) || + std::any_of(control->children.begin(), control->children.end(), + [](const auto &child) { + return HasUnknownStateEffects(child.get()); + }); +} - uint16_t mask = 0; - for (SpecialRegister reg : - {SpecialRegister::kHf32Mode, SpecialRegister::kMmadDirection}) { - uint16_t bit = static_cast(reg); - int64_t first_stage = -1; - int64_t writer_stage = -1; - bool valid = true; - for (size_t i = 0; i < nodes.size(); ++i) { - const IRStructure *node = nodes[i]; - bool writes = (node->GetSpecialWriteMask() & bit) != 0; - bool reads = (node->GetSpecialReadMask() & bit) != 0; - if (!writes && !reads) - continue; - if (node->HasConditions() || - (node->IsTask() && - static_cast(node)->IsPerCoreTask())) { - valid = false; - break; - } - if (writes) { - if (!node->IsTask() || - !IsCubeConfiguration(static_cast(node)->stmt) || - (writer_stage >= 0 && - (stages[i] < writer_stage || stages[i] > first_stage + 1))) { - valid = false; - break; - } - if (first_stage < 0) - first_stage = stages[i]; - writer_stage = stages[i]; - } - // A loop that only reads this mode stays in its writer's region as one - // outer unit. Its inner stages and buffer versions do not change the - // mode. - if (reads && (writer_stage < 0 || stages[i] != writer_stage)) { - valid = false; - break; - } +// Bind each read to its reaching definition in logical stage order, then +// protect that generation until the next physical overwrite. A single state +// has no ring of versions: each value must be consumed within one physical II. +// Unlike all-pairs logical WAW edges, these edges allow independent generations +// to interleave across any number of stages without changing their readers. +bool AddSingleStateDependencies(const std::vector &accesses, + const std::vector &stages, + const std::vector &latencies, + std::vector *deps) { + std::vector writers, logical; + for (size_t i = 0; i < accesses.size(); ++i) { + if (accesses[i].writes) { + // A conditional kill may leave an older generation live. Retain the + // conservative logical dependencies until that reaching definition is + // known, rather than assuming that the skipped writer reset the state. + if (!accesses[i].must_write) + return false; + writers.push_back(i); } - if (valid && first_stage >= 0 && writer_stage == first_stage + 1) - mask |= bit; + logical.push_back(i); } - return mask; + if (writers.empty()) + return false; + std::stable_sort(logical.begin(), logical.end(), + [&](size_t a, size_t b) { return stages[a] < stages[b]; }); + auto add = [&](size_t u, size_t v, int64_t wrap, int64_t gap) { + ICHECK(wrap == 0 || wrap == 1) + << "Manual schedule constraints are infeasible: a single-state read " + "would cross a physical overwrite; check T.Stage values"; + // Reconstruction sorts equal phases by descending stage, then source + // position. Preserve a required edge when that tie order would reverse it. + if (wrap == 0 && + (stages[u] < stages[v] || (stages[u] == stages[v] && u > v))) + gap = std::max(gap, 1); + deps->emplace_back(u, v, wrap, gap); + }; + std::vector next_writer(accesses.size()); + for (size_t k = 0; k < writers.size(); ++k) { + size_t u = writers[k], v = writers[(k + 1) % writers.size()]; + next_writer[u] = v; + add(u, v, v <= u ? 1 : 0, latencies[u]); + } + // The first read may consume the previous logical iteration's last write. + size_t writer = writers.front(); + for (size_t i : logical) { + if (accesses[i].writes) + writer = i; + } + int64_t distance = 1; + for (size_t reader : logical) { + if (accesses[reader].reads) { + int64_t wrap = distance + stages[reader] - stages[writer]; + size_t overwrite = next_writer[writer]; + int64_t overwrite_wrap = overwrite <= writer ? 1 : 0; + add(writer, reader, wrap, latencies[writer]); + add(reader, overwrite, overwrite_wrap - wrap, 0); + } + if (accesses[reader].writes) { + writer = reader; + distance = 0; + } + } + return true; } // AutoSchedule target resource limits. @@ -911,7 +1061,7 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, std::vector> data_deps; // (i, j, distance, latency) std::vector> - register_phase_deps; // (i, j, physical wrap, latency) + state_phase_deps; // (i, j, physical wrap, latency) std::vector> resource_deps; std::vector> owner_exclusion_deps; @@ -1078,22 +1228,83 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, } } - // Config regions in adjacent manual stages share one physical register. - // Keep logical dependencies for all other special state, including fences. - uint16_t phase_mask = manual_schedule && enable_offset - ? ManualRegisterPhaseMask(nodes, manual_stages) - : 0; + uint16_t phase_mask = 0; + std::unordered_set phase_storages; + if (manual_schedule && enable_offset && + !std::any_of(nodes.begin(), nodes.end(), HasUnknownStateEffects)) { + auto add_state = [&](const SingleState &state) { + std::vector state_deps; + bool has_writer = false; + for (CoreMask core : kConcreteCores) { + std::vector accesses(n); + bool core_has_writer = false; + for (size_t i = 0; i < n; ++i) { + if (NodeCoreMask(nodes[i]) & core) + accesses[i] = SummarizeState(nodes[i], state, core); + core_has_writer |= accesses[i].writes; + } + if (!core_has_writer) + continue; + has_writer = true; + if (!AddSingleStateDependencies(accesses, manual_stages, latencies, + &state_deps)) + return false; + } + state_phase_deps.insert(state_phase_deps.end(), state_deps.begin(), + state_deps.end()); + return has_writer; + }; + for (uint16_t bit = 1; bit != 0; bit <<= 1) { + if ((kConfigurationRegisters & bit) && add_state({bit, std::nullopt})) + phase_mask |= bit; + } + std::unordered_map + scalar_storages; + std::unordered_set non_scalar_storages; + for (const IRStructure *node : nodes) { + auto regions = node->GetWriteRegions(); + auto reads = node->GetReadRegions(); + regions.insert(regions.end(), reads.begin(), reads.end()); + for (const auto ®ion : regions) { + const Buffer &buffer = region->buffer; + if (buffer.scope() == "local.var" && buffer->shape.size() == 1 && + is_one(buffer->shape[0])) { + auto [it, inserted] = scalar_storages.emplace(buffer->data, buffer); + if (!inserted && !it->second.same_as(buffer)) + non_scalar_storages.insert(buffer->data); + } else { + non_scalar_storages.insert(buffer->data); + } + } + } + LetDependencyClosure let_closure = CollectLetDependencyClosure_(nodes); + for (const auto &entry : scalar_storages) { + const Var &storage = entry.first; + // Explicit rings and aliased views keep buffer-region semantics: a + // differently typed or offset alias may write only part of a value. + if (non_scalar_storages.count(storage) || + manual_buffer_versions_.count(storage) || + buffer_version_overrides_.count(storage) || + ctrl->IsMultiBufferEligible(storage)) + continue; + // Cross-stage Bind copies re-read storage in their consumer's stage. + // Keep the existing snapshot constraints until those copies have their + // own explicit reaching-definition identity. + bool copied_read = false; + for (size_t u = 0; u < n; ++u) { + for (size_t l : let_closure.dependencies[u]) + copied_read |= manual_stages[u] != manual_stages[l] && + nodes[l]->ReadsStorage(storage); + } + if (!copied_read && add_state({0, storage})) + phase_storages.insert(storage); + } + } auto add_register_dep = [&](size_t u, size_t v, int wrap, int64_t latency, uint16_t registers) { + // Fences and state with unresolved reaching definitions remain logical. if (registers & ~phase_mask) data_deps.emplace_back(u, v, wrap, latency); - if (registers & phase_mask) { - // Reconstruction issues larger stages first at equal physical times. - // A forward cross-stage WAR must therefore be strictly ordered. - if (wrap == 0 && manual_stages[u] < manual_stages[v]) - latency = std::max(latency, 1); - register_phase_deps.emplace_back(u, v, wrap, latency); - } }; for (size_t i = 0; i < n; ++i) { uint16_t wmask = nodes[i]->GetSpecialWriteMask(); @@ -1132,6 +1343,9 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, node_idx[nodes[i]] = i; } for (const auto &dep : deps) { + if (dep.kind == DependencyKind::kData && dep.storage.defined() && + phase_storages.count(dep.storage.value())) + continue; size_t i = node_idx[dep.prod_node]; size_t j = node_idx[dep.cons_node]; int64_t latency = GetDependencyLatency(dep); @@ -1231,9 +1445,9 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, ffi::Array> tvm_pipe_order_deps; ffi::Array tvm_storage_sizes; ffi::Array tvm_manual_stages; - ffi::Array> tvm_register_phase_deps; - for (const auto &[u, v, wrap, latency] : register_phase_deps) - tvm_register_phase_deps.push_back({u, v, wrap, latency}); + ffi::Array> tvm_state_phase_deps; + for (const auto &[u, v, wrap, latency] : state_phase_deps) + tvm_state_phase_deps.push_back({u, v, wrap, latency}); for (auto val : latencies) { tvm_latencies.push_back(val); @@ -1307,7 +1521,7 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, tvm_data_deps, tvm_resource_deps, tvm_owner_exclusion_deps, tvm_pipe_order_deps, tvm_storage_sizes, tvm_memory_groups, tvm_stage_order_deps, enable_offset, tvm_manual_stages, - manual_schedule, tvm_register_phase_deps) + manual_schedule, tvm_state_phase_deps) .cast, ffi::Array, int64_t>>(); ffi::Array start_times = return_val.get<0>(); diff --git a/testing/ascend/auto_schedule/test_schedule_integration.py b/testing/ascend/auto_schedule/test_schedule_integration.py index 141a51872b..9459e0d91d 100644 --- a/testing/ascend/auto_schedule/test_schedule_integration.py +++ b/testing/ascend/auto_schedule/test_schedule_integration.py @@ -834,3 +834,47 @@ def main(A: T.Tensor((16,), "float32"), O: T.Tensor((16,), "float32")): a = torch.arange(16, dtype=torch.float32) result = tilelang.compile(main, target="ascend", out_idx=-1)(a.npu()) torch.testing.assert_close(result.cpu(), 2 * (a.cumsum(0) - a), rtol=0, atol=0) + + +@tilelang.testing.requires_ascend +@pytest.mark.parametrize("first_stage, second_stage", [(0, 3), (5, 2)]) +def test_manual_scalar_generations_and_exit_value(first_stage, second_stage): + import torch + + @T.prim_func + def main(A: T.Tensor((8,), "int32"), O: T.Tensor((8, 2), "int32"), Tail: T.Tensor((1,), "int32")): + with T.Kernel(1): + x = T.alloc_var("int32", init=-7) + for i in T.Pipelined(8, annotations={"enable_offset": True}): + with T.Stage(first_stage): + x = A[i] + 10 + O[i, 0] = x + with T.Stage(second_stage): + x = A[i] + 100 + O[i, 1] = x + Tail[0] = x + + a = torch.arange(8, dtype=torch.int32) + output, tail = tilelang.compile(main, target="ascend", out_idx=[1, 2])(a.npu()) + torch.testing.assert_close(output.cpu(), torch.stack((a + 10, a + 100), dim=1), rtol=0, atol=0) + final_offset = 100 if second_stage > first_stage else 10 + torch.testing.assert_close(tail.cpu(), a[-1:] + final_offset, rtol=0, atol=0) + + +@tilelang.testing.requires_ascend +def test_manual_scalar_reads_previous_physical_round(): + import torch + + @T.prim_func + def main(A: T.Tensor((8,), "int32"), O: T.Tensor((8,), "int32")): + with T.Kernel(1): + x = T.alloc_var("int32", init=-7) + for i in T.Pipelined(8, annotations={"enable_offset": True}): + with T.Stage(0): + x = A[i] + with T.Stage(1): + O[i] = x + + a = torch.arange(8, dtype=torch.int32) * 7 + 3 + output = tilelang.compile(main, target="ascend", out_idx=1)(a.npu()) + torch.testing.assert_close(output.cpu(), a, rtol=0, atol=0) diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 049c6ebd57..90f90bcde2 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -1,4 +1,4 @@ -"""Special-register dependencies participate in scheduling and core placement.""" +"""Single-state dependencies participate in scheduling and core placement.""" from unittest.mock import patch @@ -200,8 +200,8 @@ def _manual_cube_mode_loop(stages, mode, *, guarded_writer=False, mixed_modes=Fa tirx.Evaluate(T.gemm(a, b, outputs[1], transpose_B=True, clear_accum=True)), ] if mixed_modes: - # HF32 has only its stage-4 writer, but its second reader is stage 5. - # MMAD's two complete regions can use phases; HF32 must stay logical. + # HF32's stage-5 reader still needs the previous physical round's + # setter; MMAD's same-round region order must not hide that dependency. operations[0] = seq(tirx.Evaluate(T.set_hf32_mode(None)), operations[0]) units = [] for index, (op, stage) in enumerate(zip(operations, stages)): @@ -243,9 +243,8 @@ def _manual_cube_mode_loop(stages, mode, *, guarded_writer=False, mixed_modes=Fa @pytest.mark.parametrize("mode", ["hf32", "mmad-direction"]) -@pytest.mark.parametrize("base_stage", [0, 4, 12]) -def test_manual_cube_mode_regions_use_physical_order(mode, base_stage): - stages = [base_stage, base_stage, base_stage + 1, base_stage + 1] +@pytest.mark.parametrize("stages", [[0, 0, 1, 1], [4, 4, 7, 7], [12, 12, 20, 20], [5, 5, 4, 4]]) +def test_manual_cube_mode_regions_use_physical_order(mode, stages): before, operations = _manual_cube_mode_loop(stages, mode) # Use the real solver: fixed logical 0/1 WAW distances make this infeasible. after = transform.AutoSchedule()(before) @@ -263,7 +262,7 @@ def test_manual_cube_mode_regions_keep_nested_consumers(mode): original = z3_scheduler.z3_schedule_loop_python def capture(*args, **kwargs): - phases.append(kwargs["register_phase_deps"]) + phases.append(kwargs["state_phase_deps"]) return original(*args, **kwargs) with patch.object(z3_scheduler, "z3_schedule_loop_python", side_effect=capture): @@ -280,7 +279,7 @@ def capture(*args, **kwargs): assert (1, 2, 0, 1) in phases[-1] -@pytest.mark.parametrize("effect", ["mode-writer", "opaque-header", "opaque-guard"]) +@pytest.mark.parametrize("effect", ["opaque-header", "opaque-guard"]) def test_manual_cube_mode_nested_consumers_reject_mode_changes(effect): before, _ = _manual_cube_mode_loop([4, 4, 5, 5], "mmad-direction", nested_readers=True, nested_effect=effect) with pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"): @@ -290,10 +289,8 @@ def test_manual_cube_mode_nested_consumers_reject_mode_changes(effect): @pytest.mark.parametrize( "stages, guarded_writer", [ - pytest.param([5, 5, 4, 4], False, id="exit-mode"), pytest.param([4, 5, 5, 5], False, id="reader-crosses-stage"), pytest.param([4, 4, 5, 5], True, id="conditional-writer"), - pytest.param([4, 4, 6, 6], False, id="nonadjacent-stages"), ], ) def test_manual_cube_mode_unsafe_regions_keep_logical_dependencies(stages, guarded_writer): @@ -304,26 +301,10 @@ def test_manual_cube_mode_unsafe_regions_keep_logical_dependencies(stages, guard def test_manual_cube_modes_preserve_other_register_dependencies(): before, _ = _manual_cube_mode_loop([4, 4, 5, 5], "mmad-direction", mixed_modes=True) - inputs = [] - original = z3_scheduler.z3_schedule_loop_python - - def capture(*args, **kwargs): - inputs.append((args[4], kwargs["register_phase_deps"])) - return original(*args, **kwargs) - - with ( - patch.object(z3_scheduler, "z3_schedule_loop_python", side_effect=capture), - pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"), - ): + # HF32's stage-5 reader needs the previous physical round's value, but + # MMAD's region order puts the next HF32 setter ahead of that reader. + with pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"): transform.AutoSchedule()(before) - assert len(inputs) == 1 - logical, physical = inputs[0] - # The same writer/reader pair must retain HF32's logical dependency as - # well as MMAD's physical dependency; classifying whole pairs loses one. - assert (0, 3, 0, 1) in logical - assert (0, 3, 0, 1) in physical - # At equal phases the stage-5 setter would precede the stage-4 GEMM. - assert (1, 2, 0, 1) in physical @pytest.mark.parametrize("nested", ["unrelated", "mode-writer", "opaque"]) @@ -334,9 +315,116 @@ def test_manual_cube_mode_regions_inspect_nested_effects(nested): "opaque": tirx.call_extern("int32", "custom_stateful_call"), }[nested] before, _ = _manual_cube_mode_loop([4, 4, 5, 5], "mmad-direction", nested=operation) - if nested == "unrelated": + if nested != "opaque": after = transform.AutoSchedule()(before) - assert len(calls(after, "tl.ascend_set_mmad_direction")) == 2 + assert len(calls(after, "tl.ascend_set_mmad_direction")) == (3 if nested == "mode-writer" else 2) else: with pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"): transform.AutoSchedule()(before) + + +@pytest.mark.parametrize("mode", ["hf32", "mmad-direction"]) +def test_manual_state_regions_include_nested_writers(mode): + before, _ = _manual_cube_mode_loop([0, 0, 3, 3], mode, nested_readers=True, nested_effect="mode-writer") + after = transform.AutoSchedule()(before) + children = after["main"].body.block.body.body.body.seq + assert [int(child.node["stage"]) for child in children] == [0, 0, 3, 3] + assert len(calls(after, "tl.tileop.gemm")) == 2 + + +@pytest.mark.parametrize("mode", ["padding", "atomic"]) +def test_manual_configuration_generations_use_one_physical_state(mode): + a = tirx.decl_buffer((16,), "float32", name="A") + ub = [tirx.decl_buffer((16,), "float32", name=f"ub{i}", scope="shared.dyn") for i in range(2)] + outputs = [tirx.decl_buffer((16,), "float32", name=f"O{i}") for i in range(2)] + i = tirx.Var("i", "int32") + setters = ( + [T.ascend_set_copy_pad_value(value, dtype="float32") for value in (1.0, 2.0)] + if mode == "padding" + else [T.set_atomic("add", "float32"), T.set_atomic_none()] + ) + readers = ( + [copy(a, buf, annotations={"data_select": tirx.IntImm("int32", 1)}) for buf in ub] + if mode == "padding" + else [copy(buf, out) for buf, out in zip(ub, outputs)] + ) + operations = [tirx.Evaluate(setters[0]), readers[0], tirx.Evaluate(setters[1]), readers[1]] + stages = [0, 0, 3, 3] + body = seq(*(unit(op, core=1, stage=stage, cost=(1, 1)) for op, stage in zip(operations, stages))) + loop = tirx.For(i, 0, 4, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) + before = kernel(unit(loop, core=None, stage=-1), buffers=ub, params=[a, *outputs]) + after = transform.AutoSchedule()(before) + children = after["main"].body.block.body.body.body.seq + assert [int(child.node["stage"]) for child in children] == stages + for child, operation in zip(children, operations): + tvm.ir.assert_structural_equal(child.body.body, operation) + + +def _scalar_state_loop(stages, *, accumulate=False, single_writer=False): + x = tirx.decl_buffer((1,), "int32", name="x", scope="local.var") + out = tirx.decl_buffer((8, 2), "int32", name="O") + i = tirx.Var("i", "int32") + operations = [ + tirx.BufferStore(x, (x[0] if accumulate else 0) + i + 10, [0]), + tirx.BufferStore(out, x[0], [i, 0]), + tirx.BufferStore(x, (x[0] if accumulate else 0) + i + 100, [0]), + tirx.BufferStore(out, x[0], [i, 1]), + ] + if single_writer: + operations = operations[:2] + body = seq(*(unit(op, core=1, stage=stage, cost=(1, 1)) for op, stage in zip(operations, stages))) + loop = tirx.For(i, 0, 8, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) + return kernel(unit(loop, core=None, stage=-1), buffers=[x], params=[out]), operations + + +@pytest.mark.parametrize("stages", [[0, 0, 3, 3], [4, 4, 9, 9], [5, 5, 2, 2]]) +def test_manual_scalar_generations_preserve_each_read(stages): + before, operations = _scalar_state_loop(stages) + after = transform.AutoSchedule()(before) + children = after["main"].body.block.body.body.body.seq + assert [int(child.node["stage"]) for child in children] == stages + for child, operation in zip(children, operations): + tvm.ir.assert_structural_equal(child.body.body, operation) + + +def test_manual_scalar_read_precedes_next_iteration_overwrite(): + before, operations = _scalar_state_loop([0, 1], single_writer=True) + after = transform.AutoSchedule()(before) + children = after["main"].body.block.body.body.body.seq + # At physical round q, read x from q-1 before assigning x for q. A + # write-before-read modulo schedule compiles but returns the wrong value. + assert [int(child.node["stage"]) for child in children] == [1, 0] + for child, operation in zip(children, reversed(operations)): + tvm.ir.assert_structural_equal(child.body.body, operation) + + +@pytest.mark.parametrize("accumulate, stages, single_writer", [(True, [0, 0, 3, 3], False), (False, [0, 2], True)]) +def test_manual_scalar_state_rejects_overwritten_generations(accumulate, stages, single_writer): + before, _ = _scalar_state_loop(stages, accumulate=accumulate, single_writer=single_writer) + with pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"): + transform.AutoSchedule()(before) + + +def test_manual_scalar_recurrence_keeps_its_update_chain(): + before, operations = _scalar_state_loop([4, 4, 4, 4], accumulate=True) + after = transform.AutoSchedule()(before) + children = after["main"].body.block.body.body.body.seq + for child, operation in zip(children, operations): + tvm.ir.assert_structural_equal(child.body.body, operation) + + +def test_manual_scalar_task_kills_its_incoming_value(): + before, operations = _scalar_state_loop([0, 0, 3, 3]) + loop = before["main"].body.block.body.body + grouped = seq( + unit(seq(*operations[:2]), core=1, stage=0, cost=(2, 2)), + unit(seq(*operations[2:]), core=1, stage=3, cost=(2, 2)), + ) + loop = tirx.For(loop.loop_var, loop.min, loop.extent, loop.kind, grouped, annotations=loop.annotations) + block = before["main"].body.block + before = kernel(unit(loop, core=None, stage=-1), buffers=list(block.alloc_buffers), params=list(before["main"].buffer_map.values())) + after = transform.AutoSchedule()(before) + children = after["main"].body.block.body.body.body.seq + assert [int(child.node["stage"]) for child in children] == [0, 3] + for child, expected in zip(children, [seq(*operations[:2]), seq(*operations[2:])]): + tvm.ir.assert_structural_equal(child.body.body, expected) diff --git a/tilelang/ascend/language/schedule_hint.py b/tilelang/ascend/language/schedule_hint.py index a1a95e6df8..85da72839e 100644 --- a/tilelang/ascend/language/schedule_hint.py +++ b/tilelang/ascend/language/schedule_hint.py @@ -145,11 +145,14 @@ def Stage(stage: int): ``T.Pipelined(..., num_stages=...)`` controls the automatic buffer-version search bound; it does not limit the values accepted by ``T.Stage``. - Unconditional HF32/MMAD configuration regions in adjacent stages may share - the physical register in source issue order. Each consumer must use its - preceding setter's stage, and writer stages must not decrease. This keeps - both consumer modes and the loop's outgoing mode unchanged. Nested or - conditional configuration regions retain logical iteration dependencies. + Single-element ``local.var`` buffers and hardware configuration registers + share one physical state across stages. Reads retain their reaching write + in logical stage order (source order within a stage), and must precede the + next physical overwrite, at most one pipeline round away. Independent + generations can therefore use nonadjacent stages. Read-modify-write updates + retain their value dependencies. Control and pipe fences remain logical; + unresolved conditional writes, opaque effects, and cross-stage scalar + snapshots retain conservative dependencies. Example ------- diff --git a/tilelang/ascend/transform/z3_scheduler.py b/tilelang/ascend/transform/z3_scheduler.py index 19c2675bdf..73d0364b3d 100644 --- a/tilelang/ascend/transform/z3_scheduler.py +++ b/tilelang/ascend/transform/z3_scheduler.py @@ -283,7 +283,7 @@ def z3_schedule_loop_python( manual_schedule: bool = False, seed: int | None = 42, verbose: bool = False, - register_phase_deps: list[tuple[int, int, int, int]] | None = None, + state_phase_deps: list[tuple[int, int, int, int]] | None = None, ) -> tuple[list[int], list[int], int]: """Schedule a loop body with modulo, stage, and storage constraints. @@ -341,10 +341,11 @@ def z3_schedule_loop_python( verbose : bool Print solver inputs, search progress, and the selected schedule. - register_phase_deps : list[tuple[int, int, int, int]] | None - Manual configuration ordering (u, v, physical_wrap, gap), requiring - r[v] + physical_wrap * II >= r[u] + gap. These do not carry buffer - versions or logical iteration distances. + state_phase_deps : list[tuple[int, int, int, int]] | None + Single-state value/overwrite ordering (u, v, physical_wrap, gap), + requiring r[v] + physical_wrap * II >= r[u] + gap. Reaching writes + have already been matched and logical distances converted to a + physical wrap of zero or one; these states have no buffer versions. Returns ------- @@ -375,14 +376,14 @@ def z3_schedule_loop_python( if not enable_offset and any(stage != 0 for stage in manual_stages): raise ValueError("Non-zero manual stages require enable_offset=True") - if register_phase_deps is None: - register_phase_deps = [] - if register_phase_deps: + if state_phase_deps is None: + state_phase_deps = [] + if state_phase_deps: if not manual_schedule or not enable_offset: - raise ValueError("Register phase dependencies require a manual offset schedule") - for u, v, wrap, gap in register_phase_deps: + raise ValueError("State phase dependencies require a manual offset schedule") + for u, v, wrap, gap in state_phase_deps: if not (0 <= u < n and 0 <= v < n and wrap in (0, 1) and gap >= 0): - raise ValueError(f"Invalid register phase dependency: {(u, v, wrap, gap)}") + raise ValueError(f"Invalid state phase dependency: {(u, v, wrap, gap)}") if verbose: print(f"[Python Z3 Loop] Starting scheduling for {n} tasks") @@ -390,7 +391,7 @@ def z3_schedule_loop_python( print(f"[Python Z3 Loop] IIs: {iis}") print(f"[Python Z3 Loop] Resource flags: {resource_flags}") print(f"[Python Z3 Loop] Data dependencies with distances: {data_deps}") - print(f"[Python Z3 Loop] Register phase dependencies: {register_phase_deps}") + print(f"[Python Z3 Loop] State phase dependencies: {state_phase_deps}") print(f"[Python Z3 Loop] Resource dependencies: {resource_deps}") print(f"[Python Z3 Loop] Owner exclusions: {owner_exclusion_deps}") print(f"[Python Z3 Loop] Stage-order dependencies: {stage_order_deps}") @@ -455,7 +456,7 @@ def z3_schedule_loop_python( if idxs: solver.add(z3.Sum([buffer_vars[i] * buffer_sizes[i] for i in idxs]) <= cap) - for u, v, wrap, gap in register_phase_deps: + for u, v, wrap, gap in state_phase_deps: solver.add(r_vars[v] + wrap * ii_mid >= r_vars[u] + gap) # Add data dependency constraints with distance @@ -633,7 +634,7 @@ def z3_schedule_loop_python( "iis": iis, "resource_flags": resource_flags, "data_dependencies": data_deps, - "register_phase_dependencies": register_phase_deps, + "state_phase_dependencies": state_phase_deps, "resource_dependencies": resource_deps, "owner_exclusions": owner_exclusion_deps, "manual_schedule": manual_schedule, @@ -817,7 +818,7 @@ def z3_schedule_loop_ffi( enable_offset=False, manual_stages=None, manual_schedule=False, - register_phase_deps=None, + state_phase_deps=None, ): """FFI wrapper for z3_schedule_loop_python. @@ -839,7 +840,7 @@ def z3_schedule_loop_ffi( if memory_groups is not None: for i in range(len(memory_groups)): memory_groups_list.append(list(memory_groups[i])) - register_phase_deps_list = _convert_int_tuples(register_phase_deps, 4) + state_phase_deps_list = _convert_int_tuples(state_phase_deps, 4) data_deps_list = _convert_int_tuples(data_deps, 4) resource_deps_list = _convert_int_tuples(resource_deps, 2) stage_order_deps_list = _convert_int_tuples(stage_order_deps, 2) @@ -868,7 +869,7 @@ def z3_schedule_loop_ffi( enable_offset=bool(enable_offset), manual_stages=manual_stages_list, manual_schedule=bool(manual_schedule), - register_phase_deps=register_phase_deps_list, + state_phase_deps=state_phase_deps_list, seed=42, ) From 8230a83d2744fcdf6e20233ab58a3ab73dd9ac3b Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 16:24:47 +0800 Subject: [PATCH 07/10] [Ascend] Reuse memory access summaries for scalar state --- src/ascend/transform/auto_schedule.cc | 102 ++++-------------- .../auto_schedule/test_special_registers.py | 15 ++- tilelang/ascend/language/schedule_hint.py | 10 +- 3 files changed, 38 insertions(+), 89 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 8109b0bd27..b98f8bc076 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -57,6 +57,7 @@ #include "./auto_schedule/dependency_analysis.h" #include "./auto_schedule/ir_structure.h" #include "./auto_schedule/kernel_rewriter.h" +#include "./auto_schedule/memory_detector.h" #include "./auto_schedule/multi_buffer.h" #include "./auto_schedule/scheduled_tir.h" #include "./auto_schedule/task_analysis.h" @@ -94,14 +95,6 @@ struct SingleState { uint16_t reg{0}; Optional storage; - bool Reads(const IRStructure *node) const { - return storage.defined() ? node->ReadsStorage(storage.value()) - : (node->GetSpecialReadMask() & reg) != 0; - } - bool Writes(const IRStructure *node) const { - return storage.defined() ? node->WritesStorage(storage.value()) - : (node->GetSpecialWriteMask() & reg) != 0; - } bool MustWrite(const Stmt &stmt) const { if (const auto *attr = stmt.as()) return MustWrite(attr->body); @@ -126,80 +119,13 @@ struct SingleState { }; struct StateAccess { - // Reads of the incoming value, before a guaranteed write in this unit. + // Task reads conservatively consume the incoming value, even if the task + // also writes it. Control nodes compose their children in logical order. bool reads{false}; bool writes{false}; bool must_write{false}; }; -// Task-local scalar assignments can kill earlier values before a later read. -// Summarize those effects in execution order, including branch intersections. -class ScalarStateAccessAnalyzer : public StmtExprVisitor { -public: - explicit ScalarStateAccessAnalyzer(Var storage) - : storage_(std::move(storage)) {} - - StateAccess Analyze(const Stmt &stmt) { - VisitStmt(stmt); - return access_; - } - -private: - Var storage_; - StateAccess access_; - - void VisitExpr_(const BufferLoadNode *op) final { - if (op->buffer->data.same_as(storage_)) - access_.reads |= !access_.must_write; - StmtExprVisitor::VisitExpr_(op); - } - void VisitExpr_(const VarNode *op) final { - // An address passed to an intrinsic can read the incoming scalar too. - if (GetRef(op).same_as(storage_)) - access_.reads |= !access_.must_write; - } - void VisitStmt_(const BufferStoreNode *op) final { - VisitExpr(op->value); - for (const PrimExpr &index : op->indices) - VisitExpr(index); - if (op->predicate.defined()) - VisitExpr(op->predicate.value()); - if (op->buffer->data.same_as(storage_)) { - access_.writes = true; - access_.must_write |= - !op->predicate.defined() || is_one(op->predicate.value()); - } - } - void VisitStmt_(const IfThenElseNode *op) final { - VisitExpr(op->condition); - bool incoming = access_.must_write; - VisitStmt(op->then_case); - bool then_written = access_.must_write; - access_.must_write = incoming; - if (op->else_case.defined()) - VisitStmt(op->else_case.value()); - access_.must_write &= then_written; - } - void VisitStmt_(const ForNode *op) final { - VisitExpr(op->min); - VisitExpr(op->extent); - bool incoming = access_.must_write; - VisitStmt(op->body); - const auto *extent = op->extent.as(); - access_.must_write = - incoming || (extent && extent->value > 0 && access_.must_write); - } - void VisitStmt_(const SBlockRealizeNode *op) final { - for (const PrimExpr &value : op->iter_values) - VisitExpr(value); - VisitExpr(op->predicate); - bool incoming = access_.must_write; - VisitStmt(op->block); - access_.must_write = - incoming || (is_one(op->predicate) && access_.must_write); - } -}; - CoreMask NodeCoreMask(const IRStructure *node); StateAccess SummarizeState(const IRStructure *node, const SingleState &state, @@ -210,11 +136,25 @@ StateAccess SummarizeState(const IRStructure *node, const SingleState &state, if (node->IsTask()) { const auto *task = static_cast(node); if (state.storage.defined()) { - result = - ScalarStateAccessAnalyzer(state.storage.value()).Analyze(task->stmt); - result.writes |= state.Writes(node); + MemoryAccessDetector detector( + MemoryAccessDetector::AccessMode::kPhysical); + detector.Analyze(task->stmt); + auto touches = [&](const std::vector ®ions) { + return std::any_of( + regions.begin(), regions.end(), [&](const BufferRegion ®ion) { + return region->buffer->data.same_as(state.storage.value()); + }); + }; + result = {touches(detector.GetReadRegions()), + touches(detector.GetWriteRegions()), + touches(detector.GetMustWriteRegions())}; + for (const Var &storage : detector.GetOpaqueAccessVars()) { + if (storage.same_as(state.storage.value())) + result = {true, true, false}; + } } else { - result = {state.Reads(node), state.Writes(node), + result = {(node->GetSpecialReadMask() & state.reg) != 0, + (node->GetSpecialWriteMask() & state.reg) != 0, state.MustWrite(task->stmt)}; } } else { diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 90f90bcde2..81bf80db98 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -413,18 +413,25 @@ def test_manual_scalar_recurrence_keeps_its_update_chain(): tvm.ir.assert_structural_equal(child.body.body, operation) -def test_manual_scalar_task_kills_its_incoming_value(): - before, operations = _scalar_state_loop([0, 0, 3, 3]) +@pytest.mark.parametrize("second_stage", [0, 3]) +def test_manual_scalar_task_keeps_conservative_read_dependencies(second_stage): + before, operations = _scalar_state_loop([0, 0, second_stage, second_stage]) loop = before["main"].body.block.body.body grouped = seq( unit(seq(*operations[:2]), core=1, stage=0, cost=(2, 2)), - unit(seq(*operations[2:]), core=1, stage=3, cost=(2, 2)), + unit(seq(*operations[2:]), core=1, stage=second_stage, cost=(2, 2)), ) loop = tirx.For(loop.loop_var, loop.min, loop.extent, loop.kind, grouped, annotations=loop.annotations) block = before["main"].body.block before = kernel(unit(loop, core=None, stage=-1), buffers=list(block.alloc_buffers), params=list(before["main"].buffer_map.values())) + # Whole-task access summaries retain each read even after a local write. + # The resulting update chain cannot span nonadjacent physical rounds. + if second_stage == 3: + with pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"): + transform.AutoSchedule()(before) + return after = transform.AutoSchedule()(before) children = after["main"].body.block.body.body.body.seq - assert [int(child.node["stage"]) for child in children] == [0, 3] + assert [int(child.node["stage"]) for child in children] == [0, second_stage] for child, expected in zip(children, [seq(*operations[:2]), seq(*operations[2:])]): tvm.ir.assert_structural_equal(child.body.body, expected) diff --git a/tilelang/ascend/language/schedule_hint.py b/tilelang/ascend/language/schedule_hint.py index 85da72839e..27015da078 100644 --- a/tilelang/ascend/language/schedule_hint.py +++ b/tilelang/ascend/language/schedule_hint.py @@ -149,10 +149,12 @@ def Stage(stage: int): share one physical state across stages. Reads retain their reaching write in logical stage order (source order within a stage), and must precede the next physical overwrite, at most one pipeline round away. Independent - generations can therefore use nonadjacent stages. Read-modify-write updates - retain their value dependencies. Control and pipe fences remain logical; - unresolved conditional writes, opaque effects, and cross-stage scalar - snapshots retain conservative dependencies. + generations can therefore use nonadjacent stages. Reads and writes within + one task are summarized without ordering, so even a read after a local + write conservatively retains its incoming-value dependency. + Control and pipe fences remain logical; unresolved conditional writes, + opaque effects, and cross-stage scalar snapshots retain conservative + dependencies. Example ------- From e9b1e1223a10eb357e653b9fe9e2eaad4cddd43c Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Thu, 22 Oct 2026 14:34:20 +0800 Subject: [PATCH 08/10] [Ascend] Limit pipeline phase dependencies to configuration registers --- src/ascend/transform/auto_schedule.cc | 180 ++++++------------ .../test_schedule_integration.py | 44 ----- .../auto_schedule/test_special_registers.py | 78 +------- tilelang/ascend/language/schedule_hint.py | 18 +- tilelang/ascend/transform/z3_scheduler.py | 32 ++-- 5 files changed, 89 insertions(+), 263 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index b98f8bc076..868a489a7e 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -57,7 +57,6 @@ #include "./auto_schedule/dependency_analysis.h" #include "./auto_schedule/ir_structure.h" #include "./auto_schedule/kernel_rewriter.h" -#include "./auto_schedule/memory_detector.h" #include "./auto_schedule/multi_buffer.h" #include "./auto_schedule/scheduled_tir.h" #include "./auto_schedule/task_analysis.h" @@ -91,36 +90,33 @@ constexpr uint16_t kConfigurationRegisters = using PhaseDependency = std::tuple; -struct SingleState { - uint16_t reg{0}; - Optional storage; - - bool MustWrite(const Stmt &stmt) const { - if (const auto *attr = stmt.as()) - return MustWrite(attr->body); - if (const auto *seq = stmt.as()) - return std::any_of(seq->seq.begin(), seq->seq.end(), - [&](const Stmt &child) { return MustWrite(child); }); - if (const auto *branch = stmt.as()) - return branch->else_case.defined() && MustWrite(branch->then_case) && - MustWrite(branch->else_case.value()); - const auto *evaluate = stmt.as(); - const auto *call = evaluate ? evaluate->value.as() : nullptr; - if (!call || !(AnalyzeSpecialRegisterWrites(stmt) & reg)) +bool MustSetRegister(const Stmt &stmt, uint16_t reg) { + if (const auto *attr = stmt.as()) + return MustSetRegister(attr->body, reg); + if (const auto *seq = stmt.as()) + return std::any_of( + seq->seq.begin(), seq->seq.end(), + [&](const Stmt &child) { return MustSetRegister(child, reg); }); + if (const auto *branch = stmt.as()) + return branch->else_case.defined() && + MustSetRegister(branch->then_case, reg) && + MustSetRegister(branch->else_case.value(), reg); + const auto *evaluate = stmt.as(); + const auto *call = evaluate ? evaluate->value.as() : nullptr; + if (!call || !(AnalyzeSpecialRegisterWrites(stmt) & reg)) + return false; + // Only an executed setter kills the old value. A setter nested in an + // argument (for example if_then_else) need not execute. + for (const PrimExpr &arg : call->args) { + if (AnalyzeSpecialRegisterWrites(Evaluate(arg)) & reg) return false; - // A write nested in an argument (for example if_then_else) need not - // execute. Only the call's own configuration effect is a definite kill. - for (const PrimExpr &arg : call->args) { - if (AnalyzeSpecialRegisterWrites(Evaluate(arg)) & reg) - return false; - } - return true; } -}; + return true; +} -struct StateAccess { - // Task reads conservatively consume the incoming value, even if the task - // also writes it. Control nodes compose their children in logical order. +struct RegisterAccess { + // Setters overwrite the register. Nested control nodes can also consume + // its incoming value before reaching a setter. bool reads{false}; bool writes{false}; bool must_write{false}; @@ -128,38 +124,19 @@ struct StateAccess { CoreMask NodeCoreMask(const IRStructure *node); -StateAccess SummarizeState(const IRStructure *node, const SingleState &state, - CoreMask core) { - StateAccess result; +RegisterAccess SummarizeRegister(const IRStructure *node, uint16_t reg, + CoreMask core) { + RegisterAccess result; if (!(NodeCoreMask(node) & core)) return result; if (node->IsTask()) { const auto *task = static_cast(node); - if (state.storage.defined()) { - MemoryAccessDetector detector( - MemoryAccessDetector::AccessMode::kPhysical); - detector.Analyze(task->stmt); - auto touches = [&](const std::vector ®ions) { - return std::any_of( - regions.begin(), regions.end(), [&](const BufferRegion ®ion) { - return region->buffer->data.same_as(state.storage.value()); - }); - }; - result = {touches(detector.GetReadRegions()), - touches(detector.GetWriteRegions()), - touches(detector.GetMustWriteRegions())}; - for (const Var &storage : detector.GetOpaqueAccessVars()) { - if (storage.same_as(state.storage.value())) - result = {true, true, false}; - } - } else { - result = {(node->GetSpecialReadMask() & state.reg) != 0, - (node->GetSpecialWriteMask() & state.reg) != 0, - state.MustWrite(task->stmt)}; - } + result = {(node->GetSpecialReadMask() & reg) != 0, + (node->GetSpecialWriteMask() & reg) != 0, + MustSetRegister(task->stmt, reg)}; } else { const auto *control = static_cast(node); - result = SummarizeState(control->task.get(), state, core); + result = SummarizeRegister(control->task.get(), reg, core); std::vector children; for (const auto &child : control->children) children.push_back(child.get()); @@ -170,7 +147,7 @@ StateAccess SummarizeState(const IRStructure *node, const SingleState &state, return a->GetStage() < b->GetStage(); }); for (const IRStructure *child : children) { - StateAccess access = SummarizeState(child, state, core); + RegisterAccess access = SummarizeRegister(child, reg, core); result.reads |= access.reads && !result.must_write; result.writes |= access.writes; result.must_write |= access.must_write; @@ -178,13 +155,11 @@ StateAccess SummarizeState(const IRStructure *node, const SingleState &state, const auto *extent = control->control->extent.as(); result.must_write &= extent && extent->value > 0; } - if (state.storage.defined()) - result.reads |= node->GuardsTouchStorage(state.storage.value()); result.must_write &= !node->HasConditions(); return result; } -bool HasUnknownStateEffects(const IRStructure *node) { +bool HasUnknownRegisterEffects(const IRStructure *node) { auto opaque = [](const ObjectRef &ir) { bool found = false; PostOrderVisit(ir, [&](const ObjectRef &object) { @@ -225,19 +200,19 @@ bool HasUnknownStateEffects(const IRStructure *node) { return opaque(control->task->stmt) || std::any_of(control->children.begin(), control->children.end(), [](const auto &child) { - return HasUnknownStateEffects(child.get()); + return HasUnknownRegisterEffects(child.get()); }); } // Bind each read to its reaching definition in logical stage order, then -// protect that generation until the next physical overwrite. A single state +// protect that generation until the next setter. A configuration register // has no ring of versions: each value must be consumed within one physical II. // Unlike all-pairs logical WAW edges, these edges allow independent generations // to interleave across any number of stages without changing their readers. -bool AddSingleStateDependencies(const std::vector &accesses, - const std::vector &stages, - const std::vector &latencies, - std::vector *deps) { +bool AddRegisterPhaseDependencies(const std::vector &accesses, + const std::vector &stages, + const std::vector &latencies, + std::vector *deps) { std::vector writers, logical; for (size_t i = 0; i < accesses.size(); ++i) { if (accesses[i].writes) { @@ -256,7 +231,7 @@ bool AddSingleStateDependencies(const std::vector &accesses, [&](size_t a, size_t b) { return stages[a] < stages[b]; }); auto add = [&](size_t u, size_t v, int64_t wrap, int64_t gap) { ICHECK(wrap == 0 || wrap == 1) - << "Manual schedule constraints are infeasible: a single-state read " + << "Manual schedule constraints are infeasible: a register read " "would cross a physical overwrite; check T.Stage values"; // Reconstruction sorts equal phases by descending stage, then source // position. Preserve a required edge when that tie order would reverse it. @@ -1001,7 +976,7 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, std::vector> data_deps; // (i, j, distance, latency) std::vector> - state_phase_deps; // (i, j, physical wrap, latency) + register_phase_deps; // (i, j, physical wrap, latency) std::vector> resource_deps; std::vector> owner_exclusion_deps; @@ -1169,76 +1144,34 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, } uint16_t phase_mask = 0; - std::unordered_set phase_storages; if (manual_schedule && enable_offset && - !std::any_of(nodes.begin(), nodes.end(), HasUnknownStateEffects)) { - auto add_state = [&](const SingleState &state) { - std::vector state_deps; + !std::any_of(nodes.begin(), nodes.end(), HasUnknownRegisterEffects)) { + auto add_register = [&](uint16_t reg) { + std::vector register_deps; bool has_writer = false; for (CoreMask core : kConcreteCores) { - std::vector accesses(n); + std::vector accesses(n); bool core_has_writer = false; for (size_t i = 0; i < n; ++i) { if (NodeCoreMask(nodes[i]) & core) - accesses[i] = SummarizeState(nodes[i], state, core); + accesses[i] = SummarizeRegister(nodes[i], reg, core); core_has_writer |= accesses[i].writes; } if (!core_has_writer) continue; has_writer = true; - if (!AddSingleStateDependencies(accesses, manual_stages, latencies, - &state_deps)) + if (!AddRegisterPhaseDependencies(accesses, manual_stages, latencies, + ®ister_deps)) return false; } - state_phase_deps.insert(state_phase_deps.end(), state_deps.begin(), - state_deps.end()); + register_phase_deps.insert(register_phase_deps.end(), + register_deps.begin(), register_deps.end()); return has_writer; }; for (uint16_t bit = 1; bit != 0; bit <<= 1) { - if ((kConfigurationRegisters & bit) && add_state({bit, std::nullopt})) + if ((kConfigurationRegisters & bit) && add_register(bit)) phase_mask |= bit; } - std::unordered_map - scalar_storages; - std::unordered_set non_scalar_storages; - for (const IRStructure *node : nodes) { - auto regions = node->GetWriteRegions(); - auto reads = node->GetReadRegions(); - regions.insert(regions.end(), reads.begin(), reads.end()); - for (const auto ®ion : regions) { - const Buffer &buffer = region->buffer; - if (buffer.scope() == "local.var" && buffer->shape.size() == 1 && - is_one(buffer->shape[0])) { - auto [it, inserted] = scalar_storages.emplace(buffer->data, buffer); - if (!inserted && !it->second.same_as(buffer)) - non_scalar_storages.insert(buffer->data); - } else { - non_scalar_storages.insert(buffer->data); - } - } - } - LetDependencyClosure let_closure = CollectLetDependencyClosure_(nodes); - for (const auto &entry : scalar_storages) { - const Var &storage = entry.first; - // Explicit rings and aliased views keep buffer-region semantics: a - // differently typed or offset alias may write only part of a value. - if (non_scalar_storages.count(storage) || - manual_buffer_versions_.count(storage) || - buffer_version_overrides_.count(storage) || - ctrl->IsMultiBufferEligible(storage)) - continue; - // Cross-stage Bind copies re-read storage in their consumer's stage. - // Keep the existing snapshot constraints until those copies have their - // own explicit reaching-definition identity. - bool copied_read = false; - for (size_t u = 0; u < n; ++u) { - for (size_t l : let_closure.dependencies[u]) - copied_read |= manual_stages[u] != manual_stages[l] && - nodes[l]->ReadsStorage(storage); - } - if (!copied_read && add_state({0, storage})) - phase_storages.insert(storage); - } } auto add_register_dep = [&](size_t u, size_t v, int wrap, int64_t latency, uint16_t registers) { @@ -1283,9 +1216,6 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, node_idx[nodes[i]] = i; } for (const auto &dep : deps) { - if (dep.kind == DependencyKind::kData && dep.storage.defined() && - phase_storages.count(dep.storage.value())) - continue; size_t i = node_idx[dep.prod_node]; size_t j = node_idx[dep.cons_node]; int64_t latency = GetDependencyLatency(dep); @@ -1385,9 +1315,9 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, ffi::Array> tvm_pipe_order_deps; ffi::Array tvm_storage_sizes; ffi::Array tvm_manual_stages; - ffi::Array> tvm_state_phase_deps; - for (const auto &[u, v, wrap, latency] : state_phase_deps) - tvm_state_phase_deps.push_back({u, v, wrap, latency}); + ffi::Array> tvm_register_phase_deps; + for (const auto &[u, v, wrap, latency] : register_phase_deps) + tvm_register_phase_deps.push_back({u, v, wrap, latency}); for (auto val : latencies) { tvm_latencies.push_back(val); @@ -1461,7 +1391,7 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, tvm_data_deps, tvm_resource_deps, tvm_owner_exclusion_deps, tvm_pipe_order_deps, tvm_storage_sizes, tvm_memory_groups, tvm_stage_order_deps, enable_offset, tvm_manual_stages, - manual_schedule, tvm_state_phase_deps) + manual_schedule, tvm_register_phase_deps) .cast, ffi::Array, int64_t>>(); ffi::Array start_times = return_val.get<0>(); diff --git a/testing/ascend/auto_schedule/test_schedule_integration.py b/testing/ascend/auto_schedule/test_schedule_integration.py index 9459e0d91d..141a51872b 100644 --- a/testing/ascend/auto_schedule/test_schedule_integration.py +++ b/testing/ascend/auto_schedule/test_schedule_integration.py @@ -834,47 +834,3 @@ def main(A: T.Tensor((16,), "float32"), O: T.Tensor((16,), "float32")): a = torch.arange(16, dtype=torch.float32) result = tilelang.compile(main, target="ascend", out_idx=-1)(a.npu()) torch.testing.assert_close(result.cpu(), 2 * (a.cumsum(0) - a), rtol=0, atol=0) - - -@tilelang.testing.requires_ascend -@pytest.mark.parametrize("first_stage, second_stage", [(0, 3), (5, 2)]) -def test_manual_scalar_generations_and_exit_value(first_stage, second_stage): - import torch - - @T.prim_func - def main(A: T.Tensor((8,), "int32"), O: T.Tensor((8, 2), "int32"), Tail: T.Tensor((1,), "int32")): - with T.Kernel(1): - x = T.alloc_var("int32", init=-7) - for i in T.Pipelined(8, annotations={"enable_offset": True}): - with T.Stage(first_stage): - x = A[i] + 10 - O[i, 0] = x - with T.Stage(second_stage): - x = A[i] + 100 - O[i, 1] = x - Tail[0] = x - - a = torch.arange(8, dtype=torch.int32) - output, tail = tilelang.compile(main, target="ascend", out_idx=[1, 2])(a.npu()) - torch.testing.assert_close(output.cpu(), torch.stack((a + 10, a + 100), dim=1), rtol=0, atol=0) - final_offset = 100 if second_stage > first_stage else 10 - torch.testing.assert_close(tail.cpu(), a[-1:] + final_offset, rtol=0, atol=0) - - -@tilelang.testing.requires_ascend -def test_manual_scalar_reads_previous_physical_round(): - import torch - - @T.prim_func - def main(A: T.Tensor((8,), "int32"), O: T.Tensor((8,), "int32")): - with T.Kernel(1): - x = T.alloc_var("int32", init=-7) - for i in T.Pipelined(8, annotations={"enable_offset": True}): - with T.Stage(0): - x = A[i] - with T.Stage(1): - O[i] = x - - a = torch.arange(8, dtype=torch.int32) * 7 + 3 - output = tilelang.compile(main, target="ascend", out_idx=1)(a.npu()) - torch.testing.assert_close(output.cpu(), a, rtol=0, atol=0) diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 81bf80db98..5faaad62b4 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -1,4 +1,4 @@ -"""Single-state dependencies participate in scheduling and core placement.""" +"""Special-register dependencies participate in scheduling and core placement.""" from unittest.mock import patch @@ -262,7 +262,7 @@ def test_manual_cube_mode_regions_keep_nested_consumers(mode): original = z3_scheduler.z3_schedule_loop_python def capture(*args, **kwargs): - phases.append(kwargs["state_phase_deps"]) + phases.append(kwargs["register_phase_deps"]) return original(*args, **kwargs) with patch.object(z3_scheduler, "z3_schedule_loop_python", side_effect=capture): @@ -324,7 +324,7 @@ def test_manual_cube_mode_regions_inspect_nested_effects(nested): @pytest.mark.parametrize("mode", ["hf32", "mmad-direction"]) -def test_manual_state_regions_include_nested_writers(mode): +def test_manual_configuration_regions_include_nested_writers(mode): before, _ = _manual_cube_mode_loop([0, 0, 3, 3], mode, nested_readers=True, nested_effect="mode-writer") after = transform.AutoSchedule()(before) children = after["main"].body.block.body.body.body.seq @@ -360,78 +360,20 @@ def test_manual_configuration_generations_use_one_physical_state(mode): tvm.ir.assert_structural_equal(child.body.body, operation) -def _scalar_state_loop(stages, *, accumulate=False, single_writer=False): +def test_manual_scalar_buffers_keep_logical_dependencies(): x = tirx.decl_buffer((1,), "int32", name="x", scope="local.var") out = tirx.decl_buffer((8, 2), "int32", name="O") i = tirx.Var("i", "int32") operations = [ - tirx.BufferStore(x, (x[0] if accumulate else 0) + i + 10, [0]), + tirx.BufferStore(x, i + 10, [0]), tirx.BufferStore(out, x[0], [i, 0]), - tirx.BufferStore(x, (x[0] if accumulate else 0) + i + 100, [0]), + tirx.BufferStore(x, i + 100, [0]), tirx.BufferStore(out, x[0], [i, 1]), ] - if single_writer: - operations = operations[:2] - body = seq(*(unit(op, core=1, stage=stage, cost=(1, 1)) for op, stage in zip(operations, stages))) + body = seq(*(unit(op, core=1, stage=stage, cost=(1, 1)) for op, stage in zip(operations, [0, 0, 3, 3]))) loop = tirx.For(i, 0, 8, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) - return kernel(unit(loop, core=None, stage=-1), buffers=[x], params=[out]), operations - - -@pytest.mark.parametrize("stages", [[0, 0, 3, 3], [4, 4, 9, 9], [5, 5, 2, 2]]) -def test_manual_scalar_generations_preserve_each_read(stages): - before, operations = _scalar_state_loop(stages) - after = transform.AutoSchedule()(before) - children = after["main"].body.block.body.body.body.seq - assert [int(child.node["stage"]) for child in children] == stages - for child, operation in zip(children, operations): - tvm.ir.assert_structural_equal(child.body.body, operation) - - -def test_manual_scalar_read_precedes_next_iteration_overwrite(): - before, operations = _scalar_state_loop([0, 1], single_writer=True) - after = transform.AutoSchedule()(before) - children = after["main"].body.block.body.body.body.seq - # At physical round q, read x from q-1 before assigning x for q. A - # write-before-read modulo schedule compiles but returns the wrong value. - assert [int(child.node["stage"]) for child in children] == [1, 0] - for child, operation in zip(children, reversed(operations)): - tvm.ir.assert_structural_equal(child.body.body, operation) - - -@pytest.mark.parametrize("accumulate, stages, single_writer", [(True, [0, 0, 3, 3], False), (False, [0, 2], True)]) -def test_manual_scalar_state_rejects_overwritten_generations(accumulate, stages, single_writer): - before, _ = _scalar_state_loop(stages, accumulate=accumulate, single_writer=single_writer) + before = kernel(unit(loop, core=None, stage=-1), buffers=[x], params=[out]) + # Register set/use groups may cross these stages, but ordinary scalar + # storage retains its logical WAW and reuse dependencies. with pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"): transform.AutoSchedule()(before) - - -def test_manual_scalar_recurrence_keeps_its_update_chain(): - before, operations = _scalar_state_loop([4, 4, 4, 4], accumulate=True) - after = transform.AutoSchedule()(before) - children = after["main"].body.block.body.body.body.seq - for child, operation in zip(children, operations): - tvm.ir.assert_structural_equal(child.body.body, operation) - - -@pytest.mark.parametrize("second_stage", [0, 3]) -def test_manual_scalar_task_keeps_conservative_read_dependencies(second_stage): - before, operations = _scalar_state_loop([0, 0, second_stage, second_stage]) - loop = before["main"].body.block.body.body - grouped = seq( - unit(seq(*operations[:2]), core=1, stage=0, cost=(2, 2)), - unit(seq(*operations[2:]), core=1, stage=second_stage, cost=(2, 2)), - ) - loop = tirx.For(loop.loop_var, loop.min, loop.extent, loop.kind, grouped, annotations=loop.annotations) - block = before["main"].body.block - before = kernel(unit(loop, core=None, stage=-1), buffers=list(block.alloc_buffers), params=list(before["main"].buffer_map.values())) - # Whole-task access summaries retain each read even after a local write. - # The resulting update chain cannot span nonadjacent physical rounds. - if second_stage == 3: - with pytest.raises(RuntimeError, match="Manual schedule constraints are infeasible"): - transform.AutoSchedule()(before) - return - after = transform.AutoSchedule()(before) - children = after["main"].body.block.body.body.body.seq - assert [int(child.node["stage"]) for child in children] == [0, second_stage] - for child, expected in zip(children, [seq(*operations[:2]), seq(*operations[2:])]): - tvm.ir.assert_structural_equal(child.body.body, expected) diff --git a/tilelang/ascend/language/schedule_hint.py b/tilelang/ascend/language/schedule_hint.py index 27015da078..e66cb0ced4 100644 --- a/tilelang/ascend/language/schedule_hint.py +++ b/tilelang/ascend/language/schedule_hint.py @@ -145,16 +145,14 @@ def Stage(stage: int): ``T.Pipelined(..., num_stages=...)`` controls the automatic buffer-version search bound; it does not limit the values accepted by ``T.Stage``. - Single-element ``local.var`` buffers and hardware configuration registers - share one physical state across stages. Reads retain their reaching write - in logical stage order (source order within a stage), and must precede the - next physical overwrite, at most one pipeline round away. Independent - generations can therefore use nonadjacent stages. Reads and writes within - one task are summarized without ordering, so even a read after a local - write conservatively retains its incoming-value dependency. - Control and pipe fences remain logical; unresolved conditional writes, - opaque effects, and cross-stage scalar snapshots retain conservative - dependencies. + Hardware configuration registers (HF32, MMAD direction, padding, and atomic + mode) each share one physical value across stages. Their setters overwrite + that value; readers retain the reaching setter in logical stage order and + must precede the next physical overwrite, at most one pipeline round away. + Independent set/use groups can therefore use nonadjacent stages. + Conditional setters and opaque effects retain conservative dependencies. + Scalar buffers, copied scalar snapshots, and control/pipe fences keep their + existing logical dependency rules. Example ------- diff --git a/tilelang/ascend/transform/z3_scheduler.py b/tilelang/ascend/transform/z3_scheduler.py index 73d0364b3d..51a2895121 100644 --- a/tilelang/ascend/transform/z3_scheduler.py +++ b/tilelang/ascend/transform/z3_scheduler.py @@ -283,7 +283,7 @@ def z3_schedule_loop_python( manual_schedule: bool = False, seed: int | None = 42, verbose: bool = False, - state_phase_deps: list[tuple[int, int, int, int]] | None = None, + register_phase_deps: list[tuple[int, int, int, int]] | None = None, ) -> tuple[list[int], list[int], int]: """Schedule a loop body with modulo, stage, and storage constraints. @@ -341,11 +341,11 @@ def z3_schedule_loop_python( verbose : bool Print solver inputs, search progress, and the selected schedule. - state_phase_deps : list[tuple[int, int, int, int]] | None - Single-state value/overwrite ordering (u, v, physical_wrap, gap), + register_phase_deps : list[tuple[int, int, int, int]] | None + Configuration-register ordering (u, v, physical_wrap, gap), requiring r[v] + physical_wrap * II >= r[u] + gap. Reaching writes have already been matched and logical distances converted to a - physical wrap of zero or one; these states have no buffer versions. + physical wrap of zero or one. Buffer dependencies stay in data_deps. Returns ------- @@ -376,14 +376,14 @@ def z3_schedule_loop_python( if not enable_offset and any(stage != 0 for stage in manual_stages): raise ValueError("Non-zero manual stages require enable_offset=True") - if state_phase_deps is None: - state_phase_deps = [] - if state_phase_deps: + if register_phase_deps is None: + register_phase_deps = [] + if register_phase_deps: if not manual_schedule or not enable_offset: - raise ValueError("State phase dependencies require a manual offset schedule") - for u, v, wrap, gap in state_phase_deps: + raise ValueError("Register phase dependencies require a manual offset schedule") + for u, v, wrap, gap in register_phase_deps: if not (0 <= u < n and 0 <= v < n and wrap in (0, 1) and gap >= 0): - raise ValueError(f"Invalid state phase dependency: {(u, v, wrap, gap)}") + raise ValueError(f"Invalid register phase dependency: {(u, v, wrap, gap)}") if verbose: print(f"[Python Z3 Loop] Starting scheduling for {n} tasks") @@ -391,7 +391,7 @@ def z3_schedule_loop_python( print(f"[Python Z3 Loop] IIs: {iis}") print(f"[Python Z3 Loop] Resource flags: {resource_flags}") print(f"[Python Z3 Loop] Data dependencies with distances: {data_deps}") - print(f"[Python Z3 Loop] State phase dependencies: {state_phase_deps}") + print(f"[Python Z3 Loop] Register phase dependencies: {register_phase_deps}") print(f"[Python Z3 Loop] Resource dependencies: {resource_deps}") print(f"[Python Z3 Loop] Owner exclusions: {owner_exclusion_deps}") print(f"[Python Z3 Loop] Stage-order dependencies: {stage_order_deps}") @@ -456,7 +456,7 @@ def z3_schedule_loop_python( if idxs: solver.add(z3.Sum([buffer_vars[i] * buffer_sizes[i] for i in idxs]) <= cap) - for u, v, wrap, gap in state_phase_deps: + for u, v, wrap, gap in register_phase_deps: solver.add(r_vars[v] + wrap * ii_mid >= r_vars[u] + gap) # Add data dependency constraints with distance @@ -634,7 +634,7 @@ def z3_schedule_loop_python( "iis": iis, "resource_flags": resource_flags, "data_dependencies": data_deps, - "state_phase_dependencies": state_phase_deps, + "register_phase_dependencies": register_phase_deps, "resource_dependencies": resource_deps, "owner_exclusions": owner_exclusion_deps, "manual_schedule": manual_schedule, @@ -818,7 +818,7 @@ def z3_schedule_loop_ffi( enable_offset=False, manual_stages=None, manual_schedule=False, - state_phase_deps=None, + register_phase_deps=None, ): """FFI wrapper for z3_schedule_loop_python. @@ -840,7 +840,7 @@ def z3_schedule_loop_ffi( if memory_groups is not None: for i in range(len(memory_groups)): memory_groups_list.append(list(memory_groups[i])) - state_phase_deps_list = _convert_int_tuples(state_phase_deps, 4) + register_phase_deps_list = _convert_int_tuples(register_phase_deps, 4) data_deps_list = _convert_int_tuples(data_deps, 4) resource_deps_list = _convert_int_tuples(resource_deps, 2) stage_order_deps_list = _convert_int_tuples(stage_order_deps, 2) @@ -869,7 +869,7 @@ def z3_schedule_loop_ffi( enable_offset=bool(enable_offset), manual_stages=manual_stages_list, manual_schedule=bool(manual_schedule), - state_phase_deps=state_phase_deps_list, + register_phase_deps=register_phase_deps_list, seed=42, ) From 67670fcb66a747a25a5388848b682a090c2ee397 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Thu, 22 Oct 2026 17:24:35 +0800 Subject: [PATCH 09/10] [Ascend] Summarize configuration register effects per core --- src/ascend/transform/auto_schedule.cc | 149 +++++++++++++------------- 1 file changed, 77 insertions(+), 72 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 868a489a7e..1535a9e0cc 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -90,53 +90,57 @@ constexpr uint16_t kConfigurationRegisters = using PhaseDependency = std::tuple; -bool MustSetRegister(const Stmt &stmt, uint16_t reg) { +uint16_t MustWriteConfigurationRegisters(const Stmt &stmt) { if (const auto *attr = stmt.as()) - return MustSetRegister(attr->body, reg); - if (const auto *seq = stmt.as()) - return std::any_of( - seq->seq.begin(), seq->seq.end(), - [&](const Stmt &child) { return MustSetRegister(child, reg); }); + return MustWriteConfigurationRegisters(attr->body); + if (const auto *seq = stmt.as()) { + uint16_t mask = 0; + for (const Stmt &child : seq->seq) + mask |= MustWriteConfigurationRegisters(child); + return mask; + } if (const auto *branch = stmt.as()) - return branch->else_case.defined() && - MustSetRegister(branch->then_case, reg) && - MustSetRegister(branch->else_case.value(), reg); + return branch->else_case.defined() + ? MustWriteConfigurationRegisters(branch->then_case) & + MustWriteConfigurationRegisters(branch->else_case.value()) + : 0; const auto *evaluate = stmt.as(); const auto *call = evaluate ? evaluate->value.as() : nullptr; - if (!call || !(AnalyzeSpecialRegisterWrites(stmt) & reg)) - return false; + if (!call) + return 0; + uint16_t mask = AnalyzeSpecialRegisterWrites(stmt) & kConfigurationRegisters; + if (mask == 0) + return 0; // Only an executed setter kills the old value. A setter nested in an // argument (for example if_then_else) need not execute. - for (const PrimExpr &arg : call->args) { - if (AnalyzeSpecialRegisterWrites(Evaluate(arg)) & reg) - return false; - } - return true; + for (const PrimExpr &arg : call->args) + mask &= ~AnalyzeSpecialRegisterWrites(Evaluate(arg)); + return mask; } -struct RegisterAccess { - // Setters overwrite the register. Nested control nodes can also consume - // its incoming value before reaching a setter. - bool reads{false}; - bool writes{false}; - bool must_write{false}; +// Effects on hardware configuration only. Ordinary scalar storage and variables +// use the existing buffer-region and definition-use dependencies. +struct ConfigurationEffects { + uint16_t read_mask{0}; // Incoming values that may be read. + uint16_t write_mask{0}; + uint16_t must_write_mask{0}; }; CoreMask NodeCoreMask(const IRStructure *node); -RegisterAccess SummarizeRegister(const IRStructure *node, uint16_t reg, - CoreMask core) { - RegisterAccess result; +ConfigurationEffects SummarizeConfigurationEffects(const IRStructure *node, + CoreMask core) { + ConfigurationEffects result; if (!(NodeCoreMask(node) & core)) return result; if (node->IsTask()) { const auto *task = static_cast(node); - result = {(node->GetSpecialReadMask() & reg) != 0, - (node->GetSpecialWriteMask() & reg) != 0, - MustSetRegister(task->stmt, reg)}; + result.read_mask = node->GetSpecialReadMask() & kConfigurationRegisters; + result.write_mask = node->GetSpecialWriteMask() & kConfigurationRegisters; + result.must_write_mask = MustWriteConfigurationRegisters(task->stmt); } else { const auto *control = static_cast(node); - result = SummarizeRegister(control->task.get(), reg, core); + result = SummarizeConfigurationEffects(control->task.get(), core); std::vector children; for (const auto &child : control->children) children.push_back(child.get()); @@ -147,19 +151,21 @@ RegisterAccess SummarizeRegister(const IRStructure *node, uint16_t reg, return a->GetStage() < b->GetStage(); }); for (const IRStructure *child : children) { - RegisterAccess access = SummarizeRegister(child, reg, core); - result.reads |= access.reads && !result.must_write; - result.writes |= access.writes; - result.must_write |= access.must_write; + ConfigurationEffects effects = SummarizeConfigurationEffects(child, core); + result.read_mask |= effects.read_mask & ~result.must_write_mask; + result.write_mask |= effects.write_mask; + result.must_write_mask |= effects.must_write_mask; } const auto *extent = control->control->extent.as(); - result.must_write &= extent && extent->value > 0; + if (!extent || extent->value <= 0) + result.must_write_mask = 0; } - result.must_write &= !node->HasConditions(); + if (node->HasConditions()) + result.must_write_mask = 0; return result; } -bool HasUnknownRegisterEffects(const IRStructure *node) { +bool HasUnknownConfigurationEffects(const IRStructure *node) { auto opaque = [](const ObjectRef &ir) { bool found = false; PostOrderVisit(ir, [&](const ObjectRef &object) { @@ -200,7 +206,7 @@ bool HasUnknownRegisterEffects(const IRStructure *node) { return opaque(control->task->stmt) || std::any_of(control->children.begin(), control->children.end(), [](const auto &child) { - return HasUnknownRegisterEffects(child.get()); + return HasUnknownConfigurationEffects(child.get()); }); } @@ -209,24 +215,24 @@ bool HasUnknownRegisterEffects(const IRStructure *node) { // has no ring of versions: each value must be consumed within one physical II. // Unlike all-pairs logical WAW edges, these edges allow independent generations // to interleave across any number of stages without changing their readers. -bool AddRegisterPhaseDependencies(const std::vector &accesses, - const std::vector &stages, - const std::vector &latencies, - std::vector *deps) { +bool AddConfigurationPhaseDependencies( + const std::vector &effects, uint16_t reg, + const std::vector &stages, const std::vector &latencies, + std::vector *deps) { std::vector writers, logical; - for (size_t i = 0; i < accesses.size(); ++i) { - if (accesses[i].writes) { + for (size_t i = 0; i < effects.size(); ++i) { + if (effects[i].write_mask & reg) { // A conditional kill may leave an older generation live. Retain the // conservative logical dependencies until that reaching definition is // known, rather than assuming that the skipped writer reset the state. - if (!accesses[i].must_write) + if (!(effects[i].must_write_mask & reg)) return false; writers.push_back(i); } logical.push_back(i); } if (writers.empty()) - return false; + return true; std::stable_sort(logical.begin(), logical.end(), [&](size_t a, size_t b) { return stages[a] < stages[b]; }); auto add = [&](size_t u, size_t v, int64_t wrap, int64_t gap) { @@ -240,7 +246,7 @@ bool AddRegisterPhaseDependencies(const std::vector &accesses, gap = std::max(gap, 1); deps->emplace_back(u, v, wrap, gap); }; - std::vector next_writer(accesses.size()); + std::vector next_writer(effects.size()); for (size_t k = 0; k < writers.size(); ++k) { size_t u = writers[k], v = writers[(k + 1) % writers.size()]; next_writer[u] = v; @@ -249,19 +255,19 @@ bool AddRegisterPhaseDependencies(const std::vector &accesses, // The first read may consume the previous logical iteration's last write. size_t writer = writers.front(); for (size_t i : logical) { - if (accesses[i].writes) + if (effects[i].write_mask & reg) writer = i; } int64_t distance = 1; for (size_t reader : logical) { - if (accesses[reader].reads) { + if (effects[reader].read_mask & reg) { int64_t wrap = distance + stages[reader] - stages[writer]; size_t overwrite = next_writer[writer]; int64_t overwrite_wrap = overwrite <= writer ? 1 : 0; add(writer, reader, wrap, latencies[writer]); add(reader, overwrite, overwrite_wrap - wrap, 0); } - if (accesses[reader].writes) { + if (effects[reader].write_mask & reg) { writer = reader; distance = 0; } @@ -1145,32 +1151,31 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, uint16_t phase_mask = 0; if (manual_schedule && enable_offset && - !std::any_of(nodes.begin(), nodes.end(), HasUnknownRegisterEffects)) { - auto add_register = [&](uint16_t reg) { - std::vector register_deps; - bool has_writer = false; - for (CoreMask core : kConcreteCores) { - std::vector accesses(n); - bool core_has_writer = false; - for (size_t i = 0; i < n; ++i) { - if (NodeCoreMask(nodes[i]) & core) - accesses[i] = SummarizeRegister(nodes[i], reg, core); - core_has_writer |= accesses[i].writes; - } - if (!core_has_writer) - continue; - has_writer = true; - if (!AddRegisterPhaseDependencies(accesses, manual_stages, latencies, - ®ister_deps)) - return false; + !std::any_of(nodes.begin(), nodes.end(), + HasUnknownConfigurationEffects)) { + std::vector> core_effects; + uint16_t written_registers = 0; + for (CoreMask core : kConcreteCores) { + auto &effects = core_effects.emplace_back(); + for (const IRStructure *node : nodes) { + effects.push_back(SummarizeConfigurationEffects(node, core)); + written_registers |= effects.back().write_mask; } + } + for (uint16_t reg = 1; reg != 0; reg <<= 1) { + if (!(written_registers & reg)) + continue; + std::vector register_deps; + if (!std::all_of(core_effects.begin(), core_effects.end(), + [&](const auto &effects) { + return AddConfigurationPhaseDependencies( + effects, reg, manual_stages, latencies, + ®ister_deps); + })) + continue; + phase_mask |= reg; register_phase_deps.insert(register_phase_deps.end(), register_deps.begin(), register_deps.end()); - return has_writer; - }; - for (uint16_t bit = 1; bit != 0; bit <<= 1) { - if ((kConfigurationRegisters & bit) && add_register(bit)) - phase_mask |= bit; } } auto add_register_dep = [&](size_t u, size_t v, int wrap, int64_t latency, From 9bfc7a9c6fe94c64b7de78624f14fa2ff7abce76 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Fri, 23 Oct 2026 12:32:50 +0800 Subject: [PATCH 10/10] [Ascend] Reuse configuration dependency scans --- src/ascend/transform/auto_schedule.cc | 396 +++++++++++++------------- 1 file changed, 200 insertions(+), 196 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 1535a9e0cc..92d46d1b42 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -81,200 +82,6 @@ using namespace ascend; namespace { -// Configuration values are state; loop-control and pipe registers are fences. -constexpr uint16_t kConfigurationRegisters = - static_cast(SpecialRegister::kHf32Mode) | - static_cast(SpecialRegister::kMmadDirection) | - static_cast(SpecialRegister::kPadValue) | - static_cast(SpecialRegister::kAtomicMode); - -using PhaseDependency = std::tuple; - -uint16_t MustWriteConfigurationRegisters(const Stmt &stmt) { - if (const auto *attr = stmt.as()) - return MustWriteConfigurationRegisters(attr->body); - if (const auto *seq = stmt.as()) { - uint16_t mask = 0; - for (const Stmt &child : seq->seq) - mask |= MustWriteConfigurationRegisters(child); - return mask; - } - if (const auto *branch = stmt.as()) - return branch->else_case.defined() - ? MustWriteConfigurationRegisters(branch->then_case) & - MustWriteConfigurationRegisters(branch->else_case.value()) - : 0; - const auto *evaluate = stmt.as(); - const auto *call = evaluate ? evaluate->value.as() : nullptr; - if (!call) - return 0; - uint16_t mask = AnalyzeSpecialRegisterWrites(stmt) & kConfigurationRegisters; - if (mask == 0) - return 0; - // Only an executed setter kills the old value. A setter nested in an - // argument (for example if_then_else) need not execute. - for (const PrimExpr &arg : call->args) - mask &= ~AnalyzeSpecialRegisterWrites(Evaluate(arg)); - return mask; -} - -// Effects on hardware configuration only. Ordinary scalar storage and variables -// use the existing buffer-region and definition-use dependencies. -struct ConfigurationEffects { - uint16_t read_mask{0}; // Incoming values that may be read. - uint16_t write_mask{0}; - uint16_t must_write_mask{0}; -}; - -CoreMask NodeCoreMask(const IRStructure *node); - -ConfigurationEffects SummarizeConfigurationEffects(const IRStructure *node, - CoreMask core) { - ConfigurationEffects result; - if (!(NodeCoreMask(node) & core)) - return result; - if (node->IsTask()) { - const auto *task = static_cast(node); - result.read_mask = node->GetSpecialReadMask() & kConfigurationRegisters; - result.write_mask = node->GetSpecialWriteMask() & kConfigurationRegisters; - result.must_write_mask = MustWriteConfigurationRegisters(task->stmt); - } else { - const auto *control = static_cast(node); - result = SummarizeConfigurationEffects(control->task.get(), core); - std::vector children; - for (const auto &child : control->children) - children.push_back(child.get()); - // Match buffer dependency analysis: stages establish logical iteration - // order; the sibling list establishes physical per-pipe issue order. - std::stable_sort(children.begin(), children.end(), - [](const IRStructure *a, const IRStructure *b) { - return a->GetStage() < b->GetStage(); - }); - for (const IRStructure *child : children) { - ConfigurationEffects effects = SummarizeConfigurationEffects(child, core); - result.read_mask |= effects.read_mask & ~result.must_write_mask; - result.write_mask |= effects.write_mask; - result.must_write_mask |= effects.must_write_mask; - } - const auto *extent = control->control->extent.as(); - if (!extent || extent->value <= 0) - result.must_write_mask = 0; - } - if (node->HasConditions()) - result.must_write_mask = 0; - return result; -} - -bool HasUnknownConfigurationEffects(const IRStructure *node) { - auto opaque = [](const ObjectRef &ir) { - bool found = false; - PostOrderVisit(ir, [&](const ObjectRef &object) { - if (const auto *call = object.as()) { - found |= !call->op.as() || - call->op.same_as(builtin::call_extern()) || - call->op.same_as(builtin::tvm_call_packed()) || - call->op.same_as(builtin::tvm_call_packed_lowered()) || - call->op.same_as(builtin::tvm_call_cpacked()) || - call->op.same_as(builtin::tvm_call_cpacked_lowered()) || - call->op.same_as(builtin::call_llvm_intrin()); - } - }); - return found; - }; - if (node->ContainsLoopBreak()) - return true; - for (const auto &guard : node->GetGuards()) { - if (guard->IsCondition()) { - if (opaque(static_cast(guard.get())->condition)) - return true; - } else { - const auto *attr = static_cast(guard.get()); - if (opaque(attr->value)) - return true; - if (const auto *expr = attr->node.as()) { - if (opaque(GetRef(expr))) - return true; - } - } - } - if (node->IsTask()) { - const auto *task = static_cast(node); - // PerCoreTask does not expose all internal configuration effects. - return task->IsPerCoreTask() || opaque(task->stmt); - } - const auto *control = static_cast(node); - return opaque(control->task->stmt) || - std::any_of(control->children.begin(), control->children.end(), - [](const auto &child) { - return HasUnknownConfigurationEffects(child.get()); - }); -} - -// Bind each read to its reaching definition in logical stage order, then -// protect that generation until the next setter. A configuration register -// has no ring of versions: each value must be consumed within one physical II. -// Unlike all-pairs logical WAW edges, these edges allow independent generations -// to interleave across any number of stages without changing their readers. -bool AddConfigurationPhaseDependencies( - const std::vector &effects, uint16_t reg, - const std::vector &stages, const std::vector &latencies, - std::vector *deps) { - std::vector writers, logical; - for (size_t i = 0; i < effects.size(); ++i) { - if (effects[i].write_mask & reg) { - // A conditional kill may leave an older generation live. Retain the - // conservative logical dependencies until that reaching definition is - // known, rather than assuming that the skipped writer reset the state. - if (!(effects[i].must_write_mask & reg)) - return false; - writers.push_back(i); - } - logical.push_back(i); - } - if (writers.empty()) - return true; - std::stable_sort(logical.begin(), logical.end(), - [&](size_t a, size_t b) { return stages[a] < stages[b]; }); - auto add = [&](size_t u, size_t v, int64_t wrap, int64_t gap) { - ICHECK(wrap == 0 || wrap == 1) - << "Manual schedule constraints are infeasible: a register read " - "would cross a physical overwrite; check T.Stage values"; - // Reconstruction sorts equal phases by descending stage, then source - // position. Preserve a required edge when that tie order would reverse it. - if (wrap == 0 && - (stages[u] < stages[v] || (stages[u] == stages[v] && u > v))) - gap = std::max(gap, 1); - deps->emplace_back(u, v, wrap, gap); - }; - std::vector next_writer(effects.size()); - for (size_t k = 0; k < writers.size(); ++k) { - size_t u = writers[k], v = writers[(k + 1) % writers.size()]; - next_writer[u] = v; - add(u, v, v <= u ? 1 : 0, latencies[u]); - } - // The first read may consume the previous logical iteration's last write. - size_t writer = writers.front(); - for (size_t i : logical) { - if (effects[i].write_mask & reg) - writer = i; - } - int64_t distance = 1; - for (size_t reader : logical) { - if (effects[reader].read_mask & reg) { - int64_t wrap = distance + stages[reader] - stages[writer]; - size_t overwrite = next_writer[writer]; - int64_t overwrite_wrap = overwrite <= writer ? 1 : 0; - add(writer, reader, wrap, latencies[writer]); - add(reader, overwrite, overwrite_wrap - wrap, 0); - } - if (effects[reader].write_mask & reg) { - writer = reader; - distance = 0; - } - } - return true; -} - // AutoSchedule target resource limits. // Check whether a kernel body contains any SIMT_VF blocks. SIMT_VF @@ -580,6 +387,196 @@ bool HasResourceDependency(const IRStructure *a, const IRStructure *b, return (a->GetHbmMask() & b->GetHbmMask()) != 0; } +// Configuration values are state; loop-control and pipe registers are fences. +constexpr uint16_t kConfigurationRegisters = + static_cast(SpecialRegister::kHf32Mode) | + static_cast(SpecialRegister::kMmadDirection) | + static_cast(SpecialRegister::kPadValue) | + static_cast(SpecialRegister::kAtomicMode); + +using PhaseDependency = std::tuple; + +uint16_t MustWriteConfigurationRegisters(const Stmt &stmt) { + if (const auto *attr = stmt.as()) + return MustWriteConfigurationRegisters(attr->body); + if (const auto *seq = stmt.as()) { + uint16_t mask = 0; + for (const Stmt &child : seq->seq) + mask |= MustWriteConfigurationRegisters(child); + return mask; + } + if (const auto *branch = stmt.as()) + return branch->else_case.defined() + ? MustWriteConfigurationRegisters(branch->then_case) & + MustWriteConfigurationRegisters(branch->else_case.value()) + : 0; + const auto *evaluate = stmt.as(); + const auto *call = evaluate ? evaluate->value.as() : nullptr; + if (!call) + return 0; + uint16_t mask = AnalyzeSpecialRegisterWrites(stmt) & kConfigurationRegisters; + if (mask == 0) + return 0; + // Only an executed setter kills the old value. A setter nested in an + // argument (for example if_then_else) need not execute. + for (const PrimExpr &arg : call->args) + mask &= ~AnalyzeSpecialRegisterWrites(Evaluate(arg)); + return mask; +} + +// Effects on hardware configuration only. Ordinary scalar storage and variables +// use the existing buffer-region and definition-use dependencies. +struct ConfigurationEffects { + uint16_t read_mask{0}; // Incoming values that may be read. + uint16_t write_mask{0}; + uint16_t must_write_mask{0}; +}; + +ConfigurationEffects SummarizeConfigurationEffects(const IRStructure *node, + CoreMask core) { + ConfigurationEffects result; + if (!(NodeCoreMask(node) & core)) + return result; + if (node->IsTask()) { + const auto *task = static_cast(node); + result.read_mask = node->GetSpecialReadMask() & kConfigurationRegisters; + result.write_mask = node->GetSpecialWriteMask() & kConfigurationRegisters; + result.must_write_mask = MustWriteConfigurationRegisters(task->stmt); + } else { + const auto *control = static_cast(node); + result = SummarizeConfigurationEffects(control->task.get(), core); + std::vector children; + for (const auto &child : control->children) + children.push_back(child.get()); + // Match buffer dependency analysis: stages establish logical iteration + // order; the sibling list establishes physical per-pipe issue order. + std::stable_sort(children.begin(), children.end(), + [](const IRStructure *a, const IRStructure *b) { + return a->GetStage() < b->GetStage(); + }); + for (const IRStructure *child : children) { + ConfigurationEffects effects = SummarizeConfigurationEffects(child, core); + result.read_mask |= effects.read_mask & ~result.must_write_mask; + result.write_mask |= effects.write_mask; + result.must_write_mask |= effects.must_write_mask; + } + const auto *extent = control->control->extent.as(); + if (!extent || extent->value <= 0) + result.must_write_mask = 0; + } + if (node->HasConditions()) + result.must_write_mask = 0; + return result; +} + +bool HasUnknownConfigurationEffects(const IRStructure *node) { + auto opaque = [](const ObjectRef &ir) { + bool found = false; + PostOrderVisit(ir, [&](const ObjectRef &object) { + if (const auto *call = object.as()) { + found |= !call->op.as() || + call->op.same_as(builtin::call_extern()) || + call->op.same_as(builtin::tvm_call_packed()) || + call->op.same_as(builtin::tvm_call_packed_lowered()) || + call->op.same_as(builtin::tvm_call_cpacked()) || + call->op.same_as(builtin::tvm_call_cpacked_lowered()) || + call->op.same_as(builtin::call_llvm_intrin()); + } + }); + return found; + }; + if (node->ContainsLoopBreak()) + return true; + for (const auto &guard : node->GetGuards()) { + if (guard->IsCondition()) { + if (opaque(static_cast(guard.get())->condition)) + return true; + } else { + const auto *attr = static_cast(guard.get()); + if (opaque(attr->value)) + return true; + if (const auto *expr = attr->node.as()) { + if (opaque(GetRef(expr))) + return true; + } + } + } + if (node->IsTask()) { + const auto *task = static_cast(node); + // PerCoreTask does not expose all internal configuration effects. + return task->IsPerCoreTask() || opaque(task->stmt); + } + const auto *control = static_cast(node); + return opaque(control->task->stmt) || + std::any_of(control->children.begin(), control->children.end(), + [](const auto &child) { + return HasUnknownConfigurationEffects(child.get()); + }); +} + +// Bind each read to its reaching definition in logical stage order, then +// protect that generation until the next setter. A configuration register +// has no ring of versions: each value must be consumed within one physical II. +// Unlike all-pairs logical WAW edges, these edges allow independent generations +// to interleave across any number of stages without changing their readers. +bool AddConfigurationPhaseDependencies( + const std::vector &effects, uint16_t reg, + const std::vector &logical_order, + const std::vector &stages, const std::vector &latencies, + std::vector *deps) { + std::vector writers; + size_t last_writer = 0; + for (size_t i = 0; i < effects.size(); ++i) { + if (effects[i].write_mask & reg) { + // A conditional kill may leave an older generation live. Retain the + // conservative logical dependencies until that reaching definition is + // known, rather than assuming that the skipped writer reset the state. + if (!(effects[i].must_write_mask & reg)) + return false; + // Stable stage order ends at the last source writer of the highest stage. + if (writers.empty() || stages[i] >= stages[last_writer]) + last_writer = i; + writers.push_back(i); + } + } + if (writers.empty()) + return true; + auto add = [&](size_t u, size_t v, int64_t wrap, int64_t gap) { + ICHECK(wrap == 0 || wrap == 1) + << "Manual schedule constraints are infeasible: a register read " + "would cross a physical overwrite; check T.Stage values"; + // Reconstruction sorts equal phases by descending stage, then source + // position. Preserve a required edge when that tie order would reverse it. + if (wrap == 0 && + (stages[u] < stages[v] || (stages[u] == stages[v] && u > v))) + gap = std::max(gap, 1); + deps->emplace_back(u, v, wrap, gap); + }; + std::vector next_writer(effects.size()); + for (size_t k = 0; k < writers.size(); ++k) { + size_t u = writers[k], v = writers[(k + 1) % writers.size()]; + next_writer[u] = v; + add(u, v, v <= u ? 1 : 0, latencies[u]); + } + // The first read may consume the previous logical iteration's last write. + size_t writer = last_writer; + int64_t distance = 1; + for (size_t reader : logical_order) { + if (effects[reader].read_mask & reg) { + int64_t wrap = distance + stages[reader] - stages[writer]; + size_t overwrite = next_writer[writer]; + int64_t overwrite_wrap = overwrite <= writer ? 1 : 0; + add(writer, reader, wrap, latencies[writer]); + add(reader, overwrite, overwrite_wrap - wrap, 0); + } + if (effects[reader].write_mask & reg) { + writer = reader; + distance = 0; + } + } + return true; +} + ScheduleBuilder::LetDependencyClosure ScheduleBuilder::CollectLetDependencyClosure_( const std::vector &nodes) { @@ -1162,6 +1159,13 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, written_registers |= effects.back().write_mask; } } + // Logical order depends only on stages, not on the core or configuration. + std::vector logical_order(written_registers ? n : 0); + std::iota(logical_order.begin(), logical_order.end(), size_t{0}); + std::stable_sort(logical_order.begin(), logical_order.end(), + [&](size_t a, size_t b) { + return manual_stages[a] < manual_stages[b]; + }); for (uint16_t reg = 1; reg != 0; reg <<= 1) { if (!(written_registers & reg)) continue; @@ -1169,8 +1173,8 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, if (!std::all_of(core_effects.begin(), core_effects.end(), [&](const auto &effects) { return AddConfigurationPhaseDependencies( - effects, reg, manual_stages, latencies, - ®ister_deps); + effects, reg, logical_order, manual_stages, + latencies, ®ister_deps); })) continue; phase_mask |= reg;