Fusing a 2x2 MaxPool with its elementwise Add consumer should retain an efficient schedule for the small reduction. In this case, one thread can compute an output using four input values and then apply the Add.
The graph was extracted from the official pass-test corpus for AnnotateTIROpPattern/FuseOps and adapted to the TIRx API. Its input is float32 [1,16,64,64]. Nine Add branches are concatenated in two levels, producing [1,144,64,64], followed by a 2x2 MaxPool with stride 2 and two Add operations. Both Add outputs are returned, each with shape [1,144,32,32].
These are medians of five per-trial median latencies. The five ON/OFF ratios were 5.54737, 5.50958, 5.56865, 5.52919, 5.51838. All five trials passed the 10% slowdown threshold with CV below 5%; the largest CV across all four paths was 0.73%. All paths matched an independent NumPy reference with zero maximum absolute error on three input draws per trial.
ON also uses shared memory for the reduction and its consumer. Replacing only this fused function with the SERIAL schedule reduces VM latency by 95.9%, without changing the other functions or outputs. The actual cubin reports 16 registers and no stack/local memory for both ON and SERIAL; shared memory drops from 132 bytes to zero.
This points to the default scheduling choice for a small reduction with an elementwise consumer. A schedule that parallelizes across outputs and performs the small reduction within each thread avoids the slowdown. The SERIAL intervention changes several scheduling decisions together, so it does not isolate their individual costs. The timings above are whole-graph VM timings.
The script preserves the original TIRx graph and compares the four paths above. PARTIAL places a dataflow boundary immediately after Pool. SERIAL replaces only the fused pool function before default scheduling.
import ast
import copy
import hashlib
import itertools
import json
from pathlib import Path
import statistics
import sys
import time
import numpy as np
import tvm
from tvm import relax
from tvm.relax.backend.cuda import pipeline
from tvm.script import ir as I, relax as R, tirx as T
SOURCE = r'''
@I.ir_module
class Module:
@T.prim_func(private=True, s_tir=True)
def add(x: T.Buffer((T.int64(1), T.int64(16), T.int64(64), T.int64(64)), 'float32'), B: T.Buffer((), 'float32'), T_add: T.Buffer((T.int64(1), T.int64(16), T.int64(64), T.int64(64)), 'float32')):
T.func_attr({'tirx.noalias': True})
for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(16), T.int64(64), T.int64(64)):
with T.sblock('T_add'):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap('SSSS', [ax0, ax1, ax2, ax3])
T.reads(x[v_ax0, v_ax1, v_ax2, v_ax3], B[()])
T.writes(T_add[v_ax0, v_ax1, v_ax2, v_ax3])
T_add[v_ax0, v_ax1, v_ax2, v_ax3] = x[v_ax0, v_ax1, v_ax2, v_ax3] + B[()]
@T.prim_func(private=True, s_tir=True)
def add1(lv3: T.Buffer((T.int64(1), T.int64(48), T.int64(64), T.int64(64)), 'float32'), B: T.Buffer((), 'float32'), T_add: T.Buffer((T.int64(1), T.int64(48), T.int64(64), T.int64(64)), 'float32')):
T.func_attr({'tirx.noalias': True})
for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(48), T.int64(64), T.int64(64)):
with T.sblock('T_add'):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap('SSSS', [ax0, ax1, ax2, ax3])
T.reads(lv3[v_ax0, v_ax1, v_ax2, v_ax3], B[()])
T.writes(T_add[v_ax0, v_ax1, v_ax2, v_ax3])
T_add[v_ax0, v_ax1, v_ax2, v_ax3] = lv3[v_ax0, v_ax1, v_ax2, v_ax3] + B[()]
@T.prim_func(private=True, s_tir=True)
def add2(lv16: T.Buffer((T.int64(1), T.int64(144), T.int64(32), T.int64(32)), 'float32'), B: T.Buffer((), 'float32'), T_add: T.Buffer((T.int64(1), T.int64(144), T.int64(32), T.int64(32)), 'float32')):
T.func_attr({'tirx.noalias': True})
for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(144), T.int64(32), T.int64(32)):
with T.sblock('T_add'):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap('SSSS', [ax0, ax1, ax2, ax3])
T.reads(lv16[v_ax0, v_ax1, v_ax2, v_ax3], B[()])
T.writes(T_add[v_ax0, v_ax1, v_ax2, v_ax3])
T_add[v_ax0, v_ax1, v_ax2, v_ax3] = lv16[v_ax0, v_ax1, v_ax2, v_ax3] + B[()]
@T.prim_func(private=True, s_tir=True)
def concatenate(lv: T.Buffer((T.int64(1), T.int64(16), T.int64(64), T.int64(64)), 'float32'), lv1: T.Buffer((T.int64(1), T.int64(16), T.int64(64), T.int64(64)), 'float32'), lv2: T.Buffer((T.int64(1), T.int64(16), T.int64(64), T.int64(64)), 'float32'), T_concat: T.Buffer((T.int64(1), T.int64(48), T.int64(64), T.int64(64)), 'float32')):
T.func_attr({'tirx.noalias': True})
for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(48), T.int64(64), T.int64(64)):
with T.sblock('T_concat'):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap('SSSS', [ax0, ax1, ax2, ax3])
T.reads(lv2[v_ax0, v_ax1 - T.int64(32), v_ax2, v_ax3], lv1[v_ax0, v_ax1 - T.int64(16), v_ax2, v_ax3], lv[v_ax0, v_ax1, v_ax2, v_ax3])
T.writes(T_concat[v_ax0, v_ax1, v_ax2, v_ax3])
T_concat[v_ax0, v_ax1, v_ax2, v_ax3] = T.if_then_else(T.int64(32) <= v_ax1, lv2[v_ax0, v_ax1 - T.int64(32), v_ax2, v_ax3], T.if_then_else(T.int64(16) <= v_ax1, lv1[v_ax0, v_ax1 - T.int64(16), v_ax2, v_ax3], lv[v_ax0, v_ax1, v_ax2, v_ax3]))
@T.prim_func(private=True, s_tir=True)
def concatenate1(lv4: T.Buffer((T.int64(1), T.int64(48), T.int64(64), T.int64(64)), 'float32'), lv9: T.Buffer((T.int64(1), T.int64(48), T.int64(64), T.int64(64)), 'float32'), lv14: T.Buffer((T.int64(1), T.int64(48), T.int64(64), T.int64(64)), 'float32'), T_concat: T.Buffer((T.int64(1), T.int64(144), T.int64(64), T.int64(64)), 'float32')):
T.func_attr({'tirx.noalias': True})
for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(144), T.int64(64), T.int64(64)):
with T.sblock('T_concat'):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap('SSSS', [ax0, ax1, ax2, ax3])
T.reads(lv14[v_ax0, v_ax1 - T.int64(96), v_ax2, v_ax3], lv9[v_ax0, v_ax1 - T.int64(48), v_ax2, v_ax3], lv4[v_ax0, v_ax1, v_ax2, v_ax3])
T.writes(T_concat[v_ax0, v_ax1, v_ax2, v_ax3])
T_concat[v_ax0, v_ax1, v_ax2, v_ax3] = T.if_then_else(T.int64(96) <= v_ax1, lv14[v_ax0, v_ax1 - T.int64(96), v_ax2, v_ax3], T.if_then_else(T.int64(48) <= v_ax1, lv9[v_ax0, v_ax1 - T.int64(48), v_ax2, v_ax3], lv4[v_ax0, v_ax1, v_ax2, v_ax3]))
@T.prim_func(private=True, s_tir=True)
def pool2d(lv15: T.Buffer((T.int64(1), T.int64(144), T.int64(64), T.int64(64)), 'float32'), pool_max: T.Buffer((T.int64(1), T.int64(144), T.int64(32), T.int64(32)), 'float32')):
T.func_attr({'tirx.noalias': True})
for ax0, ax1, ax2, ax3, rv0, rv1 in T.grid(T.int64(1), T.int64(144), T.int64(32), T.int64(32), T.int64(2), T.int64(2)):
with T.sblock('pool_max'):
v_ax0, v_ax1, v_ax2, v_ax3, v_rv0, v_rv1 = T.axis.remap('SSSSRR', [ax0, ax1, ax2, ax3, rv0, rv1])
T.reads(lv15[v_ax0, v_ax1, v_ax2 * T.int64(2) + v_rv0, v_ax3 * T.int64(2) + v_rv1])
T.writes(pool_max[v_ax0, v_ax1, v_ax2, v_ax3])
T.sblock_attr({'schedule_rule': 'meta_schedule.pool_max'})
with T.init():
pool_max[v_ax0, v_ax1, v_ax2, v_ax3] = T.float32(-3.4028234663852886e+38)
pool_max[v_ax0, v_ax1, v_ax2, v_ax3] = T.max(pool_max[v_ax0, v_ax1, v_ax2, v_ax3], lv15[v_ax0, v_ax1, v_ax2 * T.int64(2) + v_rv0, v_ax3 * T.int64(2) + v_rv1])
@R.function
def main(x: R.Tensor((1, 16, 64, 64), dtype='float32')) -> R.Tuple(R.Tensor((1, 144, 32, 32), dtype='float32'), R.Tensor((1, 144, 32, 32), dtype='float32')):
cls = Module
with R.dataflow():
lv = R.call_tir(cls.add, (x, R.const(1, 'float32')), out_ty=R.Tensor((1, 16, 64, 64), dtype='float32'))
lv1 = R.call_tir(cls.add, (x, R.const(1, 'float32')), out_ty=R.Tensor((1, 16, 64, 64), dtype='float32'))
lv2 = R.call_tir(cls.add, (x, R.const(1, 'float32')), out_ty=R.Tensor((1, 16, 64, 64), dtype='float32'))
lv3 = R.call_tir(cls.concatenate, (lv, lv1, lv2), out_ty=R.Tensor((1, 48, 64, 64), dtype='float32'))
lv4 = R.call_tir(cls.add1, (lv3, R.const(1, 'float32')), out_ty=R.Tensor((1, 48, 64, 64), dtype='float32'))
lv5 = R.call_tir(cls.add, (x, R.const(1, 'float32')), out_ty=R.Tensor((1, 16, 64, 64), dtype='float32'))
lv6 = R.call_tir(cls.add, (x, R.const(1, 'float32')), out_ty=R.Tensor((1, 16, 64, 64), dtype='float32'))
lv7 = R.call_tir(cls.add, (x, R.const(1, 'float32')), out_ty=R.Tensor((1, 16, 64, 64), dtype='float32'))
lv8 = R.call_tir(cls.concatenate, (lv5, lv6, lv7), out_ty=R.Tensor((1, 48, 64, 64), dtype='float32'))
lv9 = R.call_tir(cls.add1, (lv8, R.const(1, 'float32')), out_ty=R.Tensor((1, 48, 64, 64), dtype='float32'))
lv10 = R.call_tir(cls.add, (x, R.const(1, 'float32')), out_ty=R.Tensor((1, 16, 64, 64), dtype='float32'))
lv11 = R.call_tir(cls.add, (x, R.const(1, 'float32')), out_ty=R.Tensor((1, 16, 64, 64), dtype='float32'))
lv12 = R.call_tir(cls.add, (x, R.const(1, 'float32')), out_ty=R.Tensor((1, 16, 64, 64), dtype='float32'))
lv13 = R.call_tir(cls.concatenate, (lv10, lv11, lv12), out_ty=R.Tensor((1, 48, 64, 64), dtype='float32'))
lv14 = R.call_tir(cls.add1, (lv13, R.const(1, 'float32')), out_ty=R.Tensor((1, 48, 64, 64), dtype='float32'))
lv15 = R.call_tir(cls.concatenate1, (lv4, lv9, lv14), out_ty=R.Tensor((1, 144, 64, 64), dtype='float32'))
lv16 = R.call_tir(cls.pool2d, (lv15,), out_ty=R.Tensor((1, 144, 32, 32), dtype='float32'))
lv17 = R.call_tir(cls.add2, (lv16, R.const(1, 'float32')), out_ty=R.Tensor((1, 144, 32, 32), dtype='float32'))
lv18 = R.call_tir(cls.add2, (lv17, R.const(1, 'float32')), out_ty=R.Tensor((1, 144, 32, 32), dtype='float32'))
gv: R.Tuple(R.Tensor((1, 144, 32, 32), dtype='float32'), R.Tensor((1, 144, 32, 32), dtype='float32')) = (lv17, lv18)
R.output(gv)
return gv
'''
def shape(c):
return (1, c["channels"], c["height"], c["width"])
def pool_shape(c):
c0 = c["channels"] * (c["branches"]**2 if c["prelude"] else 1)
return (1,c0,(c["height"]-c["window"])//c["stride"]+1,(c["width"]-c["window"])//c["stride"]+1)
def inputs(c, seed):
return np.random.default_rng(seed).uniform(-.5,.5,shape(c)).astype("float32")
def reference(c, x):
y = x
one = np.float32(1)
if c["prelude"]:
groups = []
for _ in range(c["branches"]):
leaves = [(x+one).astype("float32") for _ in range(c["branches"])]
groups.append((np.concatenate(leaves,axis=1)+one).astype("float32"))
y = np.concatenate(groups,axis=1)
_,channels,oh,ow = pool_shape(c)
pooled = np.full((1,channels,oh,ow), np.finfo(np.float32).min, dtype="float32")
k,s = c["window"],c["stride"]
for rh in range(k):
for rw in range(k):
pooled = np.maximum(pooled,y[:,:,rh:rh+oh*s:s,rw:rw+ow*s:s])
outputs = []
for _ in range(c["tail"]):
pooled = (pooled+one).astype("float32")
outputs.append(pooled)
return (outputs[0], outputs[-1]) if c["shared_output"] else outputs[-1]
def tir_shape(dims):
return "(" + ", ".join(f"T.int64({v})" for v in dims) + ("," if len(dims)==1 else "") + ")"
def partial_source(text):
"""Split dataflow immediately after pool; leave main inputs/outputs intact."""
tree = ast.parse(text)
cls = next(n for n in tree.body if isinstance(n,ast.ClassDef))
main = next(n for n in cls.body if isinstance(n,ast.FunctionDef) and n.name=="main")
blocks = [n for n in main.body if isinstance(n,ast.With)]
if len(blocks)!=1:
raise ValueError("Expected one dataflow block")
block = blocks[0]
hits = [i for i,n in enumerate(block.body) if isinstance(n,(ast.Assign,ast.AnnAssign))
and isinstance(n.value,ast.Call) and ast.unparse(n.value.func)=="R.call_tir"
and ast.unparse(n.value.args[0])=="cls.pool2d"]
if len(hits)!=1:
raise ValueError("Expected unique pool call")
index = hits[0]; stmt = block.body[index]
var = stmt.targets[0] if isinstance(stmt,ast.Assign) else stmt.target
first = copy.deepcopy(block); second = copy.deepcopy(block)
first.body = block.body[:index+1]+[ast.parse(f"R.output({var.id})").body[0]]
second.body = block.body[index+1:]
pos = main.body.index(block); main.body[pos:pos+1] = [first,second]
return ast.unparse(ast.fix_missing_locations(tree))+"\n"
def serial_source(c, biases):
"""A diagnostic schedule: one thread computes one small pool and epilogue."""
_,channels,oh,ow = pool_shape(c)
h,w,k,s = c["height"],c["width"],c["window"],c["stride"]
total = channels*oh*ow
args = ", ".join(f"B{i}: T.Buffer((), 'float32')" for i in range(biases))
additions = "\n".join(f" acc[0] = acc[0] + B{i}[()]" for i in range(biases))
return f'''@I.ir_module
class Replacement:
@T.prim_func(private=True, s_tir=True)
def replacement(A: T.Buffer({tir_shape((1,channels,h,w))}, "float32"), {args}, O: T.Buffer({tir_shape((1,channels,oh,ow))}, "float32")):
T.func_attr({{"tirx.noalias": True, "tirx.is_scheduled": True}})
for bx in T.thread_binding(T.int64({(total+255)//256}), thread="blockIdx.x"):
for tx in T.thread_binding(T.int64(256), thread="threadIdx.x"):
with T.sblock("serial_pool_epilogue"):
v = T.axis.spatial(T.int64({total}), bx*T.int64(256)+tx)
T.where(bx*T.int64(256)+tx < T.int64({total}))
acc = T.sblock_alloc_buffer((1,), "float32", scope="local")
acc[0] = T.float32(-3.4028234663852886e+38)
for rh,rw in T.grid(T.int64({k}),T.int64({k})):
acc[0] = T.max(acc[0], A[0,v//T.int64({oh*ow}),(v//T.int64({ow})%T.int64({oh}))*T.int64({s})+rh,(v%T.int64({ow}))*T.int64({s})+rw])
{additions}
O[0,v//T.int64({oh*ow}),v//T.int64({ow})%T.int64({oh}),v%T.int64({ow})] = acc[0]
'''
def parameter_buffer_types(fn, buffer_type):
"""TIRx parameters carry BufferType in Var.ty, not a PrimFunc buffer_map."""
types = []
for index, param in enumerate(fn.params):
ty = getattr(param, "ty", None)
if not isinstance(ty, buffer_type):
raise ValueError(f"Pool parameter {index} is not a BufferType; refusing intervention")
types.append(ty)
return types
def intervention(c, out):
def apply(mod):
import tvm
from tvm import tirx
from tvm.tirx.buffer import BufferType
from tvm.script import ir as I, tirx as T
found = [(gv,fn) for gv,fn in mod.functions.items() if isinstance(fn,tirx.PrimFunc)
and "pool2d" in gv.name_hint and "add2" in gv.name_hint]
if len(found)!=1:
raise ValueError("Expected exactly one fused pool epilogue for intervention")
gv,fn = found[0]
params = list(fn.params)
buffer_types = parameter_buffer_types(fn, BufferType)
dims = [tuple(int(x) for x in ty.shape) for ty in buffer_types]
dtypes = [str(tvm.DataType(ty.dtype)) for ty in buffer_types]
expected_in = (1,pool_shape(c)[1],c["height"],c["width"])
if len(params)<3 or dims[0]!=expected_in or dims[-1]!=pool_shape(c) or not all(d==() for d in dims[1:-1]):
raise ValueError("Unrecognized fused pool signature; refusing intervention")
if any(dtype!="float32" for dtype in dtypes):
raise ValueError("Only float32 intervention validated")
text = serial_source(c,len(params)-2)
(out/"serial_kernel.py").write_text(text,encoding="utf-8")
replacement = tvm.script.from_source(text, extra_vars=dict(I=I,T=T),s_tir=True)["replacement"]
replacement_types = parameter_buffer_types(replacement, BufferType)
if ([tuple(int(x) for x in ty.shape) for ty in replacement_types] != dims
or [str(tvm.DataType(ty.dtype)) for ty in replacement_types] != dtypes):
raise ValueError("Replacement parameter shapes/dtypes differ from original")
record = dict(replaced_function=gv.name_hint, biases=len(params)-2,
parameter_api="tirx.Var.ty / BufferType", parameter_shapes=dims, parameter_dtypes=dtypes,
policy="diagnostic serial reduction per output; not an upstream compiler fix",
before_sha256=hashlib.sha256(fn.script().encode()).hexdigest(),
unchanged_functions={key.name_hint:hashlib.sha256(value.script().encode()).hexdigest()
for key,value in mod.functions.items() if key!=gv})
mod.update_func(gv,replacement)
record["other_functions_unchanged"] = all(hashlib.sha256(mod[key].script().encode()).hexdigest()==value
for key,value in record["unchanged_functions"].items())
write_json(out/"intervention.json",record)
return mod
return apply
CASE = dict(channels=16, height=64, width=64, window=2, stride=2,
branches=3, prelude=True, tail=2, shared_output=True)
VARIANTS = ("off", "partial", "on", "serial")
def write_json(path, data):
path.write_text(json.dumps(data, indent=2), encoding="utf-8")
def build_variant(key, target, directory):
directory.mkdir(parents=True, exist_ok=True)
text = partial_source(SOURCE) if key == "partial" else SOURCE
mod = tvm.script.from_source(text, extra_vars=dict(I=I, R=R, T=T), s_tir=True)
with target, tvm.transform.PassContext(opt_level=3):
steps = pipeline.library_dispatch_passes(target) + pipeline.legalize_passes(target)
for index, transform in enumerate(steps):
name = str(transform.info.name).split(".")[-1]
if key == "off" and name in ("FuseOps", "FuseTIR"):
continue
if name == "ApplyDefaultSchedule" and key == "serial":
mod = intervention(CASE, directory)(mod)
mod = transform(mod)
if name in ("FuseOps", "FuseTIR", "ApplyDefaultSchedule"):
(directory / f"{index:02d}_{name}.py").write_text(mod.script())
for transform in pipeline.dataflow_lower_passes(target) + pipeline.finalize_passes(target):
mod = transform(mod)
return relax.build(mod, target=target, relax_pipeline=None, tir_pipeline="default")
def main():
trial = int(sys.argv[1]) if len(sys.argv) > 1 else 0
out = Path(f"pool_repro_trial_{trial}")
out.mkdir(parents=True, exist_ok=True)
dev = tvm.cuda(0)
assert dev.exist
arch = "sm_" + str(dev.compute_version).replace(".", "")
target = tvm.target.Target({"kind": "cuda", "arch": arch}, host="llvm -mcpu=generic")
print(tvm.__version__, target, flush=True)
order = VARIANTS if trial % 2 == 0 else tuple(reversed(VARIANTS))
vms = {key: relax.VirtualMachine(build_variant(key, target, out/key), dev) for key in order}
for seed in range(3):
x = inputs(CASE, seed)
expected = reference(CASE, x)
for key, vm in vms.items():
vm.set_input("main", tvm.runtime.tensor(x, device=dev))
vm.invoke_stateful("main")
dev.sync()
actual = vm.get_outputs("main")
assert len(actual) == len(expected)
for got, want in zip(actual, expected):
got = got.numpy()
assert got.dtype == want.dtype and got.shape == want.shape
np.testing.assert_allclose(got, want, rtol=2e-4, atol=2e-5)
print("All four paths passed the independent NumPy reference.", flush=True)
for vm in vms.values():
vm.set_input("main", tvm.runtime.tensor(inputs(CASE, 0), device=dev))
timers = {k: vm.time_evaluator("invoke_stateful", dev, number=10,
repeat=5, min_repeat_ms=50) for k, vm in vms.items()}
permutations = list(itertools.permutations(VARIANTS))
offset = trial * 40
for key in permutations[offset % len(permutations)]:
for _ in range(75):
vms[key].invoke_stateful("main")
dev.sync()
deadline = time.monotonic() + 2
while time.monotonic() < deadline:
timers[key]("main")
samples = {key: [] for key in VARIANTS}
for block in range(40):
for key in permutations[(offset + block) % len(permutations)]:
dev.sync()
samples[key].extend(float(x) for x in timers[key]("main").results)
results = {k: dict(median_us=statistics.median(xs)*1e6,
cv=statistics.stdev(xs)/statistics.mean(xs), samples_s=xs)
for k, xs in samples.items()}
for key in VARIANTS:
print(key, "median_us=", results[key]["median_us"], "CV=", results[key]["cv"])
for a, b in (("on", "off"), ("on", "partial"), ("serial", "on")):
print(a+"/"+b, results[a]["median_us"]/results[b]["median_us"])
write_json(out/"timings.json", results)
if __name__ == "__main__":
main()
Expected behavior
Fusing a 2x2 MaxPool with its elementwise Add consumer should retain an efficient schedule for the small reduction. In this case, one thread can compute an output using four input values and then apply the Add.
Actual behavior
Enabling
FuseOpsandFuseTIRincreases end-to-end VM latency from 45.68 us to 251.70 us on an RTX A6000. Keeping the upstream Add/Concat operations fused while separating Pool from Add takes 13.00 us.The graph was extracted from the official pass-test corpus for AnnotateTIROpPattern/FuseOps and adapted to the TIRx API. Its input is float32
[1,16,64,64]. Nine Add branches are concatenated in two levels, producing[1,144,64,64], followed by a 2x2 MaxPool with stride 2 and two Add operations. Both Add outputs are returned, each with shape[1,144,32,32].These are medians of five per-trial median latencies. The five ON/OFF ratios were
5.54737, 5.50958, 5.56865, 5.52919, 5.51838. All five trials passed the 10% slowdown threshold with CV below 5%; the largest CV across all four paths was 0.73%. All paths matched an independent NumPy reference with zero maximum absolute error on three input draws per trial.The saved scheduling trace shows that
GeneralReductionacceptsfused_pool2d_add2. The resulting launch configuration differs substantially:ON also uses shared memory for the reduction and its consumer. Replacing only this fused function with the SERIAL schedule reduces VM latency by 95.9%, without changing the other functions or outputs. The actual cubin reports 16 registers and no stack/local memory for both ON and SERIAL; shared memory drops from 132 bytes to zero.
This points to the default scheduling choice for a small reduction with an elementwise consumer. A schedule that parallelizes across outputs and performs the small reduction within each thread avoids the slowdown. The SERIAL intervention changes several scheduling decisions together, so it does not isolate their individual costs. The timings above are whole-graph VM timings.
Environment
0.26.dev0+source.8f328e88f328e802cfe5e41fcc8f5c17e7582b1c28bfce4sm_86550.12012.4.131llvm -mcpu=generic6.8.0-52-generic, glibc 2.35USE_CUDA=ON,USE_LLVM=ON,USE_CUBLAS=OFF,USE_CUDNN=OFF,USE_CUTLASS=OFFSteps to reproduce
Save the following as
repro_pool_fusion.pyand run it in a TVM CUDA environment at the commit above:The script preserves the original TIRx graph and compares the four paths above. PARTIAL places a dataflow boundary immediately after Pool. SERIAL replaces only the fused pool function before default scheduling.
Inputs are sampled uniformly from
[-0.5, 0.5]with seeds 0, 1 and 2 for correctness; seed 0 is used for timing. Inputs are copied to the GPU before measurement. Timing usesinvoke_stateful, 75 warmup calls followed by two seconds of timed warmup, and 40 interleaved blocks of five repeats (number=10,min_repeat_ms=50). Compilation and transfers are outside the timer. Each process saves the scheduling IR and all 200 samples per path.Triage