diff --git a/scripts/crossrepo_contract.py b/scripts/crossrepo_contract.py index 0576f04..0f2b4ee 100644 --- a/scripts/crossrepo_contract.py +++ b/scripts/crossrepo_contract.py @@ -1845,6 +1845,287 @@ def _2d_tiled_conv_sigmoid_case() -> ContractCase: ) +def _cotiled_concat_2d_case() -> ContractCase: + """A pre-spatial channel Concat of two same-resolution skips feeding a + Conv, forced to 2D (HW) tiling at a 24K budget. + + up and skip are both NCHW [1, 32, 66, 66]. A channel-axis Concat builds + cat [1, 64, 66, 66], which a 3x3 stride-1 pad-1 Conv maps to output + [1, 32, 66, 66]. At this budget the greedy temporal partition keeps the + Concat and the Conv as separate over-budget stages. Task 1's eligibility + change is what lets the multi-input Concat stage tile on both axes at all + (any Concat stage was previously excluded from HW tiling); both the + Concat and the Conv stage now solve a TILE_AXIS_HW tile with tiles_h > 1 + and tiles_w > 1, so the plan's first tile-plan record proves 2D. This is + the co-tiled skip contract: the runtime loads up and skip at the same + tile rectangle and must reproduce the whole-op result bit-exact. + """ + h = w = 66 + c = 32 + kernel, stride, pad = 3, 1, 1 + rng = np.random.default_rng(17) + up = helper.make_tensor_value_info("up", TensorProto.FLOAT, [1, c, h, w]) + skip = helper.make_tensor_value_info( + "skip", TensorProto.FLOAT, [1, c, h, w] + ) + output = helper.make_tensor_value_info( + "output", TensorProto.FLOAT, [1, c, h, w] + ) + weights = numpy_helper.from_array( + rng.normal(0.0, 0.05, size=(c, 2 * c, kernel, kernel)).astype( + np.float32 + ), + "weights", + ) + bias = numpy_helper.from_array( + rng.normal(0.0, 0.05, size=(c,)).astype(np.float32), "bias" + ) + model = _model( + "cotiled_concat_2d", + [ + helper.make_node("Concat", ["up", "skip"], ["cat"], axis=1), + helper.make_node( + "Conv", + ["cat", "weights", "bias"], + ["output"], + name="conv0", + kernel_shape=[kernel, kernel], + strides=[stride, stride], + pads=[pad, pad, pad, pad], + ), + ], + [up, skip], + [output], + [weights, bias], + ) + inputs = { + "up": rng.uniform(-1.0, 1.0, size=(1, c, h, w)).astype(np.float32), + "skip": rng.uniform(-1.0, 1.0, size=(1, c, h, w)).astype(np.float32), + } + return ContractCase( + "float_cotiled_concat_2d", + model, + model, + inputs, + expected_operators=("Concat", "Conv"), + mem_budget="24K", + expect_tiled=True, + expect_2d=True, + ) + + +def _qdq_cotiled_concat_2d_case() -> ContractCase: + """The int8 sibling of _cotiled_concat_2d_case, built via the _qdq + QDQ pattern. + + up and skip are NCHW [1, 128, 66, 66]; a channel-axis Concat builds + cat [1, 256, 66, 66] which a 3x3 stride-1 pad-1 Conv maps to int8 output + [1, 128, 66, 66]. up, skip, and cat share one scale/zero-point so the + Concat is a lossless channel copy (the co-tiled multi-input LOAD path, + not the requant path, is what this gate exercises); the single int8 + rounding boundary is the Conv output, exactly as in _qdq_2d_tiled_conv. + Both the Concat and the Conv stage tile on TILE_AXIS_HW with tiles_h > 1 + and tiles_w > 1. + """ + h = w = 66 + c = 128 + kernel, stride, pad = 3, 1, 1 + rng = np.random.default_rng(23) + up = helper.make_tensor_value_info("up", TensorProto.FLOAT, [1, c, h, w]) + skip = helper.make_tensor_value_info( + "skip", TensorProto.FLOAT, [1, c, h, w] + ) + model_output = helper.make_tensor_value_info( + "output", TensorProto.FLOAT, [1, c, h, w] + ) + int8_output = helper.make_tensor_value_info( + "output_q", TensorProto.INT8, [1, c, h, w] + ) + + def _scalar(value: float, name: str, dtype=np.float32) -> onnx.TensorProto: + return numpy_helper.from_array(np.array([value], dtype=dtype), name) + + # up, skip, and cat share one scale so the Concat rescales nothing. + skip_scale = _scalar(0.02, "skip_scale") + skip_zero_point = _scalar(0, "skip_zero_point", np.int8) + up_scale = _scalar(0.02, "up_scale") + up_zero_point = _scalar(0, "up_zero_point", np.int8) + cat_scale = _scalar(0.02, "cat_scale") + cat_zero_point = _scalar(0, "cat_zero_point", np.int8) + output_scale = _scalar(0.05, "output_scale") + output_zero_point = _scalar(0, "output_zero_point", np.int8) + weight = numpy_helper.from_array( + rng.normal(0.0, 0.05, size=(c, 2 * c, kernel, kernel)).astype( + np.float32 + ), + "weight", + ) + weight_scale = _scalar(0.01, "weight_scale") + weight_zero_point = _scalar(0, "weight_zero_point", np.int8) + initializers = [ + up_scale, + up_zero_point, + skip_scale, + skip_zero_point, + cat_scale, + cat_zero_point, + output_scale, + output_zero_point, + weight, + weight_scale, + weight_zero_point, + ] + nodes = [ + helper.make_node( + "QuantizeLinear", ["up", "up_scale", "up_zero_point"], ["up_q"] + ), + helper.make_node( + "DequantizeLinear", + ["up_q", "up_scale", "up_zero_point"], + ["up_dq"], + ), + helper.make_node( + "QuantizeLinear", + ["skip", "skip_scale", "skip_zero_point"], + ["skip_q"], + ), + helper.make_node( + "DequantizeLinear", + ["skip_q", "skip_scale", "skip_zero_point"], + ["skip_dq"], + ), + helper.make_node("Concat", ["up_dq", "skip_dq"], ["cat"], axis=1), + helper.make_node( + "QuantizeLinear", ["cat", "cat_scale", "cat_zero_point"], ["cat_q"] + ), + helper.make_node( + "DequantizeLinear", + ["cat_q", "cat_scale", "cat_zero_point"], + ["cat_dq"], + ), + helper.make_node( + "QuantizeLinear", + ["weight", "weight_scale", "weight_zero_point"], + ["weight_q"], + ), + helper.make_node( + "DequantizeLinear", + ["weight_q", "weight_scale", "weight_zero_point"], + ["weight_dq"], + ), + helper.make_node( + "Conv", + ["cat_dq", "weight_dq"], + ["raw"], + name="conv0", + kernel_shape=[kernel, kernel], + strides=[stride, stride], + pads=[pad, pad, pad, pad], + ), + helper.make_node( + "QuantizeLinear", + ["raw", "output_scale", "output_zero_point"], + ["output_q"], + ), + helper.make_node( + "DequantizeLinear", + ["output_q", "output_scale", "output_zero_point"], + ["output"], + ), + ] + compile_model = _model( + "qdq_cotiled_concat_2d", + nodes, + [up, skip], + [model_output], + initializers, + ) + reference_model = copy.deepcopy(compile_model) + del reference_model.graph.output[:] + reference_model.graph.output.extend([int8_output]) + onnx.checker.check_model(reference_model) + + inputs = { + "up": rng.uniform(-1.0, 1.0, size=(1, c, h, w)).astype(np.float32), + "skip": rng.uniform(-1.0, 1.0, size=(1, c, h, w)).astype(np.float32), + } + return ContractCase( + "int8_cotiled_concat_2d", + compile_model, + reference_model, + inputs, + ("Concat", "Conv"), + mem_budget="24K", + expect_tiled=True, + expect_2d=True, + ) + + +def _cotiled_add_2d_case() -> ContractCase: + """A pre-spatial residual Add of two same-resolution operands feeding a + Conv, forced to 2D (HW) tiling at a 24K budget. + + x and skip are both NCHW [1, 64, 66, 66]; Add produces added + [1, 64, 66, 66], which a 3x3 stride-1 pad-1 Conv maps to output + [1, 64, 66, 66]. Like the Concat sibling, the greedy temporal partition + keeps the Add and the Conv as separate over-budget stages; Task 1's + change lets the multi-input Add stage tile on both axes (Add was + previously excluded from HW tiling). Both stages solve a TILE_AXIS_HW + tile with tiles_h > 1 and tiles_w > 1, and the runtime loads x and skip + at the same tile rectangle. + """ + h = w = 66 + c = 64 + kernel, stride, pad = 3, 1, 1 + rng = np.random.default_rng(19) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, c, h, w]) + skip = helper.make_tensor_value_info( + "skip", TensorProto.FLOAT, [1, c, h, w] + ) + output = helper.make_tensor_value_info( + "output", TensorProto.FLOAT, [1, c, h, w] + ) + weights = numpy_helper.from_array( + rng.normal(0.0, 0.05, size=(c, c, kernel, kernel)).astype(np.float32), + "weights", + ) + bias = numpy_helper.from_array( + rng.normal(0.0, 0.05, size=(c,)).astype(np.float32), "bias" + ) + model = _model( + "cotiled_add_2d", + [ + helper.make_node("Add", ["x", "skip"], ["added"]), + helper.make_node( + "Conv", + ["added", "weights", "bias"], + ["output"], + name="conv0", + kernel_shape=[kernel, kernel], + strides=[stride, stride], + pads=[pad, pad, pad, pad], + ), + ], + [x, skip], + [output], + [weights, bias], + ) + inputs = { + "x": rng.uniform(-1.0, 1.0, size=(1, c, h, w)).astype(np.float32), + "skip": rng.uniform(-1.0, 1.0, size=(1, c, h, w)).astype(np.float32), + } + return ContractCase( + "float_cotiled_add_2d", + model, + model, + inputs, + expected_operators=("Add", "Conv"), + mem_budget="24K", + expect_tiled=True, + expect_2d=True, + ) + + def _build_convtranspose_2d( c_in: int, c_out: int, h_in: int, w_in: int, seed: int ) -> tuple[onnx.ModelProto, dict[str, Array]]: @@ -2724,6 +3005,9 @@ def _run_gate(runtime: Path, work_dir: Path) -> None: _2d_tiled_conv_case(), _qdq_2d_tiled_conv_case(), _2d_tiled_conv_sigmoid_case(), + _cotiled_concat_2d_case(), + _qdq_cotiled_concat_2d_case(), + _cotiled_add_2d_case(), _convtranspose_2d_tiled_case(), _qdq_convtranspose_2d_tiled_case(), _convtranspose_2d_partial_edge_case(), diff --git a/src/tigris/analysis/partition_spatial.py b/src/tigris/analysis/partition_spatial.py index 1de8b4d..2808f5a 100644 --- a/src/tigris/analysis/partition_spatial.py +++ b/src/tigris/analysis/partition_spatial.py @@ -234,25 +234,55 @@ def tiled_peak(th: int, tw: int) -> int: return (max(th, 1), max(tw, 1)) +def _cotileable_skip_operands( + ag: AnalyzedGraph, stage: Stage, op: OpNode +) -> bool: + """Every stage-external operand of a pre-spatial Concat/Add/Mul must be a + rank-4 tensor at the same H/W as the op output, so the executor's shared + input-halo rectangle load co-tiles it correctly. Intra-stage operands + (produced by an earlier op in the stage) are fine - they are not loaded. + + A rank-4 CONSTANT operand (an initializer, e.g. a Concat against a baked + tensor) is never a stage input, so it escapes the external same-H/W check; + but the 2D executor would still have to co-tile it against the spatial + op's input-halo rectangle and cannot tile-offset a full-size constant. + Fail closed on any such operand rather than silently emit a wrong result. + """ + out = ag.tensors.get(op.outputs[0]) + if out is None or len(out.shape) != 4: + return False + out_hw = tuple(out.shape[2:4]) + external = set(stage.input_tensors) + for name in op.inputs: + info = ag.tensors.get(name) + if info is not None and info.is_constant and len(info.shape) == 4: + return False + if name in external: + if info is None or len(info.shape) != 4 or tuple(info.shape[2:4]) != out_hw: + return False + return True + + def _stage_2d_eligible( ag: AnalyzedGraph, stage: Stage, stage_ops: list[OpNode] ) -> bool: """A stage may attempt HW tiling only if it is a standalone rank-4 stage - with at most one spatial op, no binary op, and all of whose ops implement - the HW tile contract. Multi-spatial-op stages are excluded: the runtime - executor's 2D contract is audited only for a single composed spatial op - per stage. - - Binary ops (Add, Mul) are excluded even though they pass the per-op HW - contract check below: exec_stage_tiled_2d loads every stage input using - the same conv input-halo rectangle, so a binary op's second operand - (e.g. a residual skip tensor) is not guaranteed to be co-tiled with the - spatial op's output at that rectangle. Admitting a stage like - [Conv, Add(conv_out, skip)] as 2D would load the skip operand with the - wrong region and size and silently produce a wrong result. This is the - conservative, fail-closed choice: such a stage falls back to the 1D path - (or fails closed if 1D is also infeasible). Co-tiled-binary 2D tiling is - a future refinement, not attempted here. + with at most one spatial op, and all of whose ops implement the HW tile + contract. Multi-spatial-op stages are excluded: the runtime executor's + 2D contract is audited only for a single composed spatial op per stage. + + Add/Mul/Concat are admitted only when they are consumed strictly BEFORE + the spatial op (their operands then sit at the spatial op's input + resolution, which is exactly the halo rectangle exec_stage_tiled_2d + loads every stage input at) and every stage-external operand is a + same-H/W co-tileable skip (see _cotileable_skip_operands). A stage like + [Add(input, skip), Conv] is admitted this way. Post-spatial occurrences + (e.g. [Conv, Add(conv_out, skip)]) and different-resolution operands + keep the fail-closed rejection: exec_stage_tiled_2d loads every stage + input using the spatial op's own input-halo rectangle, and a post-spatial + or different-resolution operand is not guaranteed to be co-tiled with + that rectangle. Admitting such a stage as 2D would load the operand with + the wrong region and size and silently produce a wrong result. """ if _stage_io_ranks(ag, stage) != {4}: return False @@ -263,17 +293,24 @@ def _stage_2d_eligible( ) if spatial_count > 1: return False - if any(op.op_type in _BINARY_OPS for op in stage_ops): - return False - # Concat carries the same hazard as the binary ops above: it takes an - # independent second operand, and exec_stage_tiled_2d loads every stage - # input with the spatial op's own input-halo rectangle, which is not - # guaranteed to be co-tiled with a distinct Concat operand at output - # resolution. Exclude it from HW eligibility for symmetry with _BINARY_OPS; - # such a stage falls back to the 1D path or fails closed. A Concat that is - # a stage head/fan-in is unaffected (it is not a single-spatial-op stage). - if any(op.op_type == "Concat" for op in stage_ops): - return False + # Locate the single spatial op (CONV/POOL); spatial_count <= 1 is ensured above. + spatial_idx = next( + (i for i, o in enumerate(stage_ops) + if classify_op(o.op_type) in (TileCategory.CONV, TileCategory.POOL)), + None, + ) + # Concat/Add/Mul are admitted only when consumed strictly BEFORE the spatial + # op (so their operands sit at the spatial op's input resolution = the halo + # rectangle the executor loads every stage input at) and every stage-external + # operand is a same-resolution skip. Post-spatial or different-resolution + # operands are the deferred cases: fail closed. op_indices/stage_ops are in + # topological (execution) order, so list index is dependency order. + for i, op in enumerate(stage_ops): + if op.op_type in _BINARY_OPS or op.op_type == "Concat": + if spatial_idx is not None and i >= spatial_idx: + return False + if not _cotileable_skip_operands(ag, stage, op): + return False return all(_op_supports_axis(op, TILE_AXIS_HW) for op in stage_ops) diff --git a/src/tigris/analysis/validation.py b/src/tigris/analysis/validation.py index 7ba25f5..287a2c5 100644 --- a/src/tigris/analysis/validation.py +++ b/src/tigris/analysis/validation.py @@ -2,6 +2,7 @@ from dataclasses import dataclass +from tigris.analysis.lifetime import compute_lifetimes from tigris.analysis.partition_spatial import ( _back_propagate_tile_heights, _chain_fast_bytes, @@ -490,27 +491,57 @@ def slow_pool_usage(ag: AnalyzedGraph) -> SlowMemoryUsage: A stage needs tiling when its peak exceeds the TOTAL fast pool (fast + reserve), so compressed and uncompressed compiles gate identically and match analyze (where reserve is 0). + + The slow-resident set is every tensor that crosses a stage boundary (a + stage's input or output tensor). PSRAM is freed at STAGE granularity, not + at per-op granularity, so residency is measured per tiled stage over that + stage's whole op-step INTERVAL, never sampled at individual op steps. + + For a tiled stage S, let its op-step interval be + [first, last] = [min(op_indices), max(op_indices)]. A boundary tensor is + slow-resident during S iff its lifetime interval overlaps that interval: + birth_step <= last and death_step >= first. The stage's residency is the + sum of those tensors' sizes; the peak is the max over tiled stages, and a + stage overflows when its sum exceeds slow_budget. + + Interval overlap counts a stage's own inputs AND outputs concurrently + (both always overlap S) as well as a long-lived skip that spans S (its + interval still overlaps). A per-op-step sample under-counts a MULTI-OP + stage whose input dies at an early op step and whose output is born at a + later op step to max(input, output): no single sampled step sees both, + even though both occupy slow memory for the whole stage. Interval overlap + upper-bounds true concurrent residency, the conservative choice for a + fail-closed budget check. """ slow_budget = ag.budget.slow if slow_budget <= 0 or not ag.stages: return SlowMemoryUsage(0, slow_budget, ()) fast_total = ag.budget.fast + ag.budget.fast_reserve + ag = compute_lifetimes(ag) + + # Slow-resident set: every tensor that crosses a stage boundary. + slow_names: set[str] = set() + for s in ag.stages: + slow_names.update(s.input_tensors) + slow_names.update(s.output_tensors) + slow_lifetimes = [ag.lifetimes[n] for n in slow_names if n in ag.lifetimes] + + def interval_bytes(first: int, last: int) -> int: + return sum( + lt.size_bytes + for lt in slow_lifetimes + if lt.birth_step <= last and lt.death_step >= first + ) + peak = 0 overflow: list[int] = [] for s in ag.stages: - if s.peak_bytes > fast_total: - in_size = sum( - ag.tensors[n].size_bytes for n in s.input_tensors - if n in ag.tensors - ) - out_size = sum( - ag.tensors[n].size_bytes for n in s.output_tensors - if n in ag.tensors - ) - stage_slow = in_size + out_size - peak = max(peak, stage_slow) - if stage_slow > slow_budget: - overflow.append(s.stage_id) + if s.peak_bytes <= fast_total or not s.op_indices: + continue + stage_peak = interval_bytes(min(s.op_indices), max(s.op_indices)) + peak = max(peak, stage_peak) + if stage_peak > slow_budget: + overflow.append(s.stage_id) return SlowMemoryUsage(peak, slow_budget, tuple(overflow)) diff --git a/tests/test_2d_tiling.py b/tests/test_2d_tiling.py index 972147a..2711667 100644 --- a/tests/test_2d_tiling.py +++ b/tests/test_2d_tiling.py @@ -347,6 +347,138 @@ def test_conv_plus_unary_stage_still_goes_hw(): assert stage_plan.axis == TILE_AXIS_HW +# Pre-spatial co-tiled skip connections (Add/Concat consumed BEFORE the +# spatial op) are a different hazard shape than the post-spatial cases +# above: the skip operand sits at the spatial op's INPUT resolution, which +# is exactly the halo rectangle exec_stage_tiled_2d loads every stage input +# at. When every stage-external operand is same-H/W, the shared-rectangle +# load co-tiles it correctly, so these stages must be admitted. Different +# resolution operands keep the fail-closed rejection. + + +def build_high_res_add_conv_graph(h, w, c, mem_budget): + """Pre-spatial: Add(input, skip) -> Conv. The Add is consumed by the Conv, + both operands same [1,c,h,w]. Co-tileable skip -> must be 2D eligible.""" + add = OpNode(name="add", op_type="Add", + inputs=["input", "skip"], outputs=["added"], attrs={}) + conv = OpNode(name="conv", op_type="Conv", + inputs=["added"], outputs=["output"], + attrs={"kernel_shape": [3, 3], "strides": [1, 1], "dilations": [1, 1]}) + stage = Stage(stage_id=0, op_indices=[0, 1], + input_tensors=["input", "skip"], output_tensors=["output"], + peak_bytes=c * h * w) + return AnalyzedGraph( + ops=[add, conv], stages=[stage], + tensors={ + "input": TensorInfo("input", (1, c, h, w), TensorProto.INT8), + "skip": TensorInfo("skip", (1, c, h, w), TensorProto.INT8), + "added": TensorInfo("added", (1, c, h, w), TensorProto.INT8), + "output": TensorInfo("output", (1, c, h, w), TensorProto.INT8), + }, + budget=MemoryBudget(fast=mem_budget)) + + +def build_high_res_concat_conv_graph(h, w, c, mem_budget): + """Pre-spatial: Concat(up, skip) -> Conv. Concat on channel axis, both + operands same H/W. Co-tileable skip -> must be 2D eligible.""" + concat = OpNode(name="concat", op_type="Concat", + inputs=["up", "skip"], outputs=["cat"], attrs={"axis": 1}) + conv = OpNode(name="conv", op_type="Conv", + inputs=["cat"], outputs=["output"], + attrs={"kernel_shape": [3, 3], "strides": [1, 1], "dilations": [1, 1]}) + stage = Stage(stage_id=0, op_indices=[0, 1], + input_tensors=["skip", "up"], output_tensors=["output"], + peak_bytes=c * h * w) + return AnalyzedGraph( + ops=[concat, conv], stages=[stage], + tensors={ + "up": TensorInfo("up", (1, c, h, w), TensorProto.INT8), + "skip": TensorInfo("skip", (1, c, h, w), TensorProto.INT8), + "cat": TensorInfo("cat", (1, 2 * c, h, w), TensorProto.INT8), + "output": TensorInfo("output", (1, c, h, w), TensorProto.INT8), + }, + budget=MemoryBudget(fast=mem_budget)) + + +def build_high_res_concat_conv_diffres_graph(h, w, c, mem_budget): + """Pre-spatial Concat but the skip is a DIFFERENT resolution (h//2 x w//2). + Not co-tileable -> must stay rejected (clause 3 fail-closed).""" + concat = OpNode(name="concat", op_type="Concat", + inputs=["up", "skip"], outputs=["cat"], attrs={"axis": 1}) + conv = OpNode(name="conv", op_type="Conv", + inputs=["cat"], outputs=["output"], + attrs={"kernel_shape": [3, 3], "strides": [1, 1], "dilations": [1, 1]}) + stage = Stage(stage_id=0, op_indices=[0, 1], + input_tensors=["skip", "up"], output_tensors=["output"], + peak_bytes=c * h * w) + return AnalyzedGraph( + ops=[concat, conv], stages=[stage], + tensors={ + "up": TensorInfo("up", (1, c, h, w), TensorProto.INT8), + "skip": TensorInfo("skip", (1, c, h // 2, w // 2), TensorProto.INT8), + "cat": TensorInfo("cat", (1, 2 * c, h, w), TensorProto.INT8), + "output": TensorInfo("output", (1, c, h, w), TensorProto.INT8), + }, + budget=MemoryBudget(fast=mem_budget)) + + +def build_high_res_concat_const_skip_graph(h, w, c, mem_budget): + """Pre-spatial Concat whose skip operand is a rank-4 CONSTANT (initializer), + same H/W as `up`. A constant is never a stage input, so it escapes the + external same-H/W co-tile check in _cotileable_skip_operands, yet the 2D + executor cannot tile-offset a full-size constant against the spatial op's + input-halo rectangle. Must fail closed.""" + concat = OpNode(name="concat", op_type="Concat", + inputs=["up", "const_skip"], outputs=["cat"], attrs={"axis": 1}) + conv = OpNode(name="conv", op_type="Conv", + inputs=["cat"], outputs=["output"], + attrs={"kernel_shape": [3, 3], "strides": [1, 1], "dilations": [1, 1]}) + stage = Stage(stage_id=0, op_indices=[0, 1], + input_tensors=["up"], output_tensors=["output"], + peak_bytes=c * h * w) + return AnalyzedGraph( + ops=[concat, conv], stages=[stage], + tensors={ + "up": TensorInfo("up", (1, c, h, w), TensorProto.INT8), + "const_skip": TensorInfo( + "const_skip", (1, c, h, w), TensorProto.INT8, is_constant=True + ), + "cat": TensorInfo("cat", (1, 2 * c, h, w), TensorProto.INT8), + "output": TensorInfo("output", (1, c, h, w), TensorProto.INT8), + }, + budget=MemoryBudget(fast=mem_budget)) + + +def test_stage_2d_eligible_admits_pre_spatial_add_skip(): + ag = build_high_res_add_conv_graph(h=256, w=256, c=256, mem_budget=24 * 1024) + assert _stage_2d_eligible(ag, ag.stages[0], ag.ops) is True + + +def test_stage_2d_eligible_admits_pre_spatial_concat_skip(): + ag = build_high_res_concat_conv_graph(h=256, w=256, c=256, mem_budget=24 * 1024) + assert _stage_2d_eligible(ag, ag.stages[0], ag.ops) is True + + +def test_stage_2d_eligible_rejects_constant_concat_skip(): + # A rank-4 constant Concat operand is not a stage input, so it escapes the + # external same-H/W co-tile check, but it cannot be tile-offset for the 2D + # executor's shared input-halo rectangle load. The stage must fail closed. + ag = build_high_res_concat_const_skip_graph(h=256, w=256, c=256, mem_budget=24 * 1024) + assert _stage_2d_eligible(ag, ag.stages[0], ag.ops) is False + + +def test_stage_2d_eligible_rejects_diffres_skip(): + ag = build_high_res_concat_conv_diffres_graph(h=256, w=256, c=256, mem_budget=24 * 1024) + assert _stage_2d_eligible(ag, ag.stages[0], ag.ops) is False + + +def test_pre_spatial_concat_skip_goes_hw(): + # End to end: the admitted pre-spatial concat-skip stage tiles on TILE_AXIS_HW. + ag = build_high_res_concat_conv_graph(h=256, w=256, c=256, mem_budget=24 * 1024) + ag = partition_spatial(ag) + assert plan_for_single_stage(ag).axis == TILE_AXIS_HW + + # ConvTranspose 2D tiling over the OUTPUT extent. # # ConvTranspose stays UNTILEABLE in _OP_CATEGORY (so it is auto-excluded from diff --git a/tests/test_slow_pool_concurrent.py b/tests/test_slow_pool_concurrent.py new file mode 100644 index 0000000..4d8a630 --- /dev/null +++ b/tests/test_slow_pool_concurrent.py @@ -0,0 +1,173 @@ +"""Slow-tier (PSRAM) budget must model concurrent residency, not just the +coarse per-stage input+output max. A tensor that crosses several stage +boundaries (a long-lived skip) stays resident in slow memory for every +step in between, and can coexist with a middle stage's own boundary +tensors in a way no single stage's own input+output sum captures. +""" + +from onnx import TensorProto + +from tigris.analysis.validation import slow_pool_usage +from tigris.graph.ir import AnalyzedGraph, MemoryBudget, OpNode, Stage, TensorInfo + + +def build_long_lived_skip_graph() -> AnalyzedGraph: + """Four single-op stages. "skip" is born at stage 0's output and only + consumed by stage 3, so it stays slow-resident through stage 1 and + stage 2. Each stage has exactly one op, so op step index == stage + op_indices == lifetime step index. + + Tensor sizes (bytes; INT8 dtype so num_elements == size_bytes): + input=64, skip=1000 (the long-lived skip), mid1=2000, mid2=64, + output=64. + + Old coarse per-stage peak (that stage's own input_tensors + + output_tensors sizes, summed independently per stage): + stage0: input(64) + skip(1000) = 1064 + stage1: input(64) + mid1(2000) = 2064 + stage2: mid1(2000) + mid2(64) = 2064 + stage3: (mid2(64) + skip(1000)) + output(64) = 1128 + old coarse peak = max(...) = 2064 + + Concurrent peak (closed-closed birth_step <= t <= death_step, matching + _live_bytes_by_step / stage.peak_bytes in partition_temporal.py; + lifetimes: input birth=-1 death=1, skip birth=0 death=3, mid1 birth=1 + death=2, mid2 birth=2 death=3, output birth=3 death=4): + t=0 (stage0): input(64) alive; skip born this step, already live + -> 64 + 1000 = 1064 + t=1 (stage1): input(64, dies here) + skip(1000, alive) + + mid1(2000, born this step, already live) + -> 64 + 1000 + 2000 = 3064 + t=2 (stage2): skip(1000, still alive) + mid1(2000, dies here) + + mid2(64, born this step, already live) + -> 1000 + 2000 + 64 = 3064 + t=3 (stage3): skip(1000, dies here) + mid2(64, dies here) + + output(64, born this step, already live) + -> 1000 + 64 + 64 = 1128 + concurrent peak = max(1064, 3064, 3064, 1128) = 3064 + + 3064 > 2064: at stage 1's own op step, "skip" (produced by stage 0, + not yet consumed by stage 3) is concurrently resident alongside both + of stage 1's own boundary tensors ("input" and "mid1"), and the + coarse per-stage method never sums a boundary tensor from one stage + against a different stage's own boundary tensors. + """ + op0 = OpNode(name="tap", op_type="Conv", inputs=["input"], outputs=["skip"], step=0) + op1 = OpNode(name="branch", op_type="Conv", inputs=["input"], outputs=["mid1"], step=1) + op2 = OpNode(name="mid_conv", op_type="Conv", inputs=["mid1"], outputs=["mid2"], step=2) + op3 = OpNode(name="merge", op_type="Add", inputs=["mid2", "skip"], outputs=["output"], step=3) + + stage0 = Stage( + stage_id=0, op_indices=[0], + input_tensors=["input"], output_tensors=["skip"], + peak_bytes=100_000, + ) + stage1 = Stage( + stage_id=1, op_indices=[1], + input_tensors=["input"], output_tensors=["mid1"], + peak_bytes=100_000, + ) + stage2 = Stage( + stage_id=2, op_indices=[2], + input_tensors=["mid1"], output_tensors=["mid2"], + peak_bytes=100_000, + ) + stage3 = Stage( + stage_id=3, op_indices=[3], + input_tensors=["mid2", "skip"], output_tensors=["output"], + peak_bytes=100_000, + ) + + return AnalyzedGraph( + ops=[op0, op1, op2, op3], + stages=[stage0, stage1, stage2, stage3], + model_inputs=["input"], + model_outputs=["output"], + tensors={ + "input": TensorInfo("input", (1, 64), TensorProto.INT8), + "skip": TensorInfo("skip", (1, 1000), TensorProto.INT8), + "mid1": TensorInfo("mid1", (1, 2000), TensorProto.INT8), + "mid2": TensorInfo("mid2", (1, 64), TensorProto.INT8), + "output": TensorInfo("output", (1, 64), TensorProto.INT8), + }, + budget=MemoryBudget(fast=64), + ) + + +def test_slow_pool_counts_long_lived_skip_concurrently(): + # Interval-overlap over each single-op stage's op-step interval reproduces + # the same concurrent peak (3064) the closed-closed per-step model gave for + # this all-single-op graph, so this expectation is unchanged by the + # multi-op interval-overlap fix. + expected_concurrent_peak = 3064 + old_coarse_peak = 2064 + assert expected_concurrent_peak > old_coarse_peak + + ag = build_long_lived_skip_graph() + ag.budget = MemoryBudget(fast=ag.budget.fast, slow=expected_concurrent_peak) + usage = slow_pool_usage(ag) + assert usage.slow_peak_bytes == expected_concurrent_peak + assert usage.overflow_stage_ids == () + + ag.budget = MemoryBudget(fast=ag.budget.fast, slow=expected_concurrent_peak - 1) + usage_tight = slow_pool_usage(ag) + assert usage_tight.overflow_stage_ids != () + + +def build_multi_op_tiled_stage_graph() -> AnalyzedGraph: + """One tiled stage of two ops: Conv(in)->c, Sigmoid(c)->out. Both ``in`` + and ``out`` are stage-boundary tensors (equal size); ``c`` is intra-stage + (produced and consumed inside the stage), so it is not slow-resident. + + The stage is tiled: peak_bytes 100_000 exceeds the fast pool (64). PSRAM is + freed at STAGE granularity, but under the OLD per-op-step sampling ``in`` + dies at the Conv step (step 0) and ``out`` is born only at the Sigmoid step + (step 1), so no single sampled step counts both and the stage reports + max(in, out). Interval-overlap over the stage op-step interval [0, 1] + counts both, reporting in + out. + + Sizes (INT8, so num_elements == size_bytes): in=1000, out=1000. + """ + conv = OpNode(name="conv", op_type="Conv", inputs=["in"], outputs=["c"], step=0) + sig = OpNode(name="sig", op_type="Sigmoid", inputs=["c"], outputs=["out"], step=1) + stage = Stage( + stage_id=0, op_indices=[0, 1], + input_tensors=["in"], output_tensors=["out"], + peak_bytes=100_000, + ) + return AnalyzedGraph( + ops=[conv, sig], + stages=[stage], + model_inputs=["in"], + model_outputs=["out"], + tensors={ + "in": TensorInfo("in", (1, 1000), TensorProto.INT8), + "c": TensorInfo("c", (1, 1000), TensorProto.INT8), + "out": TensorInfo("out", (1, 1000), TensorProto.INT8), + }, + budget=MemoryBudget(fast=64), + ) + + +def test_slow_pool_counts_multi_op_stage_boundaries_concurrently(): + in_bytes = 1000 + out_bytes = 1000 + concurrent = in_bytes + out_bytes # 2000, interval-overlap + per_op_step = max(in_bytes, out_bytes) # 1000, the old under-count + assert concurrent > per_op_step + + ag = build_multi_op_tiled_stage_graph() + ag.budget = MemoryBudget(fast=ag.budget.fast, slow=concurrent) + usage = slow_pool_usage(ag) + assert usage.slow_peak_bytes == concurrent + assert usage.overflow_stage_ids == () + + # A slow budget just below in+out must be rejected. The old per-op-step + # model reports only max(in, out) and would have wrongly accepted it, + # letting a plan that overflows PSRAM through the fail-closed check. + ag.budget = MemoryBudget(fast=ag.budget.fast, slow=concurrent - 1) + assert slow_pool_usage(ag).overflow_stage_ids != () + + # Just above in+out fits. + ag.budget = MemoryBudget(fast=ag.budget.fast, slow=concurrent + 1) + assert slow_pool_usage(ag).overflow_stage_ids == ()