[examples] Better loop tool formatting.
@@ -38,16 +38,18 @@ def flops_after_steps(env, num_steps):
def run_one_sweep(
device: str,
k: int,
- vectorize: int = 4,
+ vectorize: int = 1,
linear: bool = False,
logdir: Optional[Path] = None,
):
"""Run a single sweep."""
logdir = logdir or create_user_logs_dir("loop_tool_sweep")
- logfile = logdir / f"k{k}-v{vectorize}-{'linear' if linear else 'log'}-{device}.txt"
+ logfile = logdir / f"k{k}-v{vectorize}-{device}-{'linear' if linear else 'log'}.txt"
+ print("Logging results to", logfile)
+ print()
+ print("Device", "K", "Inner", "Vec.", "FLOPS", sep="\t")
with open(logfile, "w") as f:
print("device", "k", "inner", "vectorize", "flops", sep=",", file=f)
- print("Logging results to", logfile)
def log(k, inner, vectorize, flops):
print(device.upper(), k, inner, vectorize, flops, sep="\t", flush=True)
@@ -91,7 +93,7 @@ def run_one_sweep(
def sweep(
device: List[str] = ["cuda"],
k: List[int] = [512, 1024, 2048, 4096, 8192],
- vectorize: List[int] = [4],
+ vectorize: List[int] = [1],
linear: List[bool] = [False],
logdir: Optional[Path] = None,
):
