Crane: Inter-Layer Scheduling Framework for DNN Inference and Training Co-Support on Tiled Architecture

---

Review Form

Reviewer: The Guardian (Adversarial Skeptic)

Summary

The authors present Crane, a framework for inter-layer scheduling on tiled architectures that aims to provide co-support for both DNN inference and training. The central thesis is that prior works are limited due to an inadequate scheduling representation. The paper proposes a hierarchical, table-format representation (ScT and MeT) to capture four key design factors: execution scheme, fusion, recomputation, and batch splitting. This representation is used to formulate the scheduling task as a Mixed-Integer Linear Programming (MILP) problem. The evaluation claims significant improvements in Energy-Delay Product (EDP) and scheduling speed over existing frameworks like SET, Tangram, and TileFlow.

However, the work's core claims of providing a "thoroughly" explored, mathematically structured optimization appear to be undermined by the methodology itself, which is a complex multi-stage heuristic that uses MILP as a component rather than constituting a single, comprehensive optimization. The scalability of the approach is not sufficiently demonstrated, and the evaluation relies on questionable "hypothetical" baselines for its most significant claims in the training domain.

Strengths

1. Expressive Representation: The proposed table-based representation using a Scheduling Table (ScT) and Memory Table (MeT) is a clear and systematic way to model the complex state space of inter-layer scheduling, including memory lifetime and sub-batch progress (Section 5, pages 6-7). This is a well-conceived abstraction.
2. Unified Scope: The ambition to unify the scheduling of inference and training workloads under a single framework that considers execution schemes, fusion, recomputation, and batch splitting (E+F+R+B) is commendable. The paper correctly identifies this as a gap in prior research.
3. Formal Problem Components: The use of MILP to solve sub-problems within the scheduling flow is a valid approach for finding locally optimal solutions for workload and memory configurations, given a fixed block structure and sub-batch size.

Weaknesses

1. Contradiction in Optimization Claims: The paper repeatedly frames its approach as a "thoroughly" explored, "mathematically structured optimization" in contrast to "heuristic search algorithms" (Abstract, Table 1). However, the overall optimization process described in Section 5.5 (Figure 7) and Section 6 is fundamentally a high-level heuristic search. The process involves:

   • Enumerating candidates for sub-batch size (BSsub).
   • Pruning the search space to the top K1 candidates based on a Cost_comp MILP.
   • Further pruning to the top K2 candidates based on a Cost_traffic MILP.
   • An iterative, gradual partitioning for the hierarchical structure that terminates based on a pre-set threshold θ (Section 6, Step-2).

   This is not a single, thorough optimization. It is a meta-heuristic that relies on pruning and arbitrary thresholds (K1, K2, θ), which directly contradicts the paper's central criticism of prior work. The claim of avoiding "incomplete and inefficient exploration" is therefore unsubstantiated.

2. Unproven Universality and Unanalyzed Scalability: Theorem 1 (page 5) claims universality for any DAG-structured model. The provided "proof" is a high-level, informal sketch. It fails to formally address the potential complexity and depth of the required hierarchy for highly branched models, nor does it prove that the pipeline-derived state basis is sufficient for all valid execution schedules. Furthermore, the paper provides no analysis of the scalability of its core MILP formulation. The state space SB for a composite block of N sub-blocks is 2N-1. The runtime of the MILP solver is highly dependent on N. The reported scheduling speedups are based on models of moderate depth, but there is no evidence to suggest this approach would remain tractable for significantly deeper or more complex models.

3. Weak and Unverifiable Baselines in Training Evaluation: The most dramatic results (e.g., up to 21.01x EDP reduction) are reported in the training evaluation (Section 7.3). However, these results are derived from comparisons against "hypothetical training-oriented SET and Tangram" (page 11). These baselines are not established works; they are the authors' own modifications of inference-only schedulers. The validity of these comparisons is highly suspect. There is no way to verify if these hypothetical baselines were implemented in a fair, robust, or optimized manner. The critical result that SET and Tangram produce an "out-of-memory error" for OPT-6.7B (Figure 11b) is entirely dependent on the authors' unverified implementation choices for these baselines. This is not a scientifically rigorous comparison.

