tl.composite: fused epilogue ops with per-op scope
Extend tl.composite() with an ordered epilogue list. Each op carries
a scope flag - output_tile (default, runs once per (m,n) before
STORE), k_tile (every K-tile right after GEMM), or kernel. Plan
generator slots MATH stages by scope; pe_math reuses pe_dma's
local-loop pattern so chained epilogues (bias->relu) skip the port
hop. op_log captures per-stage params for telemetry. Topology
gains a gemm->math edge (snapshot test updated).
API stays backward-compatible - `epilogue=` is opt-in.
Example:
h = tl.composite(
op="gemm", a=a, b=b, out_ptr=int(out),
epilogue=[
{"op": "dequant", "scale": s_per_k, "scope": "k_tile"},
{"op": "bias", "bias": bias_vec},
{"op": "relu"},
{"op": "scale", "factor": 0.5},
],
)
tl.wait(h)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -23,6 +23,7 @@ def generate_gemm_plan(
|
||||
pe_prefix: str,
|
||||
a_pinned: bool = False,
|
||||
b_pinned: bool = False,
|
||||
epilogue_specs: tuple = (),
|
||||
) -> PipelinePlan:
|
||||
"""Generate GEMM tile plan: M→N→K order.
|
||||
|
||||
@@ -46,7 +47,15 @@ def generate_gemm_plan(
|
||||
dma_id = f"{pe_prefix}.pe_dma"
|
||||
fetch_id = f"{pe_prefix}.pe_fetch_store"
|
||||
gemm_id = f"{pe_prefix}.pe_gemm"
|
||||
# math_id = f"{pe_prefix}.pe_math" # for K-accumulation if needed
|
||||
math_id = f"{pe_prefix}.pe_math"
|
||||
|
||||
# Split epilogue_specs by scope. Lazy import to avoid circular dep
|
||||
# between pe_commands ← pe_types ← tiling.
|
||||
from kernbench.common.pe_commands import Scope as _Scope
|
||||
k_tile_ops = [o for o in epilogue_specs
|
||||
if getattr(o, "scope", None) == _Scope.K_TILE]
|
||||
out_tile_ops = [o for o in epilogue_specs
|
||||
if getattr(o, "scope", None) == _Scope.OUTPUT_TILE]
|
||||
|
||||
tiles: list[TilePlan] = []
|
||||
tile_id = 0
|
||||
@@ -106,9 +115,35 @@ def generate_gemm_plan(
|
||||
},
|
||||
))
|
||||
|
||||
# K-tile-scope epilogue MATH stages (e.g. dequant) — run on
|
||||
# every K-tile right after GEMM, before the accumulator
|
||||
# advances to the next K slice.
|
||||
for op in k_tile_ops:
|
||||
stages.append(Stage(
|
||||
stage_type=StageType.MATH,
|
||||
component=math_id,
|
||||
params={
|
||||
"op_kind": op.kind,
|
||||
"num_elements": tile_m * tile_n,
|
||||
"scope": "k_tile",
|
||||
},
|
||||
))
|
||||
|
||||
# STORE + DMA_WRITE only on last K-tile per (m,n). The C
|
||||
# accumulator stays in RegFile across the K loop.
|
||||
if last_k:
|
||||
# Output-tile-scope epilogue MATH (bias, relu, ...) runs
|
||||
# ONCE per (m,n) after the final K-tile, before writeback.
|
||||
for op in out_tile_ops:
|
||||
stages.append(Stage(
|
||||
stage_type=StageType.MATH,
|
||||
component=math_id,
|
||||
params={
|
||||
"op_kind": op.kind,
|
||||
"num_elements": tile_m * tile_n,
|
||||
"scope": "output_tile",
|
||||
},
|
||||
))
|
||||
stages.append(Stage(
|
||||
stage_type=StageType.STORE,
|
||||
component=fetch_id,
|
||||
|
||||
Reference in New Issue
Block a user