Skip to content

Latest commit

 

History

History
167 lines (132 loc) · 6.88 KB

File metadata and controls

167 lines (132 loc) · 6.88 KB

NeRArchSim Scheduler

Scheduler Overview

This folder schedules the neural-rendering operator graph defined in Operator onto the hardware modules specified in Hardware, and calculates the resulting latency. It supports both inference and training workloads, including backward pass operators and training-specific hardware modules for GSArch, GBU, and Instant3D pipelines.

Table of Contents

Unified analysis (CLI)

Use the end-to-end CLI to map, schedule, and generate reports from a traced DAG:

./nerarch_sim analyze traces/.../execution_dag.pkl results \
  --hardware Hardware/examples/hardware_configs/icarus_config.json

Outputs are organized under results/ (inputs, mapping, scheduling, reports, visuals).

Optional: DOT subgraph extraction

If you have a grouped DOT (e.g., execution_dag_grouped.dot), you can extract and render a specific cluster as part of analyze:

./nerarch_sim analyze traces/.../execution_dag.pkl results \
  --hardware Hardware/examples/hardware_configs/icarus_config.json \
  --plot-dot-subgraph \
  --dot execution_dag_grouped.dot \
  --cluster-index coarse:0 \
  --subgraph-out-prefix execution_dag_component0

Subgraph PNG/SVG/DOT will be written under results/visuals/ with the given prefix. For the standalone helper, see Instrumentation/README.md.

Overview

The scheduler maps transformed operators to hardware units and produces an execution order with start/end times and resource assignments. It supports:

  • Operator-level scheduling (intra-operator ordering and constraints)
  • System-level scheduling (resource-aware placement and timing)
  • Optional PPA estimation using the C++ core
  • Training-aware scheduling with backward operator support
  • Specialized mapping for training-specific hardware modules (FRM, BUM, GradientCompute, etc.)

Components

  • Python
    • mapping/: load and validate hardware configs
    • IR/: scheduler IR utilities
    • op_sched/: operator-level scheduling utilities
    • gscore_schedule.py: example design-space exploration flow (GSCore)
  • C++ core (cpp/)
    • operator_scheduler, system_scheduler, mapping_engine, ppa_estimator
    • Python bindings built via ./build_cpp.sh

Inputs and outputs

  • Inputs
    • mapped_ir.json (from map stage) or an execution DAG transformed via the unified CLI
    • Hardware config JSON (see Hardware/examples/hardware_configs/)
    • Optional optimization config JSON (see Scheduler/optimization/configs/)
  • Outputs
    • scheduled_ir.json: per-op timing, resource assignments, and dependencies
    • Optional PPA metrics embedded in the schedule

How to run (CLI)

Run scheduling only:

./nerarch_sim schedule results/mapping/mapped_ir.json \
  -o results/scheduling/scheduled_ir.json \
  --hardware Hardware/examples/hardware_configs/icarus_config.json

End-to-end (map -> schedule -> report) is documented in the root README and above.

Enable optimizations with explicit policy/config:

./nerarch_sim schedule results/mapping/mapped_ir.json \
  -o results/scheduling/scheduled_ir_opt.json \
  --hardware Hardware/examples/hardware_configs/icarus_config.json \
  --optimization-config Scheduler/optimization/configs/standard_neural_rendering_config.json

Optimization config behavior:

  • Optimizations are enabled when --optimization-config is provided.
  • If --optimization-config is omitted, optimizations are disabled.
  • Scheduler validates each enabled technique against selected hardware.
  • Incompatible technique/hardware combinations fail fast with assertion.

SRAM/memory config behavior:

  • memory_model_version: 2 enables strict SRAM modeling.
  • In strict mode, memory_bindings are required and there is no implicit fallback to "all SRAM blocks".
  • Bound SRAM blocks must provide explicit timing/bandwidth fields (read_bw_gbps, write_bw_gbps, read_latency_cycles, write_latency_cycles), otherwise scheduling fails fast.

Build the C++ core (optional)

./build_cpp.sh

This compiles the C++ operator/system scheduler and PPA estimator used by the CLI.

How to extend

  • Scheduling policy (operator-level)
    • Add or modify policies under Scheduler/op_sched/
    • Integrate policy selection where operator-level scheduling is invoked
  • Scheduling policy (system-level / C++)
    • Extend operator_scheduler or system_scheduler in Scheduler/cpp/src/ and corresponding headers in Scheduler/cpp/include/
    • Rebuild with ./build_cpp.sh
  • Hardware units and configs
    • Update or add JSON files under Hardware/examples/hardware_configs/
    • Ensure parsing in Scheduler/mapping/hw_config.py supports new fields
  • Training pipeline support
    • Add new operator mappings in Scheduler/mapping/__init__.py and Scheduler/cpp/src/mapping/mapping_engine.cpp
    • Support backward operators by checking for "(B)" suffix in operator types
    • Hardware configs for training accelerators are in Hardware/examples/hardware_configs/ (gsarch_config.json, gbu_config.json, instant3d_config.json)

For the IR written by the scheduler, see scheduled_ir.json produced in results/scheduling/ when using the unified CLI.

Direct API (C++ bindings)

After building the C++ core, you can call the schedulers directly from Python via pybind11:

# Ensure the module path is visible
PYTHONPATH=build/Scheduler/cpp python - << 'PY'
import nerarchsim_cpp as rs

# Load mapped IR produced by the map stage
mapped_ir = rs.load_mapped_ir_from_json('results/mapping/mapped_ir.json')

# Operator-level scheduling
optlib = rs.OptimizationLibrary()
optimizer = rs.AnalyticalOperatorOptimizer(optlib)
op_scheduler = rs.OperatorLevelScheduler(optimizer)
op_scheduled_ir = op_scheduler.schedule(mapped_ir)

# System-level scheduling
sys_scheduler = rs.SystemLevelScheduler()
system_schedule = sys_scheduler.schedule(op_scheduled_ir)

print('Total cycles:', system_schedule.total_cycles)
PY

Testing

Comprehensive test suite available in the tests/ directory:

# Run all tests
cd Scheduler/tests
python run_all_tests.py

# Run specific test suites
python test_training_pipelines.py  # Training pipeline integration
python test_mapping.py              # Mapping engine tests
python test_scheduler.py            # Scheduler component tests

See tests/README.md for detailed test documentation.

Python-only prototyping

For quick experiments without modifying the C++ core, you can prototype policies in Scheduler/op_sched/ and use the unified CLI for execution. For full performance and PPA estimation, use the C++ bindings built via ./build_cpp.sh.