Skip to content

Refactor: make a Graph boundary scalar a formal parameter #2160

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChaoZheng109 merged 1 commit into from
Sep 11, 2026
Merged

Refactor: make a Graph boundary scalar a formal parameter #2160

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/investigations/2026-08-host-orch-phase-tail-is-page-faults.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ where the fault landed. Over the 447 faulting calls above 10 µs (867 faults, 19
| `recording.tensor_sources`, doubles | 24 B | 70 | 174 | **16.1%** |
| `recording.internal_fanins` (8 B) and `.predicates` (192 B), double | — | 77 | 195 | **14.7%** |
| `recording.tensor_map` entry pool, initialized on write | 128 B | 36 | 43 | 4.5% |
| `recording.scalars` (8 B) / `.scalar_sources` (16 B), double | — | 16 | 25 | 2.7% |
| `recording.scalars` (8 B) / `.scalar_sources` (16 B, since renamed `.scalar_inheritance` and narrowed to 4 B), double | — | 16 | 25 | 2.7% |

165 of the 447 faulting calls reallocated no recording-owned container at all, which
leaves the per-node vector as the allocation: it is the single largest source.
Expand Down
45 changes: 35 additions & 10 deletions simpler_setup/kernel_compiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import subprocess
import sys
import tempfile
import warnings
from functools import cache
from pathlib import Path
from typing import Optional, Union
Expand Down Expand Up @@ -363,9 +364,24 @@ def incore_compile_cache_token(self, core_type: str) -> dict[str, object]:
}