4. Heuristic Simplifications Within the Formalism: The method for handling recomputation (Section 5.3, page 8) appears to be a heuristic choice rather than an optimally derived one. The process is split into two distinct steps: a "Backward Pass-only Pre-processing" step followed by a "Forward Recomputation-then-Backward Pass" step. This split simplifies the dependency graph by ensuring recomputed activations are always available before they are needed. However, it is not proven that this decoupling is optimal. A truly optimal solution might interleave recomputation and backward propagation on a finer granularity. This simplification appears to be a concession to tractability that undermines the claim of optimality.

Questions to Address In Rebuttal

1. Please reconcile the claims of "thorough" exploration and moving beyond "heuristic search" with the methodology described in Section 5.5 and 6, which explicitly relies on heuristic pruning (K1, K2) and convergence thresholds (θ). Is the overall framework not, in fact, a meta-heuristic?
2. Provide empirical data or a formal analysis on the scalability of the MILP solver runtime as the number of sub-blocks N in a composite block increases. At what level of model complexity does Crane's scheduling time become prohibitive?
3. Regarding the training evaluation (Section 7.3), can the authors provide a detailed description of the implementation of the "hypothetical" SET and Tangram baselines? Specifically, how was recomputation support added, and what steps were taken to ensure this was a fair and competitive implementation rather than a strawman?
4. Justify why the two-step recomputation process (Section 5.3) is optimal. Have alternative, more coupled strategies for recomputation and backward pass scheduling been considered? If so, why were they discarded? If not, how can the current approach be considered part of a "thorough" optimization?

---

Review Form:

Reviewer: The Synthesizer (Contextual Analyst)

Summary

This paper presents Crane, a novel framework for inter-layer scheduling of DNN workloads on tiled architectures. The authors identify the root cause of limitations in prior work—such as incomplete optimization, scheduling inflexibility, and inefficient search—as the absence of a sufficiently powerful and structured scheduling representation.

The core contribution of this work is not merely a new scheduler, but a new representational framework for the scheduling problem itself. Crane introduces a hierarchical, table-based abstraction (ScT for scheduling, MeT for memory) that is expressive enough to unify the four key design factors: execution scheme (E), layer fusion (F), recomputation (R), and batch splitting (B). This unified representation is mathematically structured, allowing the authors to reformulate the scheduling problem as a Mixed-Integer Linear Program (MILP). This shift from heuristic search (e.g., Simulated Annealing) to a formal optimization method enables a more complete and efficient exploration of the design space, uniquely providing co-support for both DNN inference and training workloads within a single framework.

Strengths

1. A Fundamental Shift in Problem Representation: The paper's primary strength lies in its diagnosis of the core problem. Instead of treating the symptoms (e.g., slow heuristic search), the authors correctly identify the disease: an inadequate representation. The proposed hierarchical table-format is an elegant solution. It successfully captures the complex interplay between computation, data movement, and memory lifetime across layers, which has been a significant challenge in this domain. The three principles laid out in Section 3.1 (page 3)—rich expressiveness, topological flexibility, and mathematical structuredness—are precisely the right goals for such a representation, and the paper delivers on them.

2. Unification of Inference and Training Scheduling: This is a major step forward for the field. Historically, tools have been designed for either inference (e.g., SET, Tangram) or training (e.g., Checkmate), because the optimization objectives and constraints (especially memory pressure and the need for recomputation) are so different. By creating a representation that natively incorporates recomputation (R) as a first-class citizen alongside fusion and execution schemes, Crane provides a single, coherent framework for both domains. This is not just a matter of convenience; it allows for the discovery of scheduling strategies that holistically optimize the entire compute lifecycle and simplifies the software stack for hardware vendors and researchers.

3. Methodological Advancement from Heuristics to Formal Optimization: The move from sampling-based heuristics (like those in SET and TileFlow) to a MILP formulation is a significant increase in methodological rigor. While heuristics can be effective, they offer no guarantee of optimality and often suffer from slow convergence. By structuring the problem for a MILP solver, Crane can perform a more thorough search of the scheduling space. The impressive empirical results, showing a 2.82× scheduling speedup over SET despite exploring a larger search space, validate that this structured approach is not just theoretically sound but also practically efficient for the problem scale considered.

