diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 1687d6b936..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 @@ -61,6 +62,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" @@ -385,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) { @@ -786,6 +978,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 +1146,73 @@ 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. + uint16_t phase_mask = 0; + if (manual_schedule && enable_offset && + !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; + } + } + // 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; + std::vector register_deps; + if (!std::all_of(core_effects.begin(), core_effects.end(), + [&](const auto &effects) { + return AddConfigurationPhaseDependencies( + effects, reg, logical_order, manual_stages, + latencies, ®ister_deps); + })) + continue; + phase_mask |= reg; + register_phase_deps.insert(register_phase_deps.end(), + register_deps.begin(), register_deps.end()); + } + } + 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); + }; 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 || - (!(wmask & static_cast(SpecialRegister::kLoopControl)) && + uint16_t overlap = wmask & nodes[j]->GetSpecialReadMask(); + if (j == i || overlap == 0 || + (!(overlap & 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 +1324,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 +1400,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..5faaad62b4 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -181,3 +181,199 @@ 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, 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)] + 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'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)): + 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, + stage=stage, + cost=(1, 1) if index % 2 == 0 else (64, 32), + guard=i == 0 if guarded_writer and index == 2 else None, + ) + ) + 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])) + 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("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) + 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("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", ["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", + [ + pytest.param([4, 5, 5, 5], False, id="reader-crosses-stage"), + pytest.param([4, 4, 5, 5], True, id="conditional-writer"), + ], +) +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) + # 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) + + +@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 != "opaque": + after = transform.AutoSchedule()(before) + 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_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 + 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 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, i + 10, [0]), + tirx.BufferStore(out, x[0], [i, 0]), + tirx.BufferStore(x, i + 100, [0]), + tirx.BufferStore(out, x[0], [i, 1]), + ] + 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}) + 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) diff --git a/tilelang/ascend/language/schedule_hint.py b/tilelang/ascend/language/schedule_hint.py index 51e0fdb371..e66cb0ced4 100644 --- a/tilelang/ascend/language/schedule_hint.py +++ b/tilelang/ascend/language/schedule_hint.py @@ -145,6 +145,15 @@ 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``. + 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 ------- >>> 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..51a2895121 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,12 @@ 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 + 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. Buffer dependencies stay in data_deps. + Returns ------- tuple[list[int], list[int], int] @@ -369,12 +376,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 +456,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 +634,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 +818,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 +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) 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 +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, seed=42, )