pinpinpoo opened a new issue, #20401:
URL: https://github.com/apache/tvm/issues/20401
### 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 `FuseOps` and `FuseTIR` increases 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]`.
| Configuration | VM latency, us |
|---|---:|
| OFF: skip FuseOps and FuseTIR | 45.68 |
| PARTIAL: fuse the prefix, keep Pool and Add separate | 13.00 |
| ON: default fusion and scheduling | 251.70 |
| SERIAL: same fused graph, replace only the fused pool function's schedule
| 10.30 |
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 `GeneralReduction` accepts
`fused_pool2d_add2`. The resulting launch configuration differs substantially:
| Pool schedule | Blocks | Threads per block | Output mapping |
|---|---:|---:|---|
| OFF / PARTIAL | 144 | 1024 | One output per thread, serial four-element
reduction |
| ON | 147456 | 256 | One output per block; only four threads load reduction
inputs |
| SERIAL | 576 | 256 | One output per thread, serial reduction followed by
Add |
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
- TVM: `0.26.dev0+source.8f328e8`
- Commit: `8f328e802cfe5e41fcc8f5c17e7582b1c28bfce4`
- GPU: NVIDIA RTX A6000, `sm_86`
- NVIDIA driver: `550.120`
- CUDA compiler: `12.4.131`
- LLVM: 15; host target `llvm -mcpu=generic`
- OS: Linux x86_64, kernel `6.8.0-52-generic`, glibc 2.35
- Python: 3.11.16; NumPy: 2.4.6
- `USE_CUDA=ON`, `USE_LLVM=ON`, `USE_CUBLAS=OFF`, `USE_CUDNN=OFF`,
`USE_CUTLASS=OFF`
### Steps to reproduce
Save the following as `repro_pool_fusion.py` and run it in a TVM CUDA
environment at the commit above:
```bash
for trial in 0 1 2 3 4; do
python repro_pool_fusion.py "$trial"
done
```
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 uses `invoke_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.
```python
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()
```
### Triage
- needs-triage
- backend:cuda
--
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]