4. Strong Connection to the Broader Landscape: The authors do a commendable job of situating their work. Table 1 (page 2) provides a clear and concise summary of the state-of-the-art, highlighting the specific gaps Crane aims to fill. The framework builds upon a rich history of work in both intra-layer (which it pragmatically accepts as a plug-in) and inter-layer scheduling, while charting a new path forward. This work synthesizes ideas from compiler theory (hierarchical representations), optimization (MILP), and computer architecture (tiled accelerators) into a cohesive whole.

Weaknesses

My critiques are less about flaws and more about the boundaries and future implications of this approach.

1. Scalability of the MILP Formulation: The primary concern with any MILP-based solution is its scalability. While the results are excellent for the models tested, MILP problems are NP-hard in the general case. The paper’s hierarchical approach is a clever and necessary method for managing this complexity, effectively acting as a problem decomposition heuristic. However, it's unclear how this approach will scale to the next generation of truly gargantuan models (e.g., Mixture-of-Experts with complex routing, or models with thousands of layers). The runtime of the optimization process itself could become a bottleneck as the number of blocks, layers within blocks, and states grows.

2. The Optimality of the Hierarchy: Theorem 1 (page 5) establishes the universality of the representation—that any DAG-based schedule can be represented. However, the process of finding the optimal hierarchical block structure itself seems non-trivial. The gradual, cost-driven partitioning process described in Section 6 (page 10) is a heuristic. This introduces a potential weakness: the quality of the final schedule is dependent on the quality of this initial, high-level partitioning. An early suboptimal decision in structuring the blocks could prune away the true globally optimal solution.

3. Decoupling of Inter- and Intra-Layer Scheduling: The decision to treat the intra-layer scheduler as a "plug-in" that provides cost estimates (α, β, γ in Eqs. 24 & 25, page 10) is pragmatic. However, it abstracts away the deep coupling between these two problems. For instance, an inter-layer decision to fuse two layers might drastically change the optimal intra-layer tiling strategy and dataflow for that fused operator. A simple cost model might not capture these non-linear effects, potentially leading the inter-layer optimizer astray.

Questions to Address In Rebuttal

1. Regarding the scalability of the MILP solver: Could the authors comment on the growth of the MILP problem size and solver runtime as a function of model depth and width? Have they identified any "cliffs" where the problem becomes intractable, and does the hierarchical pruning strategy effectively mitigate this for the foreseeable future of model sizes?

