pinpinpoo opened a new issue, #20342:
URL: https://github.com/apache/tvm/issues/20342
Enabling `FuseOps`/`FuseTIR` causes a performance regression for an FP16 Exp
chain followed by a broadcast Add on an RTX A6000. VM latency increases by 30%
with 8 Exp layers and 60% with 16 layers compared with disabling these passes.
Fusing the Exp chain while keeping the broadcast Add separate is faster than
either configuration.
### Workload
Inputs are `x: [4096, 1]` and `b: [4096, 4096]`, both FP16:
```python
z = x
for _ in range(depth):
z = exp(x - float16(0.5) * z)
y = b + z
```
Inputs are sampled from `[-0.25, 0.25]` with fixed seeds. This recurrence
keeps the Exp chain finite. The reference rounds to FP16 after each operator.
### Results
The three configurations have the same inputs and a single `[4096, 4096]`
output:
- **OFF:** skip `FuseOps` and `FuseTIR` in the CUDA pipeline.
- **PARTIAL:** place the Exp producer and broadcast Add in separate dataflow
blocks, then run the full pipeline.
- **ON:** run the full pipeline on a single dataflow block.
| Exp layers | ON/OFF | PARTIAL/OFF | ON/PARTIAL | Trials passing the
stability and 10% slowdown thresholds |
|---|---:|---:|---:|---:|
| 8 | 1.301768 | 0.793551 | 1.641148 | 5/5 |
| 16 | 1.595166 | 0.646946 | 2.444368 | 4/5 |
Ratios are latency ratios, summarized as the median of five per-trial
ratios. For 16 layers, ON/OFF was `1.584963, 1.595166, 1.602671, 1.615977,
1.569858`. One trial exceeded the 5% CV threshold; it is included in the table.
GPU clocks were not locked.
Each trial runs in a separate process. Timing uses
`vm.time_evaluator("invoke_stateful", ...)` with GPU-resident inputs:
`number=10`, `repeat=5`, `min_repeat_ms=50`, and 40 measurement blocks, giving
200 samples per configuration. Warmup is 75 calls followed by 2 seconds of
timed batches. The six configuration orders are cycled across blocks and offset
across trials. The measured time includes VM allocations; compilation,
transfers and correctness checks are outside the timer.
All three configurations passed reference and pairwise checks on three input
draws per trial, with zero maximum absolute error. The saved IR confirms that
ON fuses the Exp-to-Add edge and PARTIAL retains a fused producer followed by a
separate Add. Cubin hashes match across the five trials for each configuration.
### Environment
- TVM: `0.26.dev0+source.8f328e8`
- Source commit recorded in the build manifest:
`8f328e802cfe5e41fcc8f5c17e7582b1c28bfce4`
- GPU: NVIDIA RTX A6000, `sm_86`; driver 550.120
- CUDA toolkit: 12.4.131; cuda-bindings: 12.9.7
- LLVM: 15.0.7; host target: `llvm -mcpu=generic`
- Linux x86_64, Python 3.11.16, NumPy 2.4.6
### Reproduction
Save the code below as `repro_broadcast_regression.py` and run in a
CUDA-enabled TVM environment with NumPy installed:
```bash
python repro_broadcast_regression.py --depth 8 --device 0 --out repro_depth8
python repro_broadcast_regression.py --depth 16 --device 0 --out
repro_depth16
```
Use an idle GPU and a new output directory for each command. Each command
runs five trials and saves per-pass IR, correctness results, timing samples and
ratios. `summary.json` contains the median ratios and the number of stable
trials with at least a 10% slowdown.
The script checks for `3 * depth + 1`, `2`, and `1` calls to
`relax.call_tir` before VM lowering in OFF, PARTIAL, and ON. A failure here
means the generated graph needs to be inspected before comparing timings.
I extracted this script from the benchmark used for the measurements above;
I haven't rerun the standalone version on CUDA yet.
```python
"""Standalone extraction for TVM 0.26.dev0+source.8f328e8 / CUDA.
Run twice: --depth 8 --out repro_d8, then --depth 16 --out repro_d16.
Each command runs five separate processes. No benchmark-package imports.
"""
import argparse
import itertools
import json
import math
from pathlib import Path
import platform
import statistics
import subprocess
import sys
import time
VARIANTS = ("off", "partial", "on")
ORDERS = list(itertools.permutations(VARIANTS))
def reference(x, b, depth):
import numpy as np
z = x
for _ in range(depth):
z = (x - (z * np.float16(0.5)).astype("float16")).astype("float16")
z = np.exp(z.astype("float64")).astype("float16")
return (z + b).astype("float16")
def graph(depth, partial):
from tvm import relax as R
from tvm.relax import op
x = R.Var("x", R.TensorType([4096, 1], "float16"))
b = R.Var("b", R.TensorType([4096, 4096], "float16"))
bb = R.BlockBuilder()
def producer():
z = x
for _ in range(depth):
z = bb.emit(op.multiply(z, R.const(0.5, "float16")))
z = bb.emit(op.subtract(x, z))
z = bb.emit(op.exp(z))
return z
with bb.function("main", [x, b]):
if partial:
with bb.dataflow():
z = bb.emit_output(producer())
with bb.dataflow():
y = bb.emit_output(bb.emit(op.add(z, b)))
else:
with bb.dataflow():
y = bb.emit_output(bb.emit(op.add(producer(), b)))
bb.emit_func_output(y)
return bb.get()
def compile_vm(depth, mode, target, dev, out):
import tvm
from tvm import relax as R
from tvm.relax.backend.cuda import pipeline
out.mkdir()
mod = graph(depth, mode == "partial")
(out / "input.py").write_text(mod.script(), encoding="utf-8")
@tvm.instrument.pass_instrument
class Audit:
def run_before_pass(self, mod, info):
if mode == "off" and str(info.name).split(".")[-1] in
("FuseOps", "FuseTIR"):
raise RuntimeError("Unexpected fusion pass in OFF")
with target, tvm.transform.PassContext(opt_level=3,
instruments=[Audit()]):
steps = pipeline.library_dispatch_passes(target) +
pipeline.legalize_passes(target)
for i, p in enumerate(steps):
name = str(p.info.name).split(".")[-1]
if mode == "off" and name in ("FuseOps", "FuseTIR"):
continue
mod = p(mod)
(out / f"{i:02d}_{name}.py").write_text(mod.script(),
encoding="utf-8")
calls = []
def visit(expr):
if isinstance(expr, R.Call) and isinstance(expr.op, tvm.ir.Op):
if expr.op.name == "relax.call_tir":
calls.append(expr.args[0].name_hint)
R.analysis.post_order_visit(mod["main"].body, visit)
expected = {"off": 3 * depth + 1, "partial": 2, "on": 1}[mode]
if len(calls) != expected:
raise RuntimeError(f"{mode}: expected {expected} call_tir calls,
got {calls}; inspect IR")
for p in pipeline.dataflow_lower_passes(target) +
pipeline.finalize_passes(target):
mod = p(mod)
ex = R.build(mod, target=target, relax_pipeline=None,
tir_pipeline="default")
return R.VirtualMachine(ex, dev)
def stats(samples):
if len(samples) != 200 or any(not math.isfinite(x) or x <= 0 for x in
samples):
raise ValueError("Expected 200 finite positive samples")
return dict(median_s=statistics.median(samples),
cv=statistics.stdev(samples) / statistics.mean(samples),
samples_s=samples)
def trial(args):
import numpy as np
import tvm
dev = tvm.cuda(args.device)
if not dev.exist:
raise RuntimeError("CUDA device unavailable")
target = tvm.target.Target(
{"kind": "cuda", "arch": "sm_" +
str(dev.compute_version).replace(".", "")},
host={"kind": "llvm", "mcpu": "generic"})
order = ORDERS[args.trial % 6]
vms = {k: compile_vm(args.depth, k, target, dev, args.out / k) for k in
order}
errors = []
for seed in range(3):
rng = np.random.default_rng(seed)
x = rng.uniform(-.25, .25, (4096, 1)).astype("float16")
b = rng.uniform(-.25, .25, (4096, 4096)).astype("float16")
arrays = [tvm.runtime.tensor(a, device=dev) for a in (x, b)]
if seed == 0:
timed_inputs = arrays
expected, outputs = reference(x, b, args.depth), {}
for k, vm in vms.items():
vm.set_input("main", *arrays)
vm.invoke_stateful("main")
dev.sync()
value = vm.get_outputs("main")
while not hasattr(value, "numpy"):
if len(value) != 1:
raise AssertionError("Expected exactly one output")
value = value[0]
outputs[k] = value.numpy()
pairs = [(k, outputs[k], expected) for k in VARIANTS]
pairs += [(f"{a}/{b}", outputs[a], outputs[b])
for a, b in itertools.combinations(VARIANTS, 2)]
for label, actual, wanted in pairs:
if actual.shape != wanted.shape or actual.dtype != wanted.dtype:
raise AssertionError(f"{label}: output shape/dtype mismatch")
if not np.isfinite(actual).all() or not
np.isfinite(wanted).all():
raise AssertionError(f"{label}: non-finite output")
np.testing.assert_allclose(actual, wanted, rtol=.02, atol=.002)
errors.append(dict(seed=seed, comparison=label,
max_abs_error=float(np.max(np.abs(
actual.astype("float32") -
wanted.astype("float32"))))))
timers, samples = {}, {k: [] for k in VARIANTS}
for k in order:
vm = vms[k]
vm.set_input("main", *timed_inputs)
timers[k] = vm.time_evaluator("invoke_stateful", dev, number=10,
repeat=5, min_repeat_ms=50)
for _ in range(75):
vm.invoke_stateful("main")
dev.sync()
deadline = time.monotonic() + 2.0
while time.monotonic() < deadline:
dev.sync()
timers[k]("main")
for block in range(40):
for k in ORDERS[(args.trial + block) % 6]:
dev.sync()
samples[k].extend(float(v) for v in timers[k]("main").results)
measured = {k: stats(v) for k, v in samples.items()}
ratios = {f"{a}/{b}": measured[a]["median_s"] / measured[b]["median_s"]
for a, b in (("on", "off"), ("partial", "off"), ("on",
"partial"))}
result = dict(depth=args.depth, trial=args.trial, tvm=tvm.__version__,
python=sys.version, platform=platform.platform(),
target=str(target),
gpu=str(dev.device_name), correctness=errors,
measured=measured, ratios=ratios,
on_off_stable=max(measured[k]["cv"] for k in ("on",
"off")) <= .05)
(args.out / "result.json").write_text(json.dumps(result, indent=2),
encoding="utf-8")
print(json.dumps({"trial": args.trial, "ratios": ratios,
"CV": {k: v["cv"] for k, v in measured.items()}}),
flush=True)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--depth", type=int, choices=(8, 16), required=True)
parser.add_argument("--device", type=int, default=0)
parser.add_argument("--out", type=Path, required=True)
parser.add_argument("--trial", type=int, help=argparse.SUPPRESS)
args = parser.parse_args()
if args.out.exists():
parser.error("--out must be a new directory")
args.out.mkdir(parents=True)
if args.trial is not None:
trial(args)
return
results = []
for i in range(5):
dest = args.out / f"trial_{i}"
print(f"Depth {args.depth}, trial {i + 1}/5", flush=True)
subprocess.run([sys.executable, str(Path(__file__).resolve()),
"--depth", str(args.depth),
"--device", str(args.device), "--out", str(dest),
"--trial", str(i)], check=True)
results.append(json.loads((dest /
"result.json").read_text(encoding="utf-8")))
summary = dict(median_ratios={k: statistics.median(r["ratios"][k] for r
in results)
for k in results[0]["ratios"]},
stable_on_off_slow_trials=sum(r["on_off_stable"] and
r["ratios"]["on/off"] >= 1.10
for r in results),
total_trials=5)
(args.out / "summary.json").write_text(json.dumps(summary, indent=2),
encoding="utf-8")
print(json.dumps(summary, indent=2))
if __name__ == "__main__":
main()
```
### Where the extra computation appears
After `FuseTIR`, the fused PrimFunc still has `[4096, 1]` intermediate
buffers and small producer loops, followed by the broadcast Add. After
`ApplyDefaultSchedule`, the Exp chain is inlined into the `[4096, 4096]` output
expression. The scheduling trace selects DLight `Fallback`.
For 16 layers:
- PARTIAL schedules the producer with 4 blocks of 1024 threads, then runs a
separate Add kernel.
- ON schedules the output with 16384 blocks of 1024 threads and embeds the
entire Exp chain in that kernel. Its SASS contains 16 `MUFU.EX2` instruction
sites.
This points to repeated producer evaluation along the broadcast dimension.
The ON cubin uses 13 registers; PARTIAL uses 11 for the producer and 10 for
Add. All report zero stack, local and shared memory, and the disassembly
contains no `LDL`/`STL` instructions. These results do not suggest register
spilling. Per-kernel timings and SFU utilization still need profiling.
Relevant code at the recorded commit:
-
[Fallback.apply](https://github.com/apache/tvm/blob/8f328e802cfe5e41fcc8f5c17e7582b1c28bfce4/python/tvm/s_tir/dlight/gpu/fallback.py#L66)
calls `try_inline` before scheduling the remaining block.
-
[try_inline](https://github.com/apache/tvm/blob/8f328e802cfe5e41fcc8f5c17e7582b1c28bfce4/python/tvm/s_tir/dlight/base/common_schedules.py#L43)
tries `compute_inline` and `reverse_compute_inline`.
-
[GraphPartitioner::RunFuse](https://github.com/apache/tvm/blob/8f328e802cfe5e41fcc8f5c17e7582b1c28bfce4/src/relax/analysis/graph_partitioner.cc#L334)
makes the earlier fusion decision.
Keeping the producer materialized reduces its output to an 8 KiB buffer and
gives the fastest result in this test. Could the fusion or inlining decision
account for producer cost and broadcast reuse? The PARTIAL configuration
provides a working comparison for that change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