def _run_subprocess(
self, cmd: list[str], label: str, error_hint: str = "Compiler not found"
self,
cmd: list[str],
label: str,
error_hint: str = "Compiler not found",
surface_diagnostics: bool = False,
) -> subprocess.CompletedProcess:
"""Run a subprocess command with standardized logging and error handling."""
"""Run a subprocess command with standardized logging and error handling.

surface_diagnostics reports a successful compile's stderr as a warning. It goes
through `warnings` rather than `logger`: pytest shows its warnings summary with no
flag at all, while logger output below ERROR is swallowed unless the run passes
--log-cli-level, and child pytest processes do not inherit that option (conftest's
_resource_child_command forwards only --manual). A diagnostic nobody sees by
default is a diagnostic that does not exist.

It is opt-in per call site rather than global because the kernel toolchains carry
pre-existing warnings that would bury the ones a caller turned this on to see.
"""
logger.debug(f"[{label}] Command: {' '.join(cmd)}")
try:
result = subprocess.run(cmd, check=False, capture_output=True, text=True, cwd=self.project_root)
Expand All @@ -374,6 +390,8 @@ def _run_subprocess(
logger.debug(f"[{label}] stdout:\n{result.stdout}")
if result.stderr and logger.isEnabledFor(10):
logger.debug(f"[{label}] stderr:\n{result.stderr}")
if result.stderr and result.returncode == 0 and surface_diagnostics:
warnings.warn(f"[{label}] compiler diagnostics:\n{result.stderr}", stacklevel=2)
Comment thread
poursoul marked this conversation as resolved.

if result.returncode != 0:
logger.error(f"[{label}] Compilation failed: {result.stderr}")
Expand All @@ -391,6 +409,7 @@ def _compile_to_bytes(
label: str,
error_hint: str = "Compiler not found",
delete_output: bool = True,
surface_diagnostics: bool = False,
) -> bytes:
"""Run compilation command, read output file, clean up, return bytes.

Expand All @@ -399,23 +418,29 @@ def _compile_to_bytes(
output_path: Path to expected output file
label: Label for log messages
error_hint: Message for FileNotFoundError
surface_diagnostics: Report a successful compile's stderr as a warning
(see _run_subprocess)

Returns:
Binary contents of the compiled output file

Raises:
RuntimeError: If compilation fails or output file not found
"""
self._run_subprocess(cmd, label, error_hint)

if not os.path.isfile(output_path):
raise RuntimeError(f"Compilation succeeded but output file not found: {output_path}")
# The cleanup is a finally because _run_subprocess reports diagnostics through
# warnings.warn, which raises under an error-level warning filter -- a compile
# that produced its output would otherwise leave the file behind.
try:
self._run_subprocess(cmd, label, error_hint, surface_diagnostics=surface_diagnostics)

with open(output_path, "rb") as f:
binary_data = f.read()
if not os.path.isfile(output_path):
raise RuntimeError(f"Compilation succeeded but output file not found: {output_path}")

if delete_output:
os.remove(output_path)
with open(output_path, "rb") as f:
binary_data = f.read()
finally:
if delete_output and os.path.isfile(output_path):
os.remove(output_path)
logger.info(f"[{label}] Compilation {output_path} successful: {len(binary_data)} bytes")
return binary_data

Expand Down
6 changes: 3 additions & 3 deletions src/a2a3/runtime/host_build_graph/orchestration/arg_with_deps.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,9 @@ class CoreTaskArgsWithDeps : private CoreTaskArgs {
using CoreTaskArgs::add_output;
using CoreTaskArgs::add_scalar;
using CoreTaskArgs::add_scalars;
using CoreTaskArgs::add_scalars_i32;
using CoreTaskArgs::allow_early_resolve; // speculative early-dispatch hint (getter)
using CoreTaskArgs::copy_scalars_from;
using CoreTaskArgs::add_static_scalar;
using CoreTaskArgs::add_static_scalars;
using CoreTaskArgs::allow_early_resolve; // speculative early-dispatch hint (getter)
using CoreTaskArgs::set_allow_early_resolve; // speculative early-dispatch hint (setter)
using CoreTaskArgs::set_task_timing_slot; // selective task-timing slot (setter)
using CoreTaskArgs::task_timing_slot; // selective task-timing slot (getter)
Expand Down
24 changes: 14 additions & 10 deletions src/a2a3/runtime/host_build_graph/orchestration/orchestration_api.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,10 +122,10 @@ static inline TaskOutputTensors alloc_tensors(const TensorCreateInfo create_info
for (uint32_t i = 0; i < count; i++) {
args.add_output(create_infos[i]);
}
if (args.has_error) {
if (args.has_error()) {
rt->ops->report_fatal(
rt, SIMPLER_ERROR_INVALID_ARGS, __FUNCTION__, "%s",
args.error_msg ? args.error_msg : "alloc_tensors failed to construct output-only Arg"
args.error_msg() ? args.error_msg() : "alloc_tensors failed to construct output-only Arg"
);
return TaskOutputTensors{};
}
Expand All @@ -145,10 +145,10 @@ static inline TaskOutputTensors alloc_tensors(const CIs &...cis) {
}
CoreTaskArgs args;
(args.add_output(cis), ...);
if (args.has_error) {
if (args.has_error()) {
rt->ops->report_fatal(
rt, SIMPLER_ERROR_INVALID_ARGS, __FUNCTION__, "%s",
args.error_msg ? args.error_msg : "alloc_tensors failed to construct output-only Arg"
args.error_msg() ? args.error_msg() : "alloc_tensors failed to construct output-only Arg"
);
return TaskOutputTensors{};
}
Expand Down Expand Up @@ -431,7 +431,7 @@ static inline uint64_t rt_graph_function_id(Function function) {
// its own boundary copy as a parameter for exactly that reason.
template <typename Invoke>
static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const GraphTaskArgs &args, Invoke invoke) {
debug_assert(!args.has_error && "Graph boundary GraphTaskArgs construction failed");
debug_assert(!args.has_error() && "Graph boundary GraphTaskArgs construction failed");
debug_assert(args.tensor_count() <= GRAPH_MAX_TENSOR_ARGS && "Graph boundary exceeds the tensor limit");
debug_assert(
args.explicit_dep_count() == 0 && "Explicit dependencies crossing the Graph boundary are not supported"
Expand Down Expand Up @@ -467,9 +467,13 @@ static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const G
const uint64_t _begun_ns = rt_orch_phase_now_ns();
if (result.recording) {
void *handle = result.recording_handle;
// The formal parameters the body reads: the entry's own copy, not `args` -- the
// caller only lends those for the duration of this call, and recording resolves a
// task slot's origin against this object's slot array.
const GraphTaskArgs &params = *result.params;
// A std::function rather than a bare lambda because the pool takes it as one
// through the ops table's void *.
std::function<void(GraphTaskArgs &)> job = [invoke, handle](GraphTaskArgs &record_args) mutable {
std::function<void(const GraphTaskArgs &)> job = [invoke, handle](const GraphTaskArgs &record_args) mutable {
try {
if (!rt_graph_prepare(handle, record_args)) {
rt_graph_abort(handle);
Expand All @@ -491,15 +495,15 @@ static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const G
// costs nothing here, because the fallback below re-runs `invoke` -- captured by
// value, so unaffected -- rather than the job.
RuntimeContext *record_rt = current_runtime();
const bool queued =
record_rt->ops->graph_record_start != nullptr && record_rt->ops->graph_record_start(record_rt, args, &job);
const bool queued = record_rt->ops->graph_record_start != nullptr &&
record_rt->ops->graph_record_start(record_rt, params, &job);
if (!queued) {
try {
if (!rt_graph_prepare(handle, args)) {
if (!rt_graph_prepare(handle, params)) {
rt_graph_abort(handle);
return result;
}
invoke(args);
invoke(params);
(void)rt_graph_end();
} catch (...) {
rt_graph_abort(handle);
Expand Down
6 changes: 3 additions & 3 deletions src/a5/runtime/host_build_graph/orchestration/arg_with_deps.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,9 @@ class CoreTaskArgsWithDeps : private CoreTaskArgs {
using CoreTaskArgs::add_output;
using CoreTaskArgs::add_scalar;
using CoreTaskArgs::add_scalars;
using CoreTaskArgs::add_scalars_i32;
using CoreTaskArgs::allow_early_resolve; // speculative early-dispatch hint (getter)
using CoreTaskArgs::copy_scalars_from;
using CoreTaskArgs::add_static_scalar;
using CoreTaskArgs::add_static_scalars;
using CoreTaskArgs::allow_early_resolve; // speculative early-dispatch hint (getter)
using CoreTaskArgs::set_allow_early_resolve; // speculative early-dispatch hint (setter)
using CoreTaskArgs::set_task_timing_slot; // selective task-timing slot (setter)
using CoreTaskArgs::task_timing_slot; // selective task-timing slot (getter)
Expand Down
24 changes: 14 additions & 10 deletions src/a5/runtime/host_build_graph/orchestration/orchestration_api.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,10 +122,10 @@ static inline TaskOutputTensors alloc_tensors(const TensorCreateInfo create_info
for (uint32_t i = 0; i < count; i++) {
args.add_output(create_infos[i]);
}
if (args.has_error) {
if (args.has_error()) {
rt->ops->report_fatal(
rt, SIMPLER_ERROR_INVALID_ARGS, __FUNCTION__, "%s",
args.error_msg ? args.error_msg : "alloc_tensors failed to construct output-only Arg"
args.error_msg() ? args.error_msg() : "alloc_tensors failed to construct output-only Arg"
);
return TaskOutputTensors{};
}
Expand All @@ -145,10 +145,10 @@ static inline TaskOutputTensors alloc_tensors(const CIs &...cis) {
}
CoreTaskArgs args;
(args.add_output(cis), ...);
if (args.has_error) {
if (args.has_error()) {
rt->ops->report_fatal(
rt, SIMPLER_ERROR_INVALID_ARGS, __FUNCTION__, "%s",
args.error_msg ? args.error_msg : "alloc_tensors failed to construct output-only Arg"
args.error_msg() ? args.error_msg() : "alloc_tensors failed to construct output-only Arg"
);
return TaskOutputTensors{};
}
Expand Down Expand Up @@ -431,7 +431,7 @@ static inline uint64_t rt_graph_function_id(Function function) {
// its own boundary copy as a parameter for exactly that reason.
template <typename Invoke>
static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const GraphTaskArgs &args, Invoke invoke) {
debug_assert(!args.has_error && "Graph boundary GraphTaskArgs construction failed");
debug_assert(!args.has_error() && "Graph boundary GraphTaskArgs construction failed");
debug_assert(args.tensor_count() <= GRAPH_MAX_TENSOR_ARGS && "Graph boundary exceeds the tensor limit");
debug_assert(
args.explicit_dep_count() == 0 && "Explicit dependencies crossing the Graph boundary are not supported"
Expand Down Expand Up @@ -467,9 +467,13 @@ static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const G
const uint64_t _begun_ns = rt_orch_phase_now_ns();
if (result.recording) {
void *handle = result.recording_handle;
// The formal parameters the body reads: the entry's own copy, not `args` -- the
// caller only lends those for the duration of this call, and recording resolves a
// task slot's origin against this object's slot array.
const GraphTaskArgs &params = *result.params;
// A std::function rather than a bare lambda because the pool takes it as one
// through the ops table's void *.
std::function<void(GraphTaskArgs &)> job = [invoke, handle](GraphTaskArgs &record_args) mutable {
std::function<void(const GraphTaskArgs &)> job = [invoke, handle](const GraphTaskArgs &record_args) mutable {
try {
if (!rt_graph_prepare(handle, record_args)) {
rt_graph_abort(handle);
Expand All @@ -491,15 +495,15 @@ static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const G
// costs nothing here, because the fallback below re-runs `invoke` -- captured by
// value, so unaffected -- rather than the job.
RuntimeContext *record_rt = current_runtime();
const bool queued =
record_rt->ops->graph_record_start != nullptr && record_rt->ops->graph_record_start(record_rt, args, &job);
const bool queued = record_rt->ops->graph_record_start != nullptr &&
record_rt->ops->graph_record_start(record_rt, params, &job);
if (!queued) {
try {
if (!rt_graph_prepare(handle, args)) {
if (!rt_graph_prepare(handle, params)) {
rt_graph_abort(handle);
return result;
}
invoke(args);
invoke(params);
(void)rt_graph_end();
} catch (...) {
rt_graph_abort(handle);
Expand Down
29 changes: 13 additions & 16 deletions src/common/host_build_graph/device/graph_execution.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,8 +62,7 @@ bool bind_graph_topology(GraphExecution &execution) {
// range overflows the increment before any bound check can see it.
if (definition.task_count <= 0 || definition.task_count > MAX_IN_GRAPH_TASKS) return false;
// GRAPH_MAX_SCALAR_ARGS, not MAX_SCALAR_ARGS: this counts the scalars the
// Graph BOUNDARY carries, which the recorder sizes with
// GraphTaskArgs = Arg<GRAPH_MAX_TENSOR_ARGS, GRAPH_MAX_SCALAR_ARGS> and the
// Graph BOUNDARY carries, which the recorder sizes with GraphTaskArgs and the
// outer Graph payload hands it to GraphExecution, never through an in-graph task
// payload. MAX_SCALAR_ARGS is the per-AICore-task cap (16) and applies to
// InGraphTaskDefinition::scalar_count below, which is checked separately; using
Expand Down Expand Up @@ -373,18 +372,19 @@ GraphMaterializeResult graph_execution_materialize_slice(
definition.scalar_arg_count == 0 ?
nullptr :
graph_definition_array<uint64_t>(definition, definition.off_scalars, definition.scalar_arg_count);
const GraphScalarSourceRef *scalar_sources =
definition.scalar_arg_count == 0 ? nullptr :
graph_definition_array<GraphScalarSourceRef>(
definition, definition.off_scalar_sources, definition.scalar_arg_count
);
const GraphScalarInheritance *scalar_inheritance =
definition.scalar_arg_count == 0 ?
nullptr :
graph_definition_array<GraphScalarInheritance>(
definition, definition.off_scalar_inheritance, definition.scalar_arg_count
);
const GraphPredicate *predicates =
definition.predicate_count == 0 ?
nullptr :
graph_definition_array<GraphPredicate>(definition, definition.off_predicates, definition.predicate_count);
if (tasks == nullptr || in_graph_task_offsets == nullptr ||
(definition.tensor_arg_count != 0 && (definition_tensors == nullptr || tensor_sources == nullptr)) ||
(definition.scalar_arg_count != 0 && (definition_scalars == nullptr || scalar_sources == nullptr)) ||
(definition.scalar_arg_count != 0 && (definition_scalars == nullptr || scalar_inheritance == nullptr)) ||
(definition.predicate_count != 0 && predicates == nullptr)) {
execution.materialize_busy.store(0, std::memory_order_release);
return GraphMaterializeResult::INVALID;
Expand Down Expand Up @@ -462,18 +462,15 @@ GraphMaterializeResult graph_execution_materialize_slice(
uint64_t *task_scalars = payload.scalar_data();
for (int32_t j = 0; j < source.scalar_count; ++j) {
const int32_t scalar_index = source.scalar_offset + j;
const GraphScalarSourceRef &ref = scalar_sources[scalar_index];
if (ref.source_kind == static_cast<uint8_t>(GraphScalarSourceKind::STATIC_VALUE)) {
const GraphScalarInheritance &ref = scalar_inheritance[scalar_index];
if (!ref.inherited()) {
task_scalars[j] = definition_scalars[scalar_index];
} else if (ref.source_kind == static_cast<uint8_t>(GraphScalarSourceKind::BOUNDARY)) {
if (ref.source_index >= execution.boundary_scalar_count || execution.boundary_scalars == nullptr) {
} else {
if (ref.boundary_index() >= execution.boundary_scalar_count || execution.boundary_scalars == nullptr) {
execution.materialize_busy.store(0, std::memory_order_release);
return GraphMaterializeResult::INVALID;
}
task_scalars[j] = execution.boundary_scalars[ref.source_index];
} else {
execution.materialize_busy.store(0, std::memory_order_release);
return GraphMaterializeResult::INVALID;
task_scalars[j] = execution.boundary_scalars[ref.boundary_index()];
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}
}
reset_graph_payload(payload);
Expand Down
Loading
Loading