2. Concerning the hierarchical partitioning: How sensitive is the final solution quality to the heuristic partitioning strategy described in Section 6? For example, if you start with a different initial block structure (e.g., one layer per block vs. the paper's dependency-based grouping), how much does the final EDP change? This would help clarify whether the framework consistently finds a near-optimal solution regardless of this initial step.

3. On the interaction between scheduling layers: Could the authors elaborate on the fidelity of the intra-layer cost model? Is a single-pass cost estimation sufficient, or have they considered an iterative approach where Crane's inter-layer decisions are used to refine the intra-layer schedules, which in turn provide more accurate costs back to Crane?

Review Form

Reviewer: The Innovator (Novelty Specialist)

Summary

This paper introduces Crane, an inter-layer scheduling framework for DNNs on tiled architectures. The authors identify the core problem in prior work as the lack of a unified, expressive representation for scheduling decisions. Their central claim to novelty is the creation of a hierarchical table-format representation that co-optimizes four key design factors: execution scheme (E), layer fusion (F), recomputation (R), and batch splitting (B).

This representation consists of two main components:

1. A Scheduling Table (ScT) that tracks the cumulative number of sub-batches processed by each layer/block over time.
2. A Memory Table (MeT) that tracks the lifetime of intermediate data in on-chip (SRAM) and off-chip (DRAM) memory by recording the lower bound of the valid sub-batch range.

The authors argue that this specific, structured representation is novel because it is the first to unify all four design factors, and its mathematical structure allows the complex scheduling problem to be formulated as a Mixed-Integer Linear Program (MILP), which can be solved more efficiently and thoroughly than prior heuristic-based methods.

Strengths

The primary strength of this work lies in the conceptual novelty of its core representation.

1. Unified Representation for Disparate Factors: The most significant novel contribution is the joint representation of execution scheme, fusion, and recomputation within a single, co-optimizable mathematical structure. Prior art has treated these as largely separate problems. For instance, frameworks like SET [5] and Tangram [9] focus on execution schemes and fusion, while Checkmate [15] is purpose-built for recomputation. These approaches are fundamentally incompatible. Crane's ScT/MeT formulation provides a common, low-level language to describe the effects of all these decisions, which appears to be a genuinely new approach.

2. Novel Abstraction of Scheduling States: The paper's formalization of execution states as being derived from a canonical pipeline pattern (Section 4, pages 4-5) is a clever and powerful abstraction. This provides a regular, predictable structure to the state space, in contrast to more ad-hoc or constrained representations like the ratio-tree in SET [5]. The claim of universality for any DAG-structured model (Theorem 1, page 5) is a strong theoretical underpinning that distinguishes this work from prior schedulers that are often limited to linear chains of computation.

3. Representation-Driven Optimization: The novelty is not simply the use of an MILP solver, which has appeared in prior work (e.g., Checkmate [15]). Rather, the novelty is in designing a representation specifically to enable a clean MILP formulation. The use of cumulative sub-batch counts (ScT) and memory range bounds (MeT) translates complex temporal dependencies and memory lifetime constraints into a set of linear inequalities. This tight coupling of representation and optimization is a significant conceptual advance over prior works that relied on heuristic search over less structured representations.

Weaknesses

While the combination of ideas is novel, the individual components are built upon existing concepts. The paper could do a better job of delineating the precise boundary of its novelty.

1. Incremental Novelty of Components: Hierarchical problem decomposition is a standard technique. The use of tables to track system state is fundamental. MILP is a known tool for optimization problems. The core novelty rests entirely on the specific semantics of the ScT and MeT tables. The authors present this as a revolutionary framework, but it can also be viewed as a very clever, new encoding scheme for a known solver.

2. Unclear Scalability of the Core Formulation: The reliance on MILP, even with a well-structured problem, raises questions about scalability that are not fully addressed. While the hierarchical search strategy (Section 6, page 10) is a practical solution to prune the search space, it is a heuristic layered on top of the core MILP formulation. What are the scaling properties of the MILP itself for a single, very large, "flat" block with many sub-components? The novelty of the MILP formulation is only valuable insofar as it remains tractable for problems of interest. Does the representation itself offer any intrinsic scaling advantages over, for example, the ILP formulation in Checkmate [15], especially given that Crane's problem space is significantly larger?

3. Complexity of the Recomputation Model: The model for recomputation (Section 5.3, page 8) requires a two-step process: a backward-only pass followed by a combined recomputation-and-backward pass. This appears to be an artifact of the representation's design. A truly universal and fully expressive representation might be expected to model the bi-directional dataflow of training in a single, unified pass. This two-step process, while functional, suggests a potential limitation in the expressive power of the novel table format and feels less elegant than the forward-pass model.

Questions to Address In Rebuttal

1. Prior work like SET [5] also uses a hierarchical (tree-based) representation. Please clarify the fundamental conceptual novelty of your pipeline-derived hierarchical state abstraction beyond the stated benefit of "flexibility." What does your abstraction allow that is fundamentally impossible to represent in a ratio-tree or a similar hierarchical graph?

2. The core claim is that the table format enables a more thorough and efficient search via MILP. Given that MILP solvers can have exponential worst-case complexity, what specific properties of the ScT/MeT constraint formulation (Eqs. 1-12, pages 6-8) make it particularly tractable for this expanded search space (E+F+R+B) compared to prior, more limited ILP/MILP formulations like Checkmate [15]?

3. Regarding the recomputation model in Section 5.3, is the two-step backward process (pre-processing pass and recomputation pass) a necessary consequence of the ScT/MeT representation, or is it a heuristic choice to simplify the formulation? Could the tables, as defined, represent a single, interleaved recomputation-and-backward process, and if not, what does this imply about the representation's completeness for training workloads?