663 lines
24 KiB
Python
663 lines
24 KiB
Python
|
import json
|
||
|
|
import math
|
||
|
|
import os
|
||
|
|
import re
|
||
|
|
from typing import Dict, List, Optional, Set
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.utils.benchmark as benchmark
|
||
|
|
from torch._C._profiler import (
|
||
|
|
_EventType,
|
||
|
|
_ExtraFields_PyCall,
|
||
|
|
_ExtraFields_PyCCall,
|
||
|
|
_ExtraFields_TorchOp,
|
||
|
|
_ProfilerEvent,
|
||
|
|
)
|
||
|
|
from torch.profiler import profile
|
||
|
|
from torch.profiler._utils import index_of_first_match, traverse_bfs, traverse_dfs
|
||
|
|
|
||
|
|
|
||
|
|
class Pattern:
|
||
|
|
"""
|
||
|
|
Base class for all patterns, subclass this class and implement match()
|
||
|
|
to define custom patterns.
|
||
|
|
|
||
|
|
In subclass, define description and skip property.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, prof: profile, should_benchmark: bool = False):
|
||
|
|
self.prof = prof
|
||
|
|
self.should_benchmark = should_benchmark
|
||
|
|
self.name = "Please specify a name for pattern"
|
||
|
|
self.description = "Please specify a description for pattern"
|
||
|
|
self.url = ""
|
||
|
|
assert prof.profiler is not None and prof.profiler.kineto_results is not None
|
||
|
|
self.event_tree = prof.profiler.kineto_results.experimental_event_tree()
|
||
|
|
self.tid_root: Dict[int, List[_ProfilerEvent]] = {}
|
||
|
|
for event in self.event_tree:
|
||
|
|
self.tid_root.setdefault(event.start_tid, []).append(event)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def skip(self):
|
||
|
|
return False
|
||
|
|
|
||
|
|
def report(self, event: _ProfilerEvent):
|
||
|
|
msg = (
|
||
|
|
f"{self.description}\n[Source Code Location] {source_code_location(event)}"
|
||
|
|
)
|
||
|
|
return msg
|
||
|
|
|
||
|
|
def eventTreeTraversal(self):
|
||
|
|
"""
|
||
|
|
Traverse the event tree and yield all events.
|
||
|
|
Override this method in subclass to customize the traversal.
|
||
|
|
"""
|
||
|
|
yield from traverse_dfs(self.event_tree)
|
||
|
|
|
||
|
|
def summary(self, events: List[_ProfilerEvent]):
|
||
|
|
default_summary = f"{self.name}: {len(events)} events matched."
|
||
|
|
if self.should_benchmark:
|
||
|
|
# If benchmark summary is not empty, use it.
|
||
|
|
return (
|
||
|
|
self.benchmark_summary(events)
|
||
|
|
if hasattr(self, "benchmark") # type: ignore[attr-defined]
|
||
|
|
else default_summary
|
||
|
|
)
|
||
|
|
return default_summary
|
||
|
|
|
||
|
|
def benchmark_summary(self, events: List[_ProfilerEvent]):
|
||
|
|
def format_time(time_ns: int):
|
||
|
|
unit_lst = ["ns", "us", "ms"]
|
||
|
|
for unit in unit_lst:
|
||
|
|
if time_ns < 1000:
|
||
|
|
return f"{time_ns:.2f} {unit}"
|
||
|
|
time_ns //= 1000
|
||
|
|
return f"{time_ns:.2f} s"
|
||
|
|
|
||
|
|
assert hasattr(self, "benchmark"), "Please implement benchmark()"
|
||
|
|
shapes_factor_map = self.benchmark(events) # type: ignore[attr-defined]
|
||
|
|
original_time = sum(event.duration_time_ns for event in events)
|
||
|
|
new_time = sum(
|
||
|
|
shapes_factor_map[input_shapes(event)] * event.duration_time_ns
|
||
|
|
for event in events
|
||
|
|
)
|
||
|
|
return (
|
||
|
|
f"{self.name}: {len(events)} events matched. "
|
||
|
|
f"Total Estimated Speedup: {format_time(original_time - new_time)} ({round(original_time/new_time, 2)}X)"
|
||
|
|
)
|
||
|
|
|
||
|
|
def match(self, event: _ProfilerEvent):
|
||
|
|
"""
|
||
|
|
Return True if the event matches the pattern.
|
||
|
|
This method should be overriden in subclass.
|
||
|
|
"""
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
def matched_events(self):
|
||
|
|
if self.skip:
|
||
|
|
return []
|
||
|
|
matched_events = []
|
||
|
|
for event in self.eventTreeTraversal():
|
||
|
|
if self.match(event):
|
||
|
|
matched_events.append(event)
|
||
|
|
return matched_events
|
||
|
|
|
||
|
|
def root_of(self, event: _ProfilerEvent):
|
||
|
|
while event.parent:
|
||
|
|
event = event.parent
|
||
|
|
return event
|
||
|
|
|
||
|
|
def siblings_of(self, event: _ProfilerEvent):
|
||
|
|
if event.parent:
|
||
|
|
children = event.parent.children
|
||
|
|
else:
|
||
|
|
children = self.tid_root[event.start_tid]
|
||
|
|
index = children.index(event)
|
||
|
|
return children[:index], children[index + 1 :]
|
||
|
|
|
||
|
|
def next_of(self, event: _ProfilerEvent):
|
||
|
|
_, next_events = self.siblings_of(event)
|
||
|
|
return next_events[0] if next_events else None
|
||
|
|
|
||
|
|
def prev_of(self, event: _ProfilerEvent):
|
||
|
|
prev_events, _ = self.siblings_of(event)
|
||
|
|
return prev_events[-1] if prev_events else None
|
||
|
|
|
||
|
|
def go_up_until(self, event: _ProfilerEvent, predicate):
|
||
|
|
if not event:
|
||
|
|
return None
|
||
|
|
while event.parent and not predicate(event):
|
||
|
|
event = event.parent
|
||
|
|
return event
|
||
|
|
|
||
|
|
|
||
|
|
# Patterns
|
||
|
|
|
||
|
|
|
||
|
|
class NamePattern(Pattern):
|
||
|
|
def __init__(self, prof: profile, name: str, should_benchmark: bool = False):
|
||
|
|
super().__init__(prof, should_benchmark)
|
||
|
|
self.description = f"Matched Name Event: {name}"
|
||
|
|
self.name = name
|
||
|
|
|
||
|
|
def match(self, event: _ProfilerEvent):
|
||
|
|
return re.search(self.name, event.name) is not None
|
||
|
|
|
||
|
|
|
||
|
|
class ExtraCUDACopyPattern(Pattern):
|
||
|
|
"""
|
||
|
|
This pattern identifies if we creates a constant tensor on CPU and immediately moves it to GPU.
|
||
|
|
example: torch.zeros((100, 100)).to("cuda")
|
||
|
|
|
||
|
|
Pattern:
|
||
|
|
build-in method |build-in method
|
||
|
|
... | aten::to
|
||
|
|
aten::fill_/aten::zero_ | aten::_to_copy
|
||
|
|
|
||
|
|
Algorithm:
|
||
|
|
We start at node aten::to, go parent events' previous events,
|
||
|
|
and check if we have a aten::fill_/aten::zero_ as we keep going down the tree.
|
||
|
|
We always select the last child in the children list when we go down the tree.
|
||
|
|
If at any step we failed, it is not a match.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, prof: profile, should_benchmark: bool = False):
|
||
|
|
super().__init__(prof, should_benchmark)
|
||
|
|
self.name = "Extra CUDA Copy Pattern"
|
||
|
|
self.description = "Filled a CPU tensor and immediately moved it to GPU. Please initialize it on GPU."
|
||
|
|
self.url = "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#create-tensors-directly-on-the-target-device"
|
||
|
|
self.init_ops = {
|
||
|
|
"aten::fill_",
|
||
|
|
"aten::zero_",
|
||
|
|
"aten::normal_",
|
||
|
|
"aten::uniform_",
|
||
|
|
}
|
||
|
|
|
||
|
|
@property
|
||
|
|
def skip(self):
|
||
|
|
return not self.prof.with_stack or not self.prof.record_shapes
|
||
|
|
|
||
|
|
def match(self, event):
|
||
|
|
# TODO: We should also check tensor identities
|
||
|
|
if event.name != "aten::to":
|
||
|
|
return False
|
||
|
|
to_event = event
|
||
|
|
if not event.children:
|
||
|
|
return False
|
||
|
|
event = event.children[-1]
|
||
|
|
if event.name != "aten::_to_copy":
|
||
|
|
return False
|
||
|
|
if not event.children:
|
||
|
|
return False
|
||
|
|
event = event.children[-1]
|
||
|
|
if event.name != "aten::copy_":
|
||
|
|
return False
|
||
|
|
# aten::copy_ should have the first 2 args dtype the same
|
||
|
|
dtypes = input_dtypes(event)
|
||
|
|
if len(dtypes) < 2:
|
||
|
|
return False
|
||
|
|
if dtypes[0] is None or dtypes[0] != dtypes[1]:
|
||
|
|
return False
|
||
|
|
event = to_event
|
||
|
|
# Up one level
|
||
|
|
event = event.parent
|
||
|
|
if event is None:
|
||
|
|
return False
|
||
|
|
# Check if we have a aten::fill_ in previous leaf
|
||
|
|
event = self.prev_of(event)
|
||
|
|
if event is None:
|
||
|
|
return False
|
||
|
|
while event.children:
|
||
|
|
event = event.children[-1]
|
||
|
|
# aten::zero_ is a special optimzation case where fill_ is not called
|
||
|
|
if event.name in self.init_ops:
|
||
|
|
return True
|
||
|
|
return event.name in self.init_ops
|
||
|
|
# TODO: Check if tensor is reused
|
||
|
|
|
||
|
|
def benchmark(self, events: List[_ProfilerEvent]):
|
||
|
|
shapes_factor_map = {input_shapes(event): 0.0 for event in events}
|
||
|
|
for shape in shapes_factor_map:
|
||
|
|
size = shape[0]
|
||
|
|
to_timer = benchmark.Timer(
|
||
|
|
stmt='torch.ones(size).to("cuda")', globals={"size": size}
|
||
|
|
)
|
||
|
|
de_timer = benchmark.Timer(
|
||
|
|
stmt='torch.ones(size, device="cuda")', globals={"size": size}
|
||
|
|
)
|
||
|
|
to_time = to_timer.timeit(10).mean
|
||
|
|
de_time = de_timer.timeit(10).mean
|
||
|
|
shapes_factor_map[shape] = de_time / to_time
|
||
|
|
return shapes_factor_map
|
||
|
|
|
||
|
|
|
||
|
|
class ForLoopIndexingPattern(Pattern):
|
||
|
|
"""
|
||
|
|
This pattern identifies if we use a for loop to index a tensor that
|
||
|
|
can be vectorized.
|
||
|
|
example:
|
||
|
|
tensor = torch.empty((100, 100))
|
||
|
|
for i in range(100):
|
||
|
|
tensor[i] = i
|
||
|
|
|
||
|
|
Pattern:
|
||
|
|
aten::select | ... | aten::select | ... (Repeat)
|
||
|
|
|
||
|
|
Algorithm:
|
||
|
|
We start at node aten::select, and we check if we can find this alternating patterns.
|
||
|
|
We also keep a dictionary to avoid duplicate match in the for loop.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, prof: profile, should_benchmark: bool = False):
|
||
|
|
super().__init__(prof, should_benchmark)
|
||
|
|
self.name = "For Loop Indexing Pattern"
|
||
|
|
self.description = "For loop indexing detected. Vectorization recommended."
|
||
|
|
self.visited: Set[int] = set()
|
||
|
|
|
||
|
|
def eventTreeTraversal(self):
|
||
|
|
"""
|
||
|
|
We need to use BFS traversal order to avoid duplicate match.
|
||
|
|
"""
|
||
|
|
yield from traverse_bfs(self.event_tree)
|
||
|
|
|
||
|
|
def match(self, event: _ProfilerEvent):
|
||
|
|
if event.name != "aten::select":
|
||
|
|
return False
|
||
|
|
if event.id in self.visited:
|
||
|
|
return False
|
||
|
|
repeat_count = 1
|
||
|
|
_, next = self.siblings_of(event)
|
||
|
|
if len(next) <= 1:
|
||
|
|
return False
|
||
|
|
|
||
|
|
# Custom event list matching
|
||
|
|
def same_ops(list1, list2):
|
||
|
|
if len(list1) != len(list2):
|
||
|
|
return False
|
||
|
|
for op1, op2 in zip(list1, list2):
|
||
|
|
if op1.name != op2.name:
|
||
|
|
return False
|
||
|
|
return True
|
||
|
|
|
||
|
|
# Record the ops between two aten::select
|
||
|
|
next_select_idx = index_of_first_match(next, lambda e: e.name == "aten::select")
|
||
|
|
if next_select_idx is None:
|
||
|
|
return False
|
||
|
|
indexing_ops = [event] + next[:next_select_idx]
|
||
|
|
next = next[len(indexing_ops) - 1 :]
|
||
|
|
for i in range(0, len(next), len(indexing_ops)):
|
||
|
|
if same_ops(indexing_ops, next[i : i + len(indexing_ops)]):
|
||
|
|
repeat_count += 1
|
||
|
|
self.visited.add(next[i].id)
|
||
|
|
else:
|
||
|
|
break
|
||
|
|
return repeat_count >= 10
|
||
|
|
|
||
|
|
|
||
|
|
class FP32MatMulPattern(Pattern):
|
||
|
|
def __init__(self, prof: profile, should_benchmark: bool = False):
|
||
|
|
super().__init__(prof, should_benchmark)
|
||
|
|
self.name = "FP32 MatMul Pattern"
|
||
|
|
self.description = (
|
||
|
|
"You are currently using GPU that supports TF32. "
|
||
|
|
"Please enable TF32 by setting 'torch.backends.cuda.matmul.allow_tf32 = True'"
|
||
|
|
)
|
||
|
|
self.url = "https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
|
||
|
|
|
||
|
|
@property
|
||
|
|
def skip(self):
|
||
|
|
if torch.version.hip is not None:
|
||
|
|
has_tf32 = False
|
||
|
|
else:
|
||
|
|
# Anything less than sm_80 is not Ampere which doesn't support TF32
|
||
|
|
has_tf32 = all(int(arch[3:]) >= 80 for arch in torch.cuda.get_arch_list())
|
||
|
|
return has_tf32 is False or super().skip or not self.prof.record_shapes
|
||
|
|
|
||
|
|
def match(self, event: _ProfilerEvent):
|
||
|
|
# If we saw this pattern once, we don't need to match it again
|
||
|
|
if event.tag != _EventType.TorchOp:
|
||
|
|
return False
|
||
|
|
assert isinstance(event.extra_fields, _ExtraFields_TorchOp)
|
||
|
|
if event.name == "aten::mm":
|
||
|
|
if event.extra_fields.allow_tf32_cublas is False:
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
def report(self, event: _ProfilerEvent):
|
||
|
|
return self.description
|
||
|
|
|
||
|
|
def benchmark(self, events: List[_ProfilerEvent]):
|
||
|
|
shapes_factor_map = {input_shapes(event): 0.0 for event in events}
|
||
|
|
for shape in shapes_factor_map:
|
||
|
|
matrixA = torch.randn(shape[0], device="cuda", dtype=torch.float32)
|
||
|
|
matrixB = torch.randn(shape[1], device="cuda", dtype=torch.float32)
|
||
|
|
fp32_timer = benchmark.Timer(
|
||
|
|
stmt="torch.mm(matrixA, matrixB)",
|
||
|
|
globals={"matrixA": matrixA, "matrixB": matrixB},
|
||
|
|
)
|
||
|
|
tf32_timer = benchmark.Timer(
|
||
|
|
stmt="torch.mm(matrixA, matrixB)",
|
||
|
|
setup="torch.backends.cuda.matmul.allow_tf32 = True",
|
||
|
|
globals={"matrixA": matrixA, "matrixB": matrixB},
|
||
|
|
)
|
||
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
||
|
|
fp32_time = fp32_timer.timeit(10).mean
|
||
|
|
tf32_time = tf32_timer.timeit(10).mean
|
||
|
|
shapes_factor_map[shape] = tf32_time / fp32_time
|
||
|
|
return shapes_factor_map
|
||
|
|
|
||
|
|
|
||
|
|
class OptimizerSingleTensorPattern(Pattern):
|
||
|
|
"""
|
||
|
|
This pattern identifies if we are using the single-tensor version of an optimizer.
|
||
|
|
example:
|
||
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
|
||
|
|
By adding foreach=True to enable multi-tensor optimizer, we can gain speedup when
|
||
|
|
the kernels are relatively small.
|
||
|
|
|
||
|
|
Pattern:
|
||
|
|
XXXXX: _single_tenser_<OPTIMIZER_NAME>
|
||
|
|
|
||
|
|
Algorithm:
|
||
|
|
String match
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, prof: profile, should_benchmark: bool = False):
|
||
|
|
super().__init__(prof, should_benchmark)
|
||
|
|
self.name = "Optimizer Single Tensor Pattern"
|
||
|
|
self.optimizers_with_foreach = ["adam", "sgd", "adamw"]
|
||
|
|
self.description = (
|
||
|
|
"Deteced optimizer running with single tensor implementation. "
|
||
|
|
"Please enable multi tensor implementation by passing 'foreach=True' into optimizer."
|
||
|
|
)
|
||
|
|
self.url = ""
|
||
|
|
|
||
|
|
def match(self, event: _ProfilerEvent):
|
||
|
|
for optimizer in self.optimizers_with_foreach:
|
||
|
|
if event.name.endswith(f"_single_tensor_{optimizer}"):
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
class SynchronizedDataLoaderPattern(Pattern):
|
||
|
|
"""
|
||
|
|
This pattern identifies if we are using num_workers=0 in DataLoader.
|
||
|
|
example:
|
||
|
|
torch.utils.data.DataLoader(dataset, batch_size=batch_size)
|
||
|
|
Add num_workers=N to the arguments. N depends on system configuration.
|
||
|
|
|
||
|
|
Pattern:
|
||
|
|
dataloader.py(...): __iter__
|
||
|
|
dataloader.py(...): _get_iterator
|
||
|
|
NOT dataloader.py(...): check_worker_number_rationality
|
||
|
|
|
||
|
|
Algorithm:
|
||
|
|
If we don't see check_worker_number_rationality call in the dataloader __iter__,
|
||
|
|
It is not an asynchronous dataloader.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, prof: profile, should_benchmark: bool = False):
|
||
|
|
super().__init__(prof, should_benchmark)
|
||
|
|
self.name = "Synchronized DataLoader Pattern"
|
||
|
|
self.description = (
|
||
|
|
"Detected DataLoader running with synchronized implementation. "
|
||
|
|
"Please enable asynchronous dataloading by setting num_workers > 0 when initializing DataLoader."
|
||
|
|
)
|
||
|
|
self.url = (
|
||
|
|
"https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html"
|
||
|
|
"#enable-async-data-loading-and-augmentation"
|
||
|
|
)
|
||
|
|
|
||
|
|
def match(self, event: _ProfilerEvent):
|
||
|
|
def is_dataloader_function(name: str, function_name: str):
|
||
|
|
return name.startswith(
|
||
|
|
os.path.join("torch", "utils", "data", "dataloader.py")
|
||
|
|
) and name.endswith(function_name)
|
||
|
|
|
||
|
|
# TODO: fixme! Due to lifetime issues of the function name, this field might
|
||
|
|
# actually point to an already freed string when the even is a PyCall.
|
||
|
|
# Just silently skip this to unblock testing.
|
||
|
|
try:
|
||
|
|
event.name
|
||
|
|
except UnicodeDecodeError:
|
||
|
|
return False
|
||
|
|
|
||
|
|
if not is_dataloader_function(event.name, "__iter__"):
|
||
|
|
return False
|
||
|
|
if not event.children:
|
||
|
|
return False
|
||
|
|
event = event.children[0]
|
||
|
|
if not is_dataloader_function(event.name, "_get_iterator"):
|
||
|
|
return False
|
||
|
|
if not event.children:
|
||
|
|
return False
|
||
|
|
event = event.children[0]
|
||
|
|
return not is_dataloader_function(event.name, "check_worker_number_rationality")
|
||
|
|
# TODO: We should also check if the loader is bottleneck.
|
||
|
|
|
||
|
|
|
||
|
|
class GradNotSetToNonePattern(Pattern):
|
||
|
|
"""
|
||
|
|
This pattern identifies if we are not setting grad to None in zero_grad.
|
||
|
|
example:
|
||
|
|
optimizer.zero_grad()
|
||
|
|
By setting set_to_none=True, we can gain speedup
|
||
|
|
|
||
|
|
Pattern:
|
||
|
|
XXXXX: _zero_grad
|
||
|
|
NOT aten::zeros
|
||
|
|
aten::zero_
|
||
|
|
|
||
|
|
aten::zero_ is called on each parameter in the model.
|
||
|
|
We also want to make sure it is not called by aten::zeros.
|
||
|
|
|
||
|
|
Algorithm:
|
||
|
|
String match
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, prof: profile, should_benchmark: bool = False):
|
||
|
|
super().__init__(prof, should_benchmark)
|
||
|
|
self.name = "Gradient Set To Zero Instead of None Pattern"
|
||
|
|
self.description = (
|
||
|
|
"Detected gradient set to zero instead of None. "
|
||
|
|
"Please add 'set_to_none=True' when calling zero_grad()."
|
||
|
|
)
|
||
|
|
self.url = (
|
||
|
|
"https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html"
|
||
|
|
"#disable-gradient-calculation-for-validation-or-inference"
|
||
|
|
)
|
||
|
|
|
||
|
|
def match(self, event: _ProfilerEvent):
|
||
|
|
if not event.name.endswith(": zero_grad"):
|
||
|
|
return False
|
||
|
|
if not event.children:
|
||
|
|
return False
|
||
|
|
|
||
|
|
for sub_event in traverse_dfs(event.children):
|
||
|
|
if (
|
||
|
|
sub_event.name == "aten::zero_"
|
||
|
|
and sub_event.parent.name != "aten::zeros"
|
||
|
|
):
|
||
|
|
return True
|
||
|
|
# TODO: We should also check if the optimizer's numerical behavior will change.
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
class Conv2dBiasFollowedByBatchNorm2dPattern(Pattern):
|
||
|
|
"""
|
||
|
|
This pattern identifies if we are enabling bias in Conv2d which is followed by BatchNorm2d.
|
||
|
|
Bias doesn't do anything when followed by batchnorm.
|
||
|
|
Pattern:
|
||
|
|
nn.Module: Conv2d | nn.Module: BatchNorm2d
|
||
|
|
...
|
||
|
|
aten::conv2d AND dtype of third argument is not null
|
||
|
|
The third argument is the bias
|
||
|
|
Algorithm:
|
||
|
|
String match
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, prof: profile, should_benchmark: bool = False):
|
||
|
|
super().__init__(prof, should_benchmark)
|
||
|
|
self.name = "Enabling Bias in Conv2d Followed By BatchNorm Pattern"
|
||
|
|
self.description = "Detected bias enabled in Conv2d that is followed by BatchNorm2d. Please set 'bias=False' in Conv2d."
|
||
|
|
self.url = (
|
||
|
|
"https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html"
|
||
|
|
"#disable-bias-for-convolutions-directly-followed-by-a-batch-norm"
|
||
|
|
)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def skip(self):
|
||
|
|
return self.prof.record_shapes is False or super().skip
|
||
|
|
|
||
|
|
def match(self, event: _ProfilerEvent):
|
||
|
|
if event.name != "aten::conv2d":
|
||
|
|
return False
|
||
|
|
if len(input_dtypes(event)) < 3 or input_dtypes(event)[2] is None:
|
||
|
|
return False
|
||
|
|
# This means bias=True
|
||
|
|
event = self.go_up_until(
|
||
|
|
event, lambda e: e.name.startswith("nn.Module: Conv2d")
|
||
|
|
)
|
||
|
|
if not event:
|
||
|
|
return False
|
||
|
|
event = self.next_of(event)
|
||
|
|
if not event:
|
||
|
|
return False
|
||
|
|
return event.name.startswith("nn.Module: BatchNorm2d")
|
||
|
|
|
||
|
|
|
||
|
|
class MatMulDimInFP16Pattern(Pattern):
|
||
|
|
def __init__(self, prof: profile, should_benchmark: bool = False):
|
||
|
|
super().__init__(prof, should_benchmark)
|
||
|
|
self.name = "Matrix Multiplication Dimension Not Aligned Pattern"
|
||
|
|
self.description = "Detected matmul with dimension not aligned. Please use matmul with aligned dimension."
|
||
|
|
self.url = "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#use-mixed-precision-and-amp"
|
||
|
|
|
||
|
|
@property
|
||
|
|
def skip(self):
|
||
|
|
return not self.prof.with_stack or not self.prof.record_shapes
|
||
|
|
|
||
|
|
def match(self, event: _ProfilerEvent):
|
||
|
|
def mutiple_of(shapes, multiple):
|
||
|
|
return all(dim % multiple == 0 for shape in shapes for dim in shape[-2:])
|
||
|
|
|
||
|
|
if event.name not in ("aten::mm", "aten::bmm", "aten::addmm"):
|
||
|
|
return False
|
||
|
|
if not input_dtypes(event):
|
||
|
|
return False
|
||
|
|
arg_dtype = input_dtypes(event)[0]
|
||
|
|
if arg_dtype in (torch.bfloat16, torch.half) and not mutiple_of(
|
||
|
|
input_shapes(event), 8
|
||
|
|
):
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
def benchmark(self, events: List[_ProfilerEvent]):
|
||
|
|
def closest_multiple(shapes, multiple):
|
||
|
|
return [multiple * math.ceil(shape / multiple) for shape in shapes]
|
||
|
|
|
||
|
|
shapes_factor_map = {input_shapes(event): 0.0 for event in events}
|
||
|
|
for shape in shapes_factor_map:
|
||
|
|
matrixA = torch.randn(shape[0], device="cuda", dtype=torch.float16)
|
||
|
|
matrixB = torch.randn(shape[1], device="cuda", dtype=torch.float16)
|
||
|
|
not_aligned_dim_timer = benchmark.Timer(
|
||
|
|
stmt="torch.mm(matrixA, matrixB)",
|
||
|
|
globals={"matrixA": matrixA, "matrixB": matrixB},
|
||
|
|
)
|
||
|
|
matrixA = torch.randn(
|
||
|
|
closest_multiple(shape[0], 8), device="cuda", dtype=torch.float16
|
||
|
|
)
|
||
|
|
matrixB = torch.randn(
|
||
|
|
closest_multiple(shape[1], 8), device="cuda", dtype=torch.float16
|
||
|
|
)
|
||
|
|
aligned_dim_timer = benchmark.Timer(
|
||
|
|
stmt="torch.mm(matrixA, matrixB)",
|
||
|
|
globals={"matrixA": matrixA, "matrixB": matrixB},
|
||
|
|
)
|
||
|
|
not_aligned_dim_time = not_aligned_dim_timer.timeit(10).mean
|
||
|
|
aligned_dim_time = aligned_dim_timer.timeit(10).mean
|
||
|
|
shapes_factor_map[shape] = aligned_dim_time / not_aligned_dim_time
|
||
|
|
return shapes_factor_map
|
||
|
|
|
||
|
|
|
||
|
|
def source_code_location(event: Optional[_ProfilerEvent]):
|
||
|
|
while event:
|
||
|
|
if event.tag == _EventType.PyCall or event.tag == _EventType.PyCCall:
|
||
|
|
assert isinstance(
|
||
|
|
event.extra_fields, (_ExtraFields_PyCall, _ExtraFields_PyCCall)
|
||
|
|
)
|
||
|
|
if not event.extra_fields.caller.file_name.startswith("torch" + os.sep):
|
||
|
|
return f"{event.extra_fields.caller.file_name}:{event.extra_fields.caller.line_number}"
|
||
|
|
event = event.parent
|
||
|
|
return "No source code location found"
|
||
|
|
|
||
|
|
|
||
|
|
def input_shapes(event: _ProfilerEvent):
|
||
|
|
assert isinstance(event.extra_fields, _ExtraFields_TorchOp)
|
||
|
|
return tuple(tuple(getattr(i, "sizes", ())) for i in event.extra_fields.inputs)
|
||
|
|
|
||
|
|
|
||
|
|
def input_dtypes(event: _ProfilerEvent):
|
||
|
|
assert isinstance(event.extra_fields, _ExtraFields_TorchOp)
|
||
|
|
return tuple(getattr(i, "dtype", None) for i in event.extra_fields.inputs)
|
||
|
|
|
||
|
|
|
||
|
|
def report_all_anti_patterns(
|
||
|
|
prof,
|
||
|
|
should_benchmark: bool = False,
|
||
|
|
print_enable: bool = True,
|
||
|
|
json_report_dir: Optional[str] = None,
|
||
|
|
):
|
||
|
|
report_dict: Dict = {}
|
||
|
|
anti_patterns = [
|
||
|
|
ExtraCUDACopyPattern(prof, should_benchmark),
|
||
|
|
# ForLoopIndexingPattern(prof, should_benchmark),
|
||
|
|
FP32MatMulPattern(prof, should_benchmark),
|
||
|
|
OptimizerSingleTensorPattern(prof, should_benchmark),
|
||
|
|
SynchronizedDataLoaderPattern(prof, should_benchmark),
|
||
|
|
GradNotSetToNonePattern(prof, should_benchmark),
|
||
|
|
Conv2dBiasFollowedByBatchNorm2dPattern(prof, should_benchmark),
|
||
|
|
MatMulDimInFP16Pattern(prof, should_benchmark),
|
||
|
|
]
|
||
|
|
reported = set()
|
||
|
|
summaries = []
|
||
|
|
message_list = [f"{'-'*40}TorchTidy Report{'-'*40}"]
|
||
|
|
message_list.append("Matched Events:")
|
||
|
|
|
||
|
|
for anti_pattern in anti_patterns:
|
||
|
|
matched_events = anti_pattern.matched_events()
|
||
|
|
if not matched_events:
|
||
|
|
continue
|
||
|
|
summaries.append(anti_pattern.summary(matched_events))
|
||
|
|
for event in matched_events:
|
||
|
|
report_msg = anti_pattern.report(event)
|
||
|
|
if report_msg not in reported:
|
||
|
|
message_list.append(report_msg)
|
||
|
|
reported.add(report_msg)
|
||
|
|
src_location, line_no = source_code_location(event).split(":")
|
||
|
|
report_dict.setdefault(src_location, []).append(
|
||
|
|
{
|
||
|
|
"line_number": int(line_no),
|
||
|
|
"name": anti_pattern.name,
|
||
|
|
"url": anti_pattern.url,
|
||
|
|
"message": anti_pattern.description,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
if json_report_dir is not None:
|
||
|
|
json_report_path = os.path.join(json_report_dir, "torchtidy_report.json")
|
||
|
|
if os.path.exists(json_report_path):
|
||
|
|
with open(json_report_path) as f:
|
||
|
|
exisiting_report = json.load(f)
|
||
|
|
exisiting_report.update(report_dict)
|
||
|
|
report_dict = exisiting_report
|
||
|
|
with open(json_report_path, "w") as f:
|
||
|
|
json.dump(report_dict, f, indent=4)
|
||
|
|
|
||
|
|
message_list.append("Summary:")
|
||
|
|
message_list += summaries
|
||
|
|
message_list.append(f"{'-'*40}TorchTidy Report{'-'*40}")
|
||
|
|
if print_enable:
|
||
|
|
print("\n".join(message_list))
|