Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,362 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Post-process a profiler trace to add CUDA graph kernel annotations.
|
||||
|
||||
Reads a profiler trace (gzipped or plain JSON) and a kernel annotations
|
||||
pickle, matches kernel events by their graph node id, and writes an
|
||||
annotated trace with the annotation fields added to each kernel event's
|
||||
args (displayed alongside grid/block size in trace viewers).
|
||||
|
||||
The annotations pickle is auto-discovered from the trace file's parent
|
||||
directory (one level up, matching the rank from the trace filename).
|
||||
|
||||
Usage:
|
||||
python -m torch.cuda._annotate_cuda_graph_trace <trace_file> [-a <annotations_pkl>] [-o <output_file>]
|
||||
|
||||
Examples:
|
||||
# Auto-discover annotations pickle from trace location
|
||||
python -m torch.cuda._annotate_cuda_graph_trace \\
|
||||
traces/step_000000000014/000000.*.pt.trace.json.gz
|
||||
|
||||
# Explicit annotations pickle
|
||||
python -m torch.cuda._annotate_cuda_graph_trace trace.json.gz -a annotations.pkl
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import gzip
|
||||
import json
|
||||
import pickle
|
||||
import re
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
_WORK_CATEGORIES = {"kernel", "gpu_memcpy", "gpu_memset"}
|
||||
|
||||
|
||||
def _move_overlapping_to_stream(
|
||||
trace: dict, default_stream: int = 7, overlap_stream: int = 8
|
||||
) -> int:
|
||||
"""Move graphed kernels that overlap with their predecessor to a separate stream.
|
||||
|
||||
Perfetto cannot display overlapping (non-nested) events on the same
|
||||
stream -- they get hidden. This pass detects graphed kernel events on
|
||||
*default_stream* whose start timestamp falls before the previous
|
||||
kernel's end, and moves them to *overlap_stream* so they're visible.
|
||||
|
||||
Returns the number of events moved.
|
||||
"""
|
||||
graphed_on_default = [
|
||||
e
|
||||
for e in trace["traceEvents"]
|
||||
if e.get("cat") == "kernel"
|
||||
and e.get("tid") == default_stream
|
||||
and e.get("args", {}).get("graph node id", 0) != 0
|
||||
]
|
||||
graphed_on_default.sort(key=lambda e: e["ts"])
|
||||
|
||||
moved = 0
|
||||
prev_end = 0.0
|
||||
for event in graphed_on_default:
|
||||
ts = event["ts"]
|
||||
dur = event.get("dur", 0)
|
||||
if ts < prev_end:
|
||||
event["tid"] = overlap_stream
|
||||
event.get("args", {})["stream"] = overlap_stream
|
||||
moved += 1
|
||||
else:
|
||||
prev_end = ts + dur
|
||||
|
||||
return moved
|
||||
|
||||
|
||||
def _fix_overlapping_timestamps(trace: dict, max_adjust_us: float = 1.0) -> int:
|
||||
"""Clamp graphed kernel/memcpy timestamps so they don't overlap on the same stream.
|
||||
|
||||
CUPTI can produce slightly overlapping timestamps for consecutive graphed
|
||||
events, causing Perfetto to hide events that sit entirely "under" their
|
||||
neighbours. This pass sorts graphed work events per stream and ensures
|
||||
each event starts at or after the previous event's end.
|
||||
|
||||
Overlaps larger than *max_adjust_us* are flagged as warnings and left
|
||||
unchanged, since they likely indicate a real issue rather than CUPTI
|
||||
timestamp jitter.
|
||||
|
||||
Returns the number of events adjusted.
|
||||
"""
|
||||
per_stream: dict[int, list[dict]] = defaultdict(list)
|
||||
for event in trace["traceEvents"]:
|
||||
if (
|
||||
event.get("cat") in _WORK_CATEGORIES
|
||||
and event.get("args", {}).get("graph node id", 0) != 0
|
||||
):
|
||||
per_stream[event.get("tid")].append(event)
|
||||
|
||||
adjusted = 0
|
||||
for tid, events in per_stream.items():
|
||||
events.sort(key=lambda e: e["ts"])
|
||||
prev_end = 0.0
|
||||
for event in events:
|
||||
ts = event["ts"]
|
||||
dur = event.get("dur", 0)
|
||||
if ts < prev_end:
|
||||
overlap = prev_end - ts
|
||||
if overlap > max_adjust_us:
|
||||
print(
|
||||
f"WARNING: large overlap {overlap:.3f}us on stream {tid} "
|
||||
f"for {event.get('name', '?')[:60]}, skipping adjustment"
|
||||
)
|
||||
else:
|
||||
event["ts"] = prev_end
|
||||
adjusted += 1
|
||||
prev_end = event["ts"] + dur
|
||||
|
||||
return adjusted
|
||||
|
||||
|
||||
def annotate_trace(
|
||||
trace: dict,
|
||||
annotations: dict[int, list[Any]],
|
||||
default_stream: int = 7,
|
||||
) -> int:
|
||||
"""Add annotation fields to kernel events matching the annotations dict.
|
||||
|
||||
Each annotation entry is a list (from nested ``mark_kernels`` scopes).
|
||||
Fields from all annotations are merged into the event args; if multiple
|
||||
annotations define the same key, later entries in the list win.
|
||||
|
||||
For graphed events (graph_node_id != 0), reassigns ``tid`` and
|
||||
``args["stream"]`` to the stream recorded in annotations, or to
|
||||
*default_stream* if there is no annotation. Also moves the
|
||||
corresponding ``ac2g`` flow-finish events to the new tid so that
|
||||
CPU-to-GPU correlation arrows are preserved.
|
||||
|
||||
Removes ``gpu_user_annotation`` events and orphaned ``ac2g`` events
|
||||
from streams that have no kernel or memcpy events after reassignment,
|
||||
since CUPTI replicates these onto every stream during graph replay.
|
||||
|
||||
Returns the number of events annotated.
|
||||
"""
|
||||
# Build an index of ac2g 'f' events keyed by (tid, ts) so we can
|
||||
# move them together with the kernel events they correspond to.
|
||||
ac2g_f_index: dict[tuple, list] = {}
|
||||
for event in trace["traceEvents"]:
|
||||
if event.get("cat") == "ac2g" and event.get("ph") == "f":
|
||||
key = (event.get("tid"), event.get("ts"))
|
||||
ac2g_f_index.setdefault(key, []).append(event)
|
||||
|
||||
annotated = 0
|
||||
for event in trace.get("traceEvents", []):
|
||||
args = event.get("args", {})
|
||||
graph_node_id = args.get("graph node id")
|
||||
if graph_node_id is None or graph_node_id == 0:
|
||||
continue
|
||||
stream_id = None
|
||||
if graph_node_id in annotations:
|
||||
for ann in annotations[graph_node_id]:
|
||||
if isinstance(ann, dict):
|
||||
for key, value in ann.items():
|
||||
args[key] = str(value)
|
||||
if "stream" in ann:
|
||||
stream_id = int(ann["stream"])
|
||||
else:
|
||||
args["annotation"] = str(ann)
|
||||
annotated += 1
|
||||
|
||||
# Reassign stream: use annotated stream if available, else default
|
||||
if stream_id is None:
|
||||
stream_id = default_stream
|
||||
old_key = (event.get("tid"), event.get("ts"))
|
||||
event["tid"] = stream_id
|
||||
args["stream"] = stream_id
|
||||
|
||||
# Move the matching ac2g 'f' event(s) to the same new tid
|
||||
for ac2g_event in ac2g_f_index.get(old_key, ()):
|
||||
ac2g_event["tid"] = stream_id
|
||||
|
||||
# Remove gpu_user_annotation events and ac2g flow-finish events from
|
||||
# streams that have no real kernel/memcpy/memset work -- these are
|
||||
# noise replicated by CUPTI onto every stream during graph replay.
|
||||
tids_with_work = set()
|
||||
for event in trace["traceEvents"]:
|
||||
if event.get("cat") in _WORK_CATEGORIES:
|
||||
tids_with_work.add(event.get("tid"))
|
||||
|
||||
def _is_noise(event: dict) -> bool:
|
||||
cat = event.get("cat")
|
||||
if cat == "gpu_user_annotation":
|
||||
return event.get("tid") not in tids_with_work
|
||||
if cat == "ac2g" and event.get("ph") == "f":
|
||||
return event.get("tid") not in tids_with_work
|
||||
return False
|
||||
|
||||
original_count = len(trace["traceEvents"])
|
||||
trace["traceEvents"] = [
|
||||
event for event in trace["traceEvents"] if not _is_noise(event)
|
||||
]
|
||||
removed = original_count - len(trace["traceEvents"])
|
||||
if removed:
|
||||
print(f"Removed {removed} noise events from empty streams")
|
||||
|
||||
# Clean up metadata: remove thread_name / thread_sort_index entries
|
||||
# for noise streams that have no real (non-metadata) events, and add
|
||||
# thread_name entries for our new annotation streams.
|
||||
all_tids_in_trace = {
|
||||
e.get("tid") for e in trace["traceEvents"] if e.get("ph") != "M"
|
||||
}
|
||||
# Find the GPU process pid from existing thread_name metadata
|
||||
gpu_pid = 0
|
||||
for event in trace["traceEvents"]:
|
||||
if (
|
||||
event.get("ph") == "M"
|
||||
and event.get("name") == "thread_name"
|
||||
and str(event.get("args", {}).get("name", "")).startswith("stream ")
|
||||
):
|
||||
gpu_pid = event.get("pid", 0)
|
||||
break
|
||||
|
||||
# Remove metadata entries for tids with no non-metadata events
|
||||
trace["traceEvents"] = [
|
||||
event
|
||||
for event in trace["traceEvents"]
|
||||
if event.get("ph") != "M" or event.get("tid") in all_tids_in_trace
|
||||
]
|
||||
|
||||
# Add thread_name metadata for new annotation tids that lack one
|
||||
existing_thread_names = {
|
||||
e.get("tid")
|
||||
for e in trace["traceEvents"]
|
||||
if e.get("ph") == "M" and e.get("name") == "thread_name"
|
||||
}
|
||||
for tid in sorted(tids_with_work - existing_thread_names):
|
||||
trace["traceEvents"].append(
|
||||
{
|
||||
"ph": "M",
|
||||
"pid": gpu_pid,
|
||||
"tid": tid,
|
||||
"name": "thread_name",
|
||||
"args": {"name": f"stream {tid}"},
|
||||
}
|
||||
)
|
||||
|
||||
return annotated
|
||||
|
||||
|
||||
def load_trace(path: Path) -> dict:
|
||||
if path.suffix == ".gz" or path.name.endswith(".json.gz"):
|
||||
with gzip.open(path, "rt") as f:
|
||||
return json.load(f)
|
||||
else:
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def save_trace(trace: dict, path: Path) -> None:
|
||||
if path.suffix == ".gz" or path.name.endswith(".json.gz"):
|
||||
with gzip.open(path, "wt") as f:
|
||||
json.dump(trace, f)
|
||||
else:
|
||||
with open(path, "w") as f:
|
||||
json.dump(trace, f)
|
||||
|
||||
|
||||
def _find_annotations_pkl(trace_file: Path) -> Path | None:
|
||||
"""Auto-discover the annotations pickle from the trace file location.
|
||||
|
||||
Trace files live in e.g. ``traces/step_000000000014/000000.<id>.pt.trace.json.gz``
|
||||
where the leading digits are the rank. The pickle lives one level up:
|
||||
``traces/kernel_annotations_rank0_*.pkl``.
|
||||
"""
|
||||
match = re.match(r"^(\d+)", trace_file.name)
|
||||
if not match:
|
||||
return None
|
||||
rank = int(match.group(1))
|
||||
|
||||
traces_dir = trace_file.parent.parent
|
||||
candidates = sorted(traces_dir.glob(f"kernel_annotations_rank{rank}_*.pkl"))
|
||||
if candidates:
|
||||
return candidates[0]
|
||||
return None
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Annotate a profiler trace with CUDA graph kernel annotations."
|
||||
)
|
||||
parser.add_argument(
|
||||
"trace_file", type=Path, help="Input trace file (.json or .json.gz)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-a",
|
||||
"--annotations",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Kernel annotations pickle file. Auto-discovered if omitted.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--output",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Output file path. Defaults to <trace_file>.annotated.<ext>",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--default-stream",
|
||||
type=int,
|
||||
default=7,
|
||||
help="Stream ID to assign to unannotated graphed events (default: 7).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
annotations_pkl = args.annotations
|
||||
if annotations_pkl is None:
|
||||
annotations_pkl = _find_annotations_pkl(args.trace_file)
|
||||
if annotations_pkl is None:
|
||||
print(
|
||||
f"Could not auto-discover annotations pickle for {args.trace_file}. "
|
||||
f"Use -a to specify it explicitly.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
print(f"Auto-discovered annotations: {annotations_pkl}")
|
||||
|
||||
with open(annotations_pkl, "rb") as f:
|
||||
annotations = pickle.load(f)
|
||||
print(f"Loaded {len(annotations)} kernel annotations")
|
||||
|
||||
trace = load_trace(args.trace_file)
|
||||
total_events = len(trace.get("traceEvents", []))
|
||||
print(f"Loaded trace with {total_events} events")
|
||||
|
||||
count = annotate_trace(trace, annotations, default_stream=args.default_stream)
|
||||
print(f"Annotated {count} kernel events")
|
||||
|
||||
overlap_moved = _move_overlapping_to_stream(
|
||||
trace, default_stream=args.default_stream
|
||||
)
|
||||
if overlap_moved:
|
||||
print(f"Moved {overlap_moved} overlapping events to stream 8")
|
||||
|
||||
ts_fixed = _fix_overlapping_timestamps(trace)
|
||||
if ts_fixed:
|
||||
print(f"Fixed {ts_fixed} overlapping graphed event timestamps")
|
||||
|
||||
output = args.output
|
||||
if output is None:
|
||||
name = args.trace_file.name
|
||||
if name.endswith(".json.gz"):
|
||||
output = args.trace_file.with_name(
|
||||
name.replace(".json.gz", ".annotated.json.gz")
|
||||
)
|
||||
elif name.endswith(".json"):
|
||||
output = args.trace_file.with_suffix(".annotated.json")
|
||||
else:
|
||||
output = args.trace_file.with_suffix(args.trace_file.suffix + ".annotated")
|
||||
|
||||
save_trace(trace, output)
|
||||
print(f"Saved annotated trace to {output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,140 @@
|
||||
import torch
|
||||
from torch._C import dtype
|
||||
|
||||
|
||||
__all__ = ["GPULimits"]
|
||||
|
||||
|
||||
class GPULimits:
|
||||
r"""Utility class that provides the theoretical limits of Nvidia GPU devices. The
|
||||
limits don't take into account thermal throttling (assume that the GPU run at its
|
||||
peak rated frequency). This is because user hardware configuration may influence
|
||||
power behavior.
|
||||
"""
|
||||
|
||||
def __init__(self, target_device: torch.device):
|
||||
# The device properties object is obtained by calling 'cudaGetDeviceProperties' CUDA
|
||||
# runtime function. We need the total memory bus width and the memory clock rate to
|
||||
# calculate the memory bandwidth.
|
||||
self.device_properties = torch.cuda.get_device_properties(target_device)
|
||||
|
||||
# The compute capability is needed to determine the number of FLOPs per cycle per SM
|
||||
self.compute_capability = int(
|
||||
f"{self.device_properties.major}{self.device_properties.minor}"
|
||||
)
|
||||
|
||||
# FLOPs per cycle information derived from Table 2 in:
|
||||
# https://resources.nvidia.com/en-us-hopper-architecture/nvidia-h100-tensor-c
|
||||
|
||||
# Returns the number of FMA instructions retired per cycle per SM for a given
|
||||
# data type, when tensor cores are NOT used
|
||||
def get_fma_per_cycle_per_sm_cuda_cores(self, data_type: dtype) -> int:
|
||||
hardcoded_device_values = {
|
||||
# Ampere Architecture
|
||||
"fp16_80": 256,
|
||||
"fp32_80": 64,
|
||||
"fp64_80": 32,
|
||||
# Hopper Architecture
|
||||
"fp16_90": 64,
|
||||
"fp32_90": 128,
|
||||
"fp64_90": 64,
|
||||
# Blackwell Architecture
|
||||
"fp16_100": 256,
|
||||
"fp32_100": 128,
|
||||
"fp64_100": 64,
|
||||
}
|
||||
dict_key = ""
|
||||
if data_type is torch.float16:
|
||||
dict_key = f"fp16_{self.compute_capability}"
|
||||
elif data_type is torch.float32:
|
||||
dict_key = f"fp32_{self.compute_capability}"
|
||||
elif data_type is torch.float64:
|
||||
dict_key = f"fp64_{self.compute_capability}"
|
||||
else:
|
||||
dict_key = "unknown"
|
||||
|
||||
if dict_key not in hardcoded_device_values:
|
||||
raise RuntimeError(
|
||||
f"No data for sm_{self.compute_capability} and {data_type}."
|
||||
)
|
||||
|
||||
return hardcoded_device_values[dict_key]
|
||||
|
||||
# Returns the number of FMA instructions retired per cycle per SM for a given
|
||||
# data type, when tensor cores ARE used
|
||||
def get_fma_per_cycle_per_sm_tensor_cores(self, data_type: dtype) -> int:
|
||||
hardcoded_device_values = {
|
||||
# Ampere Architecture
|
||||
"int8_80": 2048,
|
||||
"fp16_80": 1024,
|
||||
"fp32_80": 512,
|
||||
"fp64_80": 64,
|
||||
# Hopper Architecture
|
||||
"int8_90": 4096,
|
||||
"fp8_90": 4096,
|
||||
"fp16_90": 2048,
|
||||
"fp32_90": 1024,
|
||||
"fp64_90": 128,
|
||||
# Blackwell Architecture
|
||||
"int8_100": 8192,
|
||||
"fp8_100": 8192,
|
||||
"fp16_100": 4096,
|
||||
"fp32_100": 2048,
|
||||
}
|
||||
dict_key = ""
|
||||
if data_type is torch.float16:
|
||||
dict_key = f"fp16_{self.compute_capability}"
|
||||
elif data_type is torch.bfloat16:
|
||||
# FP16 and BF16 are equivalent in terms of FLOPs per cycle per SM
|
||||
dict_key = f"fp16_{self.compute_capability}"
|
||||
elif data_type is torch.float32:
|
||||
dict_key = f"fp32_{self.compute_capability}"
|
||||
elif data_type is torch.int8:
|
||||
dict_key = f"int8_{self.compute_capability}"
|
||||
elif data_type is torch.float64:
|
||||
dict_key = f"fp64_{self.compute_capability}"
|
||||
else:
|
||||
dict_key = "unknown"
|
||||
|
||||
if dict_key not in hardcoded_device_values:
|
||||
raise RuntimeError(
|
||||
f"No data for sm_{self.compute_capability} and {data_type}."
|
||||
)
|
||||
|
||||
return hardcoded_device_values[dict_key]
|
||||
|
||||
def get_tflops_per_second(
|
||||
self, data_type: dtype, use_tensor_cores: bool = True
|
||||
) -> float:
|
||||
num_sms = self.device_properties.multi_processor_count
|
||||
clock_rate = self.device_properties.clock_rate # KHz
|
||||
|
||||
fma_per_cycle = 0
|
||||
if use_tensor_cores:
|
||||
fma_per_cycle = self.get_fma_per_cycle_per_sm_tensor_cores(data_type)
|
||||
else:
|
||||
fma_per_cycle = self.get_fma_per_cycle_per_sm_cuda_cores(data_type)
|
||||
|
||||
# 1 FMA counts as 2 floating point operations
|
||||
# Clock rate is in KHz
|
||||
tflops_per_second = num_sms * fma_per_cycle * 2 * clock_rate / 1e9
|
||||
return tflops_per_second
|
||||
|
||||
def get_memory_bandwidth_Bps(self) -> int:
|
||||
# DRAM devices are Double-Data which means they provide an output at both fronts of
|
||||
# a clock beat
|
||||
bus_bytes_per_cycle = int(2 * self.device_properties.memory_bus_width / 8)
|
||||
mem_clock_rate_Hz = self.device_properties.memory_clock_rate * 1000
|
||||
bytes_per_second = bus_bytes_per_cycle * mem_clock_rate_Hz
|
||||
return bytes_per_second
|
||||
|
||||
def get_shared_memory_bandwidth_Bps(self) -> int:
|
||||
# Each warp can LD or ST 32 x 4 bytes per cycle. To calculate the
|
||||
# device's throughput we need to multiply with frequency and number of SMs.
|
||||
num_sms = self.device_properties.multi_processor_count
|
||||
bytes_per_cycle_per_sm = 128
|
||||
bytes_per_cycle_per_device = num_sms * bytes_per_cycle_per_sm
|
||||
bytes_per_second = (
|
||||
bytes_per_cycle_per_device * self.device_properties.clock_rate * 1000
|
||||
)
|
||||
return bytes_per_second
|
||||
@@ -0,0 +1,73 @@
|
||||
from collections.abc import Callable
|
||||
|
||||
from torch._utils import CallbackRegistry
|
||||
|
||||
|
||||
EventCreationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
|
||||
"CUDA event creation"
|
||||
)
|
||||
EventDeletionCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
|
||||
"CUDA event deletion"
|
||||
)
|
||||
EventRecordCallbacks: "CallbackRegistry[int, int]" = CallbackRegistry(
|
||||
"CUDA event record"
|
||||
)
|
||||
EventWaitCallbacks: "CallbackRegistry[int, int]" = CallbackRegistry("CUDA event wait")
|
||||
MemoryAllocationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
|
||||
"CUDA memory allocation"
|
||||
)
|
||||
MemoryDeallocationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
|
||||
"CUDA memory deallocation"
|
||||
)
|
||||
StreamCreationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
|
||||
"CUDA stream creation"
|
||||
)
|
||||
DeviceSynchronizationCallbacks: "CallbackRegistry[[]]" = CallbackRegistry(
|
||||
"CUDA device synchronization"
|
||||
)
|
||||
StreamSynchronizationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
|
||||
"CUDA stream synchronization"
|
||||
)
|
||||
EventSynchronizationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
|
||||
"CUDA event synchronization"
|
||||
)
|
||||
|
||||
|
||||
def register_callback_for_event_creation(cb: Callable[[int], None]) -> None:
|
||||
EventCreationCallbacks.add_callback(cb)
|
||||
|
||||
|
||||
def register_callback_for_event_deletion(cb: Callable[[int], None]) -> None:
|
||||
EventDeletionCallbacks.add_callback(cb)
|
||||
|
||||
|
||||
def register_callback_for_event_record(cb: Callable[[int, int], None]) -> None:
|
||||
EventRecordCallbacks.add_callback(cb)
|
||||
|
||||
|
||||
def register_callback_for_event_wait(cb: Callable[[int, int], None]) -> None:
|
||||
EventWaitCallbacks.add_callback(cb)
|
||||
|
||||
|
||||
def register_callback_for_memory_allocation(cb: Callable[[int], None]) -> None:
|
||||
MemoryAllocationCallbacks.add_callback(cb)
|
||||
|
||||
|
||||
def register_callback_for_memory_deallocation(cb: Callable[[int], None]) -> None:
|
||||
MemoryDeallocationCallbacks.add_callback(cb)
|
||||
|
||||
|
||||
def register_callback_for_stream_creation(cb: Callable[[int], None]) -> None:
|
||||
StreamCreationCallbacks.add_callback(cb)
|
||||
|
||||
|
||||
def register_callback_for_device_synchronization(cb: Callable[[], None]) -> None:
|
||||
DeviceSynchronizationCallbacks.add_callback(cb)
|
||||
|
||||
|
||||
def register_callback_for_stream_synchronization(cb: Callable[[int], None]) -> None:
|
||||
StreamSynchronizationCallbacks.add_callback(cb)
|
||||
|
||||
|
||||
def register_callback_for_event_synchronization(cb: Callable[[int], None]) -> None:
|
||||
EventSynchronizationCallbacks.add_callback(cb)
|
||||
@@ -0,0 +1,434 @@
|
||||
"""Annotate CUDA graph kernel nodes during capture.
|
||||
|
||||
During CUDA graph capture, ``mark_kernels`` uses ``cudaGraphGetNodes``
|
||||
to count nodes before and after the wrapped region. Nodes at indices
|
||||
``[before, after)`` are the ones added within the scope. Each kernel
|
||||
or memcpy node found is annotated by its ``toolsId`` so it can later
|
||||
be matched to profiler trace events.
|
||||
|
||||
The annotations can be pickled and later merged into a Chrome profiler
|
||||
trace using ``torch.cuda._annotate_cuda_graph_trace``.
|
||||
|
||||
Requires ``cuda.bindings`` package and a CUDA driver that supports
|
||||
``cudaGraphNodeGetToolsId`` (CUDA >= 13.1 or appropriate cuda-compat).
|
||||
When unavailable, ``mark_kernels`` silently becomes a no-op.
|
||||
|
||||
Usage during capture::
|
||||
|
||||
from torch.cuda._graph_annotations import (
|
||||
enable_annotations,
|
||||
mark_kernels,
|
||||
resolve_pending_annotations,
|
||||
remap_to_exec_graph,
|
||||
)
|
||||
|
||||
enable_annotations()
|
||||
|
||||
with torch.cuda.graph(graph):
|
||||
with mark_kernels("phase_A"):
|
||||
y = workload_a(x)
|
||||
with mark_kernels("phase_B"):
|
||||
z = workload_b(y)
|
||||
resolve_pending_annotations()
|
||||
|
||||
remap_to_exec_graph(graph)
|
||||
"""
|
||||
|
||||
from collections import defaultdict
|
||||
from contextlib import contextmanager
|
||||
from logging import getLogger
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.cuda._utils import _check_cuda_bindings, _HAS_CUDA_BINDINGS
|
||||
|
||||
|
||||
try:
|
||||
from cuda.bindings import ( # pyrefly: ignore[missing-import]
|
||||
runtime as _cuda_runtime,
|
||||
)
|
||||
except ImportError:
|
||||
_cuda_runtime = None # type: ignore[assignment]
|
||||
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
|
||||
# Tri-state: None = not probed, True = available, False = unavailable.
|
||||
# Deferred to first use to avoid premature CUDA initialization.
|
||||
_tools_id_available: bool | None = None
|
||||
|
||||
# Global kill switch. When False, mark_kernels and mark_stream are no-ops.
|
||||
_annotations_enabled: bool = False
|
||||
|
||||
|
||||
def enable_annotations() -> None:
|
||||
"""Enable kernel annotation recording."""
|
||||
global _annotations_enabled
|
||||
_annotations_enabled = True
|
||||
|
||||
|
||||
def disable_annotations() -> None:
|
||||
"""Disable kernel annotation recording."""
|
||||
global _annotations_enabled
|
||||
_annotations_enabled = False
|
||||
|
||||
|
||||
def _is_tools_id_unavailable() -> bool:
|
||||
"""Return True if we already know cudaGraphNodeGetToolsId is missing."""
|
||||
if not _HAS_CUDA_BINDINGS:
|
||||
return True
|
||||
if _tools_id_available is False:
|
||||
return True
|
||||
if not hasattr(_cuda_runtime, "cudaGraphNodeGetToolsId"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _get_tools_id(node: Any) -> int | None:
|
||||
"""Return the toolsId for a graph node, or None if unavailable."""
|
||||
global _tools_id_available
|
||||
if _tools_id_available is None:
|
||||
try:
|
||||
tools_id = _check_cuda_bindings(
|
||||
_cuda_runtime.cudaGraphNodeGetToolsId( # pyrefly: ignore[missing-attribute]
|
||||
node
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
_tools_id_available = False
|
||||
logger.info(
|
||||
"cudaGraphNodeGetToolsId not available; "
|
||||
"CUDA graph kernel annotations will be disabled"
|
||||
)
|
||||
return None
|
||||
_tools_id_available = True
|
||||
return tools_id
|
||||
return _check_cuda_bindings(
|
||||
_cuda_runtime.cudaGraphNodeGetToolsId( # pyrefly: ignore[missing-attribute]
|
||||
node
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _get_capture_graph(stream: Any) -> Any:
|
||||
"""Return the graph handle for the active capture, or None."""
|
||||
status, _id, graph, _deps, _edge_data, _num_deps = _check_cuda_bindings(
|
||||
_cuda_runtime.cudaStreamGetCaptureInfo( # pyrefly: ignore[missing-attribute]
|
||||
stream
|
||||
)
|
||||
)
|
||||
if (
|
||||
status
|
||||
!= _cuda_runtime.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive # pyrefly: ignore[missing-attribute]
|
||||
):
|
||||
return None
|
||||
return graph
|
||||
|
||||
|
||||
def _get_node_count(graph: Any) -> int:
|
||||
"""Return the number of nodes currently in the graph."""
|
||||
_, num = _check_cuda_bindings(
|
||||
_cuda_runtime.cudaGraphGetNodes( # pyrefly: ignore[missing-attribute]
|
||||
graph, numNodes=0
|
||||
)
|
||||
)
|
||||
return num
|
||||
|
||||
|
||||
# toolsId -> list of annotation objects.
|
||||
_kernel_annotations: defaultdict[int, list[Any]] = defaultdict(list)
|
||||
|
||||
# Node types we annotate. Initialized lazily to avoid touching cuda.bindings
|
||||
# at import time.
|
||||
_ANNOTATABLE_TYPES: set[Any] | None = None
|
||||
|
||||
|
||||
def _get_annotatable_types() -> set[Any]:
|
||||
global _ANNOTATABLE_TYPES
|
||||
if _ANNOTATABLE_TYPES is None:
|
||||
_ANNOTATABLE_TYPES = {
|
||||
_cuda_runtime.cudaGraphNodeType.cudaGraphNodeTypeKernel, # pyrefly: ignore[missing-attribute]
|
||||
_cuda_runtime.cudaGraphNodeType.cudaGraphNodeTypeMemcpy, # pyrefly: ignore[missing-attribute]
|
||||
}
|
||||
return _ANNOTATABLE_TYPES
|
||||
|
||||
|
||||
# Pending scopes: (annotation, start_node_index, end_node_index).
|
||||
_pending_scopes: list[tuple[Any, int, int]] = []
|
||||
|
||||
# Graph handle saved during capture for post-capture resolution.
|
||||
_capture_graph: Any = None
|
||||
|
||||
# Capture graph ID saved by resolve_pending_annotations for remap_to_exec_graph.
|
||||
_last_capture_graph_id: int | None = None
|
||||
|
||||
|
||||
@contextmanager # type: ignore[arg-type]
|
||||
def mark_kernels(annotation: str | dict[str, Any]):
|
||||
"""Context manager that records node index ranges for later annotation.
|
||||
|
||||
During capture, calls ``cudaGraphGetNodes`` to count graph nodes before
|
||||
and after the scope. Nodes at indices ``[before, after)`` were added
|
||||
inside the scope. After capture, ``resolve_pending_annotations``
|
||||
enumerates all nodes and annotates kernel/memcpy nodes in those ranges.
|
||||
|
||||
Must be called inside an active ``torch.cuda.graph()`` capture. If the
|
||||
current stream is not capturing, or if ``cudaGraphNodeGetToolsId`` is not
|
||||
available, the context manager is a no-op.
|
||||
|
||||
Args:
|
||||
annotation: Arbitrary object appended to the annotation list for
|
||||
every kernel/memcpy node whose index falls within this scope.
|
||||
"""
|
||||
if not _annotations_enabled or _is_tools_id_unavailable():
|
||||
yield
|
||||
return
|
||||
|
||||
if isinstance(annotation, str):
|
||||
annotation = {"str": annotation}
|
||||
|
||||
stream = _cuda_runtime.cudaStream_t( # pyrefly: ignore[missing-attribute]
|
||||
init_value=torch.cuda.current_stream().cuda_stream
|
||||
)
|
||||
graph = _get_capture_graph(stream)
|
||||
if graph is None:
|
||||
yield
|
||||
return
|
||||
|
||||
global _capture_graph
|
||||
_capture_graph = graph
|
||||
|
||||
start_count = _get_node_count(graph)
|
||||
|
||||
yield
|
||||
|
||||
end_count = _get_node_count(graph)
|
||||
|
||||
if end_count > start_count:
|
||||
_pending_scopes.append((annotation, start_count, end_count))
|
||||
|
||||
|
||||
def resolve_pending_annotations() -> None:
|
||||
"""Resolve pending scope index ranges into kernel annotations.
|
||||
|
||||
Enumerates all graph nodes and annotates kernel/memcpy nodes whose
|
||||
indices fall within recorded scope ranges. Must be called while still
|
||||
inside the ``torch.cuda.graph()`` capture context.
|
||||
"""
|
||||
global _capture_graph
|
||||
if not _pending_scopes:
|
||||
_capture_graph = None
|
||||
return
|
||||
|
||||
# Get a fresh graph handle from the active capture.
|
||||
stream = _cuda_runtime.cudaStream_t( # pyrefly: ignore[missing-attribute]
|
||||
init_value=torch.cuda.current_stream().cuda_stream
|
||||
)
|
||||
graph = _get_capture_graph(stream)
|
||||
if graph is None:
|
||||
graph = _capture_graph
|
||||
if graph is None:
|
||||
logger.warning("resolve_pending_annotations: no graph handle available")
|
||||
_pending_scopes.clear()
|
||||
return
|
||||
|
||||
try:
|
||||
num = _get_node_count(graph)
|
||||
if num == 0:
|
||||
_pending_scopes.clear()
|
||||
_capture_graph = None
|
||||
return
|
||||
|
||||
nodes, num = _check_cuda_bindings(
|
||||
_cuda_runtime.cudaGraphGetNodes( # pyrefly: ignore[missing-attribute]
|
||||
graph, numNodes=num
|
||||
)
|
||||
)
|
||||
|
||||
# Save capture graph ID for remap_to_exec_graph.
|
||||
global _last_capture_graph_id
|
||||
if num > 0:
|
||||
first_tid = _get_tools_id(nodes[0])
|
||||
_last_capture_graph_id = (first_tid >> 32) if first_tid else None
|
||||
|
||||
annotatable = _get_annotatable_types()
|
||||
|
||||
# Sort by (start, -end, -append_index). The append index encodes
|
||||
# nesting depth: inner context managers exit first, so they are
|
||||
# appended to _pending_scopes first (smaller index). Using
|
||||
# -append_index as tiebreaker ensures that for same-range scopes
|
||||
# the outer scope (larger index) sorts first and is pushed onto
|
||||
# the stack first, leaving the inner scope on top.
|
||||
sorted_scopes = sorted(
|
||||
(
|
||||
(ann, start, end, i)
|
||||
for i, (ann, start, end) in enumerate(_pending_scopes)
|
||||
),
|
||||
key=lambda s: (s[1], -s[2], -s[3]),
|
||||
)
|
||||
scope_ptr = 0
|
||||
active_stack: list[tuple[int, Any]] = [] # (end_idx, annotation)
|
||||
|
||||
for i in range(num):
|
||||
# Pop scopes whose range ended.
|
||||
while active_stack and active_stack[-1][0] <= i:
|
||||
active_stack.pop()
|
||||
|
||||
# Push scopes that start at or before this index.
|
||||
while scope_ptr < len(sorted_scopes) and sorted_scopes[scope_ptr][1] <= i:
|
||||
ann, _start_idx, end_idx, _idx = sorted_scopes[scope_ptr]
|
||||
if end_idx > i:
|
||||
active_stack.append((end_idx, ann))
|
||||
scope_ptr += 1
|
||||
|
||||
if not active_stack:
|
||||
continue
|
||||
|
||||
node = nodes[i]
|
||||
node_type = _check_cuda_bindings(
|
||||
_cuda_runtime.cudaGraphNodeGetType( # pyrefly: ignore[missing-attribute]
|
||||
node
|
||||
)
|
||||
)
|
||||
if node_type not in annotatable:
|
||||
continue
|
||||
|
||||
tools_id = _get_tools_id(node)
|
||||
if tools_id is None:
|
||||
logger.warning(
|
||||
"resolve_pending_annotations: toolsId unavailable, aborting"
|
||||
)
|
||||
_pending_scopes.clear()
|
||||
_capture_graph = None
|
||||
return
|
||||
|
||||
if len(active_stack) == 1:
|
||||
_kernel_annotations[tools_id].append(active_stack[0][1])
|
||||
else:
|
||||
# Merge all active scopes into one dict. Inner scopes sit
|
||||
# on top of the stack. Iterating reversed (inner first)
|
||||
# with setdefault lets the inner scope's values win for
|
||||
# overlapping keys (e.g. name, stream) while outer scopes
|
||||
# fill in any missing keys.
|
||||
merged: dict[str, Any] = {}
|
||||
for _, ann in reversed(active_stack):
|
||||
if isinstance(ann, dict):
|
||||
for ak, av in ann.items():
|
||||
merged.setdefault(ak, av)
|
||||
else:
|
||||
merged.setdefault("name", ann)
|
||||
_kernel_annotations[tools_id].append(merged)
|
||||
except Exception:
|
||||
logger.exception("resolve_pending_annotations failed")
|
||||
finally:
|
||||
_pending_scopes.clear()
|
||||
_capture_graph = None
|
||||
|
||||
|
||||
def remap_to_exec_graph(torch_cuda_graph: torch.cuda.CUDAGraph) -> None:
|
||||
"""Remap annotation keys from capture graph ID to exec graph ID.
|
||||
|
||||
During capture, toolsId encodes the capture graph's ID in the upper
|
||||
32 bits. After instantiation, the profiler uses the exec graph's ID.
|
||||
This function rewrites the keys so annotations match the trace.
|
||||
|
||||
Must be called after the ``torch.cuda.graph()`` context exits.
|
||||
"""
|
||||
if not _kernel_annotations:
|
||||
return
|
||||
|
||||
exec_handle = _cuda_runtime.cudaGraphExec_t( # pyrefly: ignore[missing-attribute]
|
||||
init_value=torch_cuda_graph.raw_cuda_graph_exec()
|
||||
)
|
||||
exec_graph_id = _check_cuda_bindings(
|
||||
_cuda_runtime.cudaGraphExecGetId( # pyrefly: ignore[missing-attribute]
|
||||
exec_handle
|
||||
)
|
||||
)
|
||||
|
||||
# Only remap annotations from the most recent capture graph.
|
||||
# Previously remapped annotations (from earlier captures) keep their
|
||||
# correct exec graph IDs.
|
||||
capture_graph_id = _last_capture_graph_id
|
||||
remapped: dict[int, list[Any]] = {}
|
||||
for tools_id, ann_list in _kernel_annotations.items():
|
||||
graph_id = tools_id >> 32
|
||||
if capture_graph_id is not None and graph_id != capture_graph_id:
|
||||
# Belongs to a different graph — keep as-is.
|
||||
remapped[tools_id] = ann_list
|
||||
continue
|
||||
node_id = tools_id & 0xFFFFFFFF
|
||||
new_tools_id = (exec_graph_id << 32) | node_id
|
||||
if new_tools_id in remapped:
|
||||
remapped[new_tools_id].extend(ann_list)
|
||||
else:
|
||||
remapped[new_tools_id] = list(ann_list)
|
||||
|
||||
_kernel_annotations.clear()
|
||||
_kernel_annotations.update(remapped)
|
||||
|
||||
|
||||
def get_kernel_annotations() -> dict[int, list[Any]]:
|
||||
"""Return the current kernel annotation map (toolsId -> annotations)."""
|
||||
return _kernel_annotations
|
||||
|
||||
|
||||
def clear_kernel_annotations() -> None:
|
||||
"""Clear all recorded kernel annotations and pending scopes."""
|
||||
global _capture_graph
|
||||
_kernel_annotations.clear()
|
||||
_pending_scopes.clear()
|
||||
_capture_graph = None
|
||||
|
||||
|
||||
# Counter-based stream ID registry. IDs start at 60 (above the highest
|
||||
# observed non-graphed CUDA stream ID) so every assigned lane is visually
|
||||
# distinct in Perfetto and doesn't collide with real streams.
|
||||
_stream_id_counter: int = 60
|
||||
_stream_id_map: dict[int, int] = {}
|
||||
|
||||
|
||||
def _get_stream_id(stream: torch.cuda.Stream) -> int:
|
||||
"""Return a small, stable stream ID for the given CUDA stream."""
|
||||
global _stream_id_counter
|
||||
key = stream.cuda_stream
|
||||
if key not in _stream_id_map:
|
||||
_stream_id_map[key] = _stream_id_counter
|
||||
_stream_id_counter += 1
|
||||
return _stream_id_map[key]
|
||||
|
||||
|
||||
def get_stream_for_pg(pg_key: str) -> int:
|
||||
"""Return a unique stream ID for the given process group key."""
|
||||
global _stream_id_counter
|
||||
if pg_key not in _stream_id_map:
|
||||
_stream_id_map[pg_key] = _stream_id_counter # type: ignore[assignment]
|
||||
_stream_id_counter += 1
|
||||
return _stream_id_map[pg_key] # type: ignore[return-value]
|
||||
|
||||
|
||||
@contextmanager # type: ignore[arg-type]
|
||||
def mark_stream(stream: torch.cuda.Stream, annotation: str | dict[str, Any]):
|
||||
"""Switch to stream, inject its ID into annotation, and mark kernels.
|
||||
|
||||
If *stream* is already the current stream, no stream switch or stream ID
|
||||
injection happens — the kernels stay on whatever stream is active (which
|
||||
keeps the trace faithful when e.g. FSDP uses the current stream for
|
||||
copy-in instead of a separate one).
|
||||
"""
|
||||
if not _annotations_enabled:
|
||||
with torch.cuda.stream(stream):
|
||||
yield
|
||||
return
|
||||
if stream.cuda_stream == torch.cuda.current_stream().cuda_stream:
|
||||
with mark_kernels(annotation):
|
||||
yield
|
||||
else:
|
||||
if isinstance(annotation, str):
|
||||
annotation = {"str": annotation}
|
||||
if isinstance(annotation, dict):
|
||||
annotation["stream"] = _get_stream_id(stream)
|
||||
with torch.cuda.stream(stream):
|
||||
with mark_kernels(annotation):
|
||||
yield
|
||||
@@ -0,0 +1,802 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import operator
|
||||
import os
|
||||
import pickle
|
||||
import subprocess
|
||||
import sys
|
||||
import warnings
|
||||
from functools import lru_cache
|
||||
from itertools import groupby
|
||||
from typing import Any
|
||||
|
||||
|
||||
cache = lru_cache(None)
|
||||
|
||||
__all__ = ["format_flamegraph", "segments", "memory", "compare"]
|
||||
|
||||
|
||||
def _frame_fmt(f, full_filename=False):
|
||||
i = f["line"]
|
||||
fname = f["filename"]
|
||||
if not full_filename:
|
||||
fname = fname.split("/")[-1]
|
||||
func = f["name"]
|
||||
return f"{fname}:{i}:{func}"
|
||||
|
||||
|
||||
@cache
|
||||
def _frame_filter(name, filename):
|
||||
omit_functions = [
|
||||
"unwind::unwind",
|
||||
"CapturedTraceback::gather",
|
||||
"gather_with_cpp",
|
||||
"_start",
|
||||
"__libc_start_main",
|
||||
"PyEval_",
|
||||
"PyObject_",
|
||||
"PyFunction_",
|
||||
]
|
||||
omit_filenames = [
|
||||
"core/boxing",
|
||||
"/Register",
|
||||
"/Redispatch",
|
||||
"pythonrun.c",
|
||||
"Modules/main.c",
|
||||
"Objects/call.c",
|
||||
"Objects/methodobject.c",
|
||||
"pycore_ceval.h",
|
||||
"ceval.c",
|
||||
"cpython/abstract.h",
|
||||
]
|
||||
for of in omit_functions:
|
||||
if of in name:
|
||||
return False
|
||||
for of in omit_filenames:
|
||||
if of in filename:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _frames_fmt(frames, full_filename=False, reverse=False):
|
||||
if reverse:
|
||||
frames = reversed(frames)
|
||||
return [
|
||||
_frame_fmt(f, full_filename)
|
||||
for f in frames
|
||||
if _frame_filter(f["name"], f["filename"])
|
||||
]
|
||||
|
||||
|
||||
def _block_extra_legacy(b):
|
||||
if "history" in b:
|
||||
frames = b["history"][0].get("frames", [])
|
||||
real_size = b["history"][0]["real_size"]
|
||||
else:
|
||||
real_size = b.get("requested_size", b["size"])
|
||||
frames = []
|
||||
return frames, real_size
|
||||
|
||||
|
||||
def _block_extra(b):
|
||||
if "frames" not in b:
|
||||
# old snapshot format made it more complicated to get frames/allocated size
|
||||
return _block_extra_legacy(b)
|
||||
return b["frames"], b["requested_size"]
|
||||
|
||||
|
||||
def format_flamegraph(flamegraph_lines, flamegraph_script=None):
|
||||
if flamegraph_script is None:
|
||||
cache_dir = os.path.expanduser("~/.cache/")
|
||||
os.makedirs(cache_dir, exist_ok=True)
|
||||
flamegraph_script = f"{cache_dir}/flamegraph.pl"
|
||||
if not os.path.exists(flamegraph_script):
|
||||
import tempfile
|
||||
import urllib.request
|
||||
|
||||
print(f"Downloading flamegraph.pl to: {flamegraph_script}")
|
||||
with tempfile.NamedTemporaryFile(mode="wb", suffix=".pl") as f:
|
||||
urllib.request.urlretrieve(
|
||||
"https://raw.githubusercontent.com/brendangregg/FlameGraph/master/flamegraph.pl",
|
||||
f.name,
|
||||
)
|
||||
try:
|
||||
os.chmod(f.name, 0o755)
|
||||
os.rename(f.name, flamegraph_script)
|
||||
except OSError: # noqa: B001,E722
|
||||
# Ok to skip, the file will be removed by tempfile
|
||||
pass
|
||||
args = [flamegraph_script, "--countname", "bytes"]
|
||||
with subprocess.Popen(
|
||||
args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, encoding="utf-8"
|
||||
) as p:
|
||||
if p.stdin is None:
|
||||
raise AssertionError("p.stdin is None")
|
||||
if p.stdout is None:
|
||||
raise AssertionError("p.stdout is None")
|
||||
p.stdin.write(flamegraph_lines)
|
||||
p.stdin.close()
|
||||
result = p.stdout.read()
|
||||
p.stdout.close()
|
||||
p.wait()
|
||||
if p.wait() != 0:
|
||||
raise AssertionError(f"flamegraph process exited with code {p.wait()}")
|
||||
return result
|
||||
|
||||
|
||||
def _write_blocks(f, prefix, blocks):
|
||||
def frames_fragment(frames):
|
||||
if not frames:
|
||||
return "<non-python>"
|
||||
return ";".join(_frames_fmt(frames, reverse=True))
|
||||
|
||||
for b in blocks:
|
||||
if "history" not in b:
|
||||
frames, accounted_for_size = _block_extra(b)
|
||||
f.write(
|
||||
f"{prefix};{b['state']};{frames_fragment(frames)} {accounted_for_size}\n"
|
||||
)
|
||||
else:
|
||||
accounted_for_size = 0
|
||||
for h in b["history"]:
|
||||
sz = h["real_size"]
|
||||
accounted_for_size += sz
|
||||
if "frames" in h:
|
||||
frames = h["frames"]
|
||||
f.write(f"{prefix};{b['state']};{frames_fragment(frames)} {sz}\n")
|
||||
else:
|
||||
f.write(f"{prefix};{b['state']};<no-context> {sz}\n")
|
||||
gaps = b["size"] - accounted_for_size
|
||||
if gaps:
|
||||
f.write(f"{prefix};{b['state']};<gaps> {gaps}\n")
|
||||
|
||||
|
||||
def segments(snapshot, format_flamegraph=format_flamegraph):
|
||||
f = io.StringIO()
|
||||
for seg in snapshot["segments"]:
|
||||
prefix = f"stream_{seg['stream']};seg_{seg['address']}"
|
||||
_write_blocks(f, prefix, seg["blocks"])
|
||||
return format_flamegraph(f.getvalue())
|
||||
|
||||
|
||||
def memory(snapshot, format_flamegraph=format_flamegraph):
|
||||
f = io.StringIO()
|
||||
for seg in snapshot["segments"]:
|
||||
prefix = f"stream_{seg['stream']}"
|
||||
_write_blocks(f, prefix, seg["blocks"])
|
||||
return format_flamegraph(f.getvalue())
|
||||
|
||||
|
||||
def compare(before, after, format_flamegraph=format_flamegraph):
|
||||
def _seg_key(seg):
|
||||
return (seg["address"], seg["total_size"])
|
||||
|
||||
def _seg_info(seg):
|
||||
return f"stream_{seg['stream']};seg_{seg['address']}"
|
||||
|
||||
f = io.StringIO()
|
||||
|
||||
before_segs = {_seg_key(seg) for seg in before}
|
||||
after_segs = {_seg_key(seg) for seg in after}
|
||||
|
||||
print(f"only_before = {[a for a, _ in (before_segs - after_segs)]}")
|
||||
print(f"only_after = {[a for a, _ in (after_segs - before_segs)]}")
|
||||
|
||||
for seg in before:
|
||||
if _seg_key(seg) not in after_segs:
|
||||
_write_blocks(f, f"only_before;{_seg_info(seg)}", seg["blocks"])
|
||||
|
||||
for seg in after:
|
||||
if _seg_key(seg) not in before_segs:
|
||||
_write_blocks(f, f"only_after;{_seg_info(seg)}", seg["blocks"])
|
||||
|
||||
return format_flamegraph(f.getvalue())
|
||||
|
||||
|
||||
def _format_size(num):
|
||||
# https://stackoverflow.com/questions/1094841/get-human-readable-version-of-file-size
|
||||
for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
|
||||
if abs(num) < 1024.0:
|
||||
return f"{num:3.1f}{unit}B"
|
||||
num /= 1024.0
|
||||
return f"{num:.1f}YiB"
|
||||
|
||||
|
||||
class Bytes:
|
||||
def __init__(self, value):
|
||||
self.value = value
|
||||
|
||||
def __add__(self, rhs):
|
||||
return Bytes(self.value + rhs)
|
||||
|
||||
def __repr__(self):
|
||||
return _format_size(self.value)
|
||||
|
||||
|
||||
def calc_active(seg):
|
||||
return sum(b["size"] for b in seg["blocks"] if b["state"] == "active_allocated")
|
||||
|
||||
|
||||
def _report_free(free_external, free_internal):
|
||||
total = free_external + free_internal
|
||||
suffix = ""
|
||||
if total != 0:
|
||||
pct = (free_internal / total) * 100
|
||||
suffix = f" ({pct:.1f}% internal)"
|
||||
return f"{Bytes(total)}{suffix}"
|
||||
|
||||
|
||||
PAGE_SIZE = 1024 * 1024 * 20
|
||||
legend = f"""\
|
||||
|
||||
Legend:
|
||||
[a ] - a segment in the allocator
|
||||
^-- a page {Bytes(PAGE_SIZE)} of memory in the segment
|
||||
a-z: pages filled with a single block's content
|
||||
' ': page is completely free
|
||||
*: page if completely full with multiple blocks
|
||||
0-9: page is partially full with tensors of multiple blocks (9 == 90% full)
|
||||
(X% internal) - of the free memory, X% is free because we rounded the size of the allocation.
|
||||
"""
|
||||
|
||||
|
||||
def segsum(data):
|
||||
r"""Visually reports how the allocator has filled its segments.
|
||||
|
||||
This printout can help debug fragmentation issues since free fragments
|
||||
will appear as gaps in this printout. The amount of free space is reported
|
||||
for each segment.
|
||||
We distinguish between internal free memory which occurs because the
|
||||
allocator rounds the allocation size, and external free memory, which are
|
||||
the gaps between allocations in a segment.
|
||||
Args:
|
||||
data: snapshot dictionary created from _snapshot()
|
||||
"""
|
||||
out = io.StringIO()
|
||||
out.write(f"Summary of segments >= {Bytes(PAGE_SIZE)} in size\n")
|
||||
total_reserved = 0
|
||||
total_allocated = 0
|
||||
free_external = 0
|
||||
free_internal = 0
|
||||
for seg in sorted(
|
||||
data["segments"], key=lambda x: (x["total_size"], calc_active(x))
|
||||
):
|
||||
total_reserved += seg["total_size"]
|
||||
|
||||
seg_free_external = 0
|
||||
seg_free_internal = 0
|
||||
seg_allocated = 0
|
||||
all_ranges = []
|
||||
boffset = 0
|
||||
for b in seg["blocks"]:
|
||||
active = b["state"] == "active_allocated"
|
||||
if active:
|
||||
_, allocated_size = _block_extra(b)
|
||||
all_ranges.append((boffset, allocated_size, True))
|
||||
seg_allocated += allocated_size
|
||||
seg_free_internal += b["size"] - allocated_size
|
||||
else:
|
||||
seg_free_external += b["size"]
|
||||
|
||||
boffset += b["size"]
|
||||
|
||||
total_allocated += seg_allocated
|
||||
free_external += seg_free_external
|
||||
free_internal += seg_free_internal
|
||||
|
||||
nseg = (seg["total_size"] - 1) // PAGE_SIZE + 1
|
||||
occupied = [" " for _ in range(nseg)]
|
||||
frac = [0.0 for _ in range(nseg)]
|
||||
active_size = 0
|
||||
for i, (start_, size, active) in enumerate(all_ranges):
|
||||
active_size += size
|
||||
finish_ = start_ + size
|
||||
start = start_ // PAGE_SIZE
|
||||
finish = (finish_ - 1) // PAGE_SIZE + 1
|
||||
m = chr(ord("a" if active else "A") + (i % 26))
|
||||
for j in range(start, finish):
|
||||
s = max(start_, j * PAGE_SIZE)
|
||||
e = min(finish_, (j + 1) * PAGE_SIZE)
|
||||
frac[j] += (e - s) / PAGE_SIZE
|
||||
if occupied[j] != " ":
|
||||
occupied[j] = "0123456789*"[int(frac[j] * 10)]
|
||||
else:
|
||||
occupied[j] = m
|
||||
stream = "" if seg["stream"] == 0 else f", stream_{seg['stream']}"
|
||||
body = "".join(occupied)
|
||||
if seg_free_external + seg_free_internal + seg_allocated != seg["total_size"]:
|
||||
raise AssertionError(
|
||||
f"Segment size mismatch: {seg_free_external} + {seg_free_internal} + {seg_allocated} != {seg['total_size']}"
|
||||
)
|
||||
stream = f" stream_{seg['stream']}" if seg["stream"] != 0 else ""
|
||||
if seg["total_size"] >= PAGE_SIZE:
|
||||
out.write(
|
||||
f"[{body}] {Bytes(seg['total_size'])} allocated, "
|
||||
f"{_report_free(seg_free_external, seg_free_internal)} free{stream}\n"
|
||||
)
|
||||
out.write(f"segments: {len(data['segments'])}\n")
|
||||
out.write(f"total_reserved: {Bytes(total_reserved)}\n")
|
||||
out.write(f"total_allocated: {Bytes(total_allocated)}\n")
|
||||
out.write(f"total_free: {_report_free(free_external, free_internal)}\n")
|
||||
out.write(legend)
|
||||
if free_internal + free_external + total_allocated != total_reserved:
|
||||
raise AssertionError(
|
||||
f"Memory accounting error: {free_internal} + {free_external} + {total_allocated} != {total_reserved}"
|
||||
)
|
||||
return out.getvalue()
|
||||
|
||||
|
||||
def trace(data):
|
||||
out = io.StringIO()
|
||||
|
||||
def format(entries):
|
||||
segment_intervals: list = []
|
||||
segment_addr_to_name = {}
|
||||
allocation_addr_to_name = {}
|
||||
|
||||
free_names: list = []
|
||||
next_name = 0
|
||||
|
||||
def _name():
|
||||
nonlocal next_name
|
||||
if free_names:
|
||||
return free_names.pop()
|
||||
r, m = next_name // 26, next_name % 26
|
||||
next_name += 1
|
||||
return f"{chr(ord('a') + m)}{'' if r == 0 else r}"
|
||||
|
||||
def find_segment(addr):
|
||||
for name, saddr, size in segment_intervals:
|
||||
if addr >= saddr and addr < saddr + size:
|
||||
return name, saddr
|
||||
for i, seg in enumerate(data["segments"]):
|
||||
saddr = seg["address"]
|
||||
size = seg["allocated_size"]
|
||||
if addr >= saddr and addr < saddr + size:
|
||||
return f"seg_{i}", saddr
|
||||
return None, None
|
||||
|
||||
count = 0
|
||||
out.write(f"{len(entries)} entries\n")
|
||||
|
||||
total_reserved = 0
|
||||
for seg in data["segments"]:
|
||||
total_reserved += seg["total_size"]
|
||||
|
||||
for count, e in enumerate(entries):
|
||||
if e["action"] == "alloc":
|
||||
addr, size = e["addr"], e["size"]
|
||||
n = _name()
|
||||
seg_name, seg_addr = find_segment(addr)
|
||||
if seg_name is None:
|
||||
seg_name = "MEM"
|
||||
offset = addr
|
||||
else:
|
||||
offset = addr - seg_addr
|
||||
out.write(f"{n} = {seg_name}[{offset}:{Bytes(size)}]\n")
|
||||
allocation_addr_to_name[addr] = (n, size, count)
|
||||
count += size
|
||||
elif e["action"] == "free_requested":
|
||||
addr, size = e["addr"], e["size"]
|
||||
name, _, _ = allocation_addr_to_name.get(addr, (addr, None, None))
|
||||
out.write(f"del {name} # {Bytes(size)}\n")
|
||||
elif e["action"] == "free_completed":
|
||||
addr, size = e["addr"], e["size"]
|
||||
count -= size
|
||||
name, _, _ = allocation_addr_to_name.get(addr, (addr, None, None))
|
||||
out.write(f"# free completed for {name} {Bytes(size)}\n")
|
||||
if name in allocation_addr_to_name:
|
||||
free_names.append(name)
|
||||
del allocation_addr_to_name[name]
|
||||
elif e["action"] == "segment_alloc":
|
||||
addr, size = e["addr"], e["size"]
|
||||
name = _name()
|
||||
out.write(f"{name} = cudaMalloc({addr}, {Bytes(size)})\n")
|
||||
segment_intervals.append((name, addr, size))
|
||||
segment_addr_to_name[addr] = name
|
||||
elif e["action"] == "segment_free":
|
||||
addr, size = e["addr"], e["size"]
|
||||
name = segment_addr_to_name.get(addr, addr)
|
||||
out.write(f"cudaFree({name}) # {Bytes(size)}\n")
|
||||
if name in segment_addr_to_name:
|
||||
free_names.append(name)
|
||||
del segment_addr_to_name[name]
|
||||
elif e["action"] == "oom":
|
||||
size = e["size"]
|
||||
free = e["device_free"]
|
||||
out.write(
|
||||
f"raise OutOfMemoryError # {Bytes(size)} requested, {Bytes(free)} free in CUDA\n"
|
||||
)
|
||||
else:
|
||||
out.write(f"{e}\n")
|
||||
out.write(f"TOTAL MEM: {Bytes(count)}")
|
||||
|
||||
for i, d in enumerate(data["device_traces"]):
|
||||
if d:
|
||||
out.write(f"Device {i} ----------------\n")
|
||||
format(d)
|
||||
return out.getvalue()
|
||||
|
||||
|
||||
_memory_viz_template = r"""
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
</head>
|
||||
<body>
|
||||
<script type="module">
|
||||
import {add_local_files} from "https://cdn.jsdelivr.net/gh/pytorch/pytorch@main/torch/utils/viz/MemoryViz.js"
|
||||
const local_files = $SNAPSHOT
|
||||
add_local_files(local_files, $VIZ_KIND)
|
||||
</script>
|
||||
</body>
|
||||
"""
|
||||
|
||||
|
||||
def _format_viz(data, viz_kind, device):
|
||||
if device is not None:
|
||||
warnings.warn(
|
||||
"device argument is deprecated, plots now contain all device",
|
||||
FutureWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
buffer = pickle.dumps(data)
|
||||
buffer += b"\x00" * (3 - len(buffer) % 3)
|
||||
# Encode the buffer with base64
|
||||
encoded_buffer = base64.b64encode(buffer).decode("utf-8")
|
||||
|
||||
json_format = json.dumps([{"name": "snapshot.pickle", "base64": encoded_buffer}])
|
||||
return _memory_viz_template.replace("$VIZ_KIND", repr(viz_kind)).replace(
|
||||
"$SNAPSHOT", json_format
|
||||
)
|
||||
|
||||
|
||||
def filter_alloc_free_pairs(data):
|
||||
for dev_id in range(len(data["device_traces"])):
|
||||
# set of indexes of trace events for alloc-free pairs
|
||||
filterSet = set()
|
||||
# map from addr to index of alloc event
|
||||
allocMap = {}
|
||||
# set of addrs from free_requested events
|
||||
freeRequested = set()
|
||||
for idx, event in enumerate(data["device_traces"][dev_id]):
|
||||
if event["action"] == "alloc":
|
||||
allocMap[event["addr"]] = idx
|
||||
elif event["action"] == "free_requested":
|
||||
freeRequested.add(event["addr"])
|
||||
if allocMap.get(event["addr"]) is not None:
|
||||
filterSet.add(idx)
|
||||
filterSet.add(allocMap[event["addr"]])
|
||||
allocMap.pop(event["addr"])
|
||||
elif event["action"] == "free_completed":
|
||||
if event["addr"] in freeRequested:
|
||||
freeRequested.remove(event["addr"])
|
||||
filterSet.add(idx)
|
||||
else:
|
||||
print(f"free_completed without free_requested: {event}")
|
||||
|
||||
# Remove events whose index is in filterSet
|
||||
if filterSet:
|
||||
# Create a new list excluding events with indices in filterSet
|
||||
data["device_traces"][dev_id] = [
|
||||
event
|
||||
for idx, event in enumerate(data["device_traces"][dev_id])
|
||||
if idx not in filterSet
|
||||
]
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def trace_plot(data, device=None, plot_segments=False, filter_freed=False):
|
||||
"""Generate a visualization over time of the memory usage recorded by the trace as an html file.
|
||||
|
||||
Args:
|
||||
data: Memory snapshot as generated from torch.cuda.memory._snapshot()
|
||||
device (torch.device, optional): Generate the trace for this device, needed if multiple devices have allocations.
|
||||
plot_segments (bool, optional): Plots memory returned from cudaMalloc, rather than individual allocations.
|
||||
Defaults to False.
|
||||
filter_freed (bool, optional): Filter out alloc-free paired events to only plot allocations that are not freed yet.
|
||||
Defaults to False to plot all trace events.
|
||||
|
||||
Returns:
|
||||
str: HTML of visualization
|
||||
"""
|
||||
if filter_freed:
|
||||
data = filter_alloc_free_pairs(data)
|
||||
|
||||
return _format_viz(
|
||||
data,
|
||||
"Active Memory Timeline"
|
||||
if not plot_segments
|
||||
else "Active Cached Memory Timeline",
|
||||
device,
|
||||
)
|
||||
|
||||
|
||||
def _profile_to_snapshot(profile):
|
||||
import torch
|
||||
from torch._C._profiler import _EventType
|
||||
from torch.profiler._memory_profiler import Action, TensorKey
|
||||
|
||||
memory_profile = profile._memory_profile()
|
||||
|
||||
allocation_stacks = {}
|
||||
for event in memory_profile._op_tree.sorted_nodes:
|
||||
if event.tag == _EventType.Allocation:
|
||||
parent = event.parent
|
||||
python_parents = []
|
||||
while parent:
|
||||
if parent.tag in (_EventType.PyCall, _EventType.PyCCall):
|
||||
python_parents.append(parent)
|
||||
parent = parent.parent
|
||||
key = TensorKey.from_allocation(event.extra_fields)
|
||||
|
||||
# Corner case: If allocation doesn't have an ID (can't prove it was used as a Tensor)
|
||||
# key will be None. I should add some way to identify these, I just haven't yet.
|
||||
if key and event.extra_fields.alloc_size > 0:
|
||||
allocation_stacks[key] = python_parents
|
||||
|
||||
device_count = torch.cuda.device_count()
|
||||
snapshot: dict[str, list[Any]] = {
|
||||
"device_traces": [[] for _ in range(device_count + 1)],
|
||||
"segments": [
|
||||
{
|
||||
"device": device,
|
||||
"address": None,
|
||||
"total_size": 0,
|
||||
"stream": 0,
|
||||
"blocks": [],
|
||||
}
|
||||
for device in range(device_count + 1)
|
||||
],
|
||||
}
|
||||
|
||||
def to_device(device):
|
||||
if device.type == "cuda":
|
||||
return device.index
|
||||
else:
|
||||
return device_count
|
||||
|
||||
def allocate(size, tensor_key, version, during_trace=True):
|
||||
device = to_device(tensor_key.device)
|
||||
addr = tensor_key.storage.ptr
|
||||
|
||||
seg = snapshot["segments"][device] # type: ignore[index]
|
||||
if seg["address"] is None or seg["address"] > addr:
|
||||
seg["address"] = addr
|
||||
seg["total_size"] = max(
|
||||
seg["total_size"], addr + size
|
||||
) # record max addr for now, we will make it the size later
|
||||
category = memory_profile._categories.get(tensor_key, version)
|
||||
category = category.name.lower() if category is not None else "unknown"
|
||||
stack = allocation_stacks.get(tensor_key, ())
|
||||
stack = [{"filename": "none", "line": 0, "name": p.name} for p in stack]
|
||||
r = {
|
||||
"action": "alloc",
|
||||
"addr": addr,
|
||||
"size": size,
|
||||
"stream": 0,
|
||||
"frames": stack,
|
||||
"category": category,
|
||||
}
|
||||
if during_trace:
|
||||
snapshot["device_traces"][device].append(r)
|
||||
return r
|
||||
|
||||
def free(alloc, device):
|
||||
for e in ("free_requested", "free_completed"):
|
||||
snapshot["device_traces"][device].append(
|
||||
{
|
||||
"action": e,
|
||||
"addr": alloc["addr"],
|
||||
"size": alloc["size"],
|
||||
"stream": 0,
|
||||
"frames": alloc["frames"],
|
||||
}
|
||||
)
|
||||
|
||||
kv_to_elem = {}
|
||||
|
||||
# create the device trace
|
||||
for _time, action, (tensor_key, version), size in memory_profile.timeline:
|
||||
if not isinstance(tensor_key, TensorKey):
|
||||
continue
|
||||
if action == Action.CREATE:
|
||||
kv_to_elem[(tensor_key, version)] = allocate(size, tensor_key, version)
|
||||
elif action == Action.DESTROY:
|
||||
free(kv_to_elem.pop((tensor_key, version)), to_device(tensor_key.device))
|
||||
elif action == Action.INCREMENT_VERSION:
|
||||
free(kv_to_elem.pop((tensor_key, version)), to_device(tensor_key.device))
|
||||
kv_to_elem[(tensor_key, version + 1)] = allocate(
|
||||
size, tensor_key, version + 1
|
||||
)
|
||||
elif action == Action.PREEXISTING:
|
||||
kv_to_elem[(tensor_key, version)] = allocate(
|
||||
size, tensor_key, version, during_trace=False
|
||||
)
|
||||
|
||||
# create the final snapshot state
|
||||
blocks_at_end = [
|
||||
(to_device(tensor_key.device), event["addr"], event["size"], event["frames"])
|
||||
for (tensor_key, version), event in kv_to_elem.items()
|
||||
]
|
||||
for device, blocks in groupby(sorted(blocks_at_end), key=operator.itemgetter(0)):
|
||||
seg = snapshot["segments"][device] # type: ignore[index]
|
||||
last_addr = seg["address"]
|
||||
for _, addr, size, frames in blocks:
|
||||
if last_addr < addr:
|
||||
seg["blocks"].append({"size": addr - last_addr, "state": "inactive"})
|
||||
seg["blocks"].append(
|
||||
{
|
||||
"size": size,
|
||||
"state": "active_allocated",
|
||||
"requested_size": size,
|
||||
"frames": frames,
|
||||
}
|
||||
)
|
||||
last_addr = addr + size
|
||||
if last_addr < seg["total_size"]:
|
||||
seg["blocks"].append(
|
||||
{"size": seg["total_size"] - last_addr, "state": "inactive"}
|
||||
)
|
||||
|
||||
snapshot["segments"] = [seg for seg in snapshot["segments"] if seg["blocks"]] # type: ignore[attr-defined]
|
||||
for seg in snapshot["segments"]: # type: ignore[attr-defined, name-defined, no-redef]
|
||||
seg["total_size"] -= seg["address"]
|
||||
if not seg["blocks"]:
|
||||
seg["blocks"].append({"size": seg["total_size"], "state": "inactive"})
|
||||
|
||||
return snapshot
|
||||
|
||||
|
||||
def profile_plot(profile, device=None):
|
||||
"""Generate a visualization over time of the memory usage recorded by kineto memory profiling as an html file.
|
||||
|
||||
Args:
|
||||
profile: profile as generated by `torch.profiler.profile(profile_memory=True)`
|
||||
device (torch.device, optional): Generate the trace for this device, needed if multiple devices have allocations.
|
||||
|
||||
Returns:
|
||||
str: HTML of visualization
|
||||
"""
|
||||
snapshot = _profile_to_snapshot(profile)
|
||||
return _format_viz(snapshot, "Active Memory Timeline", device)
|
||||
|
||||
|
||||
def segment_plot(data: Any, device=None):
|
||||
return _format_viz(data, "Allocator State History", device)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import os.path
|
||||
|
||||
thedir = os.path.realpath(os.path.dirname(__file__))
|
||||
if thedir in sys.path:
|
||||
# otherwise we find cuda/random.py as random...
|
||||
sys.path.remove(thedir)
|
||||
import argparse
|
||||
|
||||
fn_name = "torch.cuda.memory._snapshot()"
|
||||
pickled = f"pickled memory statistics from {fn_name}"
|
||||
parser = argparse.ArgumentParser(
|
||||
description=f"Visualize memory dumps produced by {fn_name}"
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest="action")
|
||||
|
||||
def _output(p):
|
||||
p.add_argument(
|
||||
"-o",
|
||||
"--output",
|
||||
default="output.svg",
|
||||
help="flamegraph svg (default: output.svg)",
|
||||
)
|
||||
|
||||
description = "Prints overall allocation statistics and a visualization of how the allocators segments are currently filled."
|
||||
stats_a = subparsers.add_parser("stats", description=description)
|
||||
stats_a.add_argument("input", help=pickled)
|
||||
|
||||
description = "Prints buffer of the most recent allocation events embedded in the snapshot in a Pythonic style."
|
||||
trace_a = subparsers.add_parser("trace", description=description)
|
||||
trace_a.add_argument("input", help=pickled)
|
||||
|
||||
description = "Generate a flamegraph that visualizes what memory is stored in each allocator segment (aka block)"
|
||||
segments_a = subparsers.add_parser("segments", description=description)
|
||||
segments_a.add_argument("input", help=pickled)
|
||||
_output(segments_a)
|
||||
|
||||
description = (
|
||||
"Generate a flamegraph the program locations contributing to CUDA memory usage."
|
||||
)
|
||||
memory_a = subparsers.add_parser("memory", description=description)
|
||||
memory_a.add_argument("input", help=pickled)
|
||||
_output(memory_a)
|
||||
|
||||
description = (
|
||||
"Generate a flamegraph that shows segments (aka blocks) that have been added "
|
||||
"or removed between two different memorys snapshots."
|
||||
)
|
||||
compare_a = subparsers.add_parser("compare", description=description)
|
||||
compare_a.add_argument("before", help=pickled)
|
||||
compare_a.add_argument("after", help=pickled)
|
||||
_output(compare_a)
|
||||
|
||||
plots = (
|
||||
(
|
||||
"trace_plot",
|
||||
"Generate a visualization over time of the memory usage recorded by the trace as an html file.",
|
||||
),
|
||||
(
|
||||
"segment_plot",
|
||||
"Visualize how allocations are packed into allocator segments at each point in a trace as an html file.",
|
||||
),
|
||||
)
|
||||
for cmd, description in plots:
|
||||
trace_plot_a = subparsers.add_parser(cmd, description=description)
|
||||
trace_plot_a.add_argument("input", help=pickled)
|
||||
help = "visualize trace from this device (default: chooses the only device with trace info or errors)"
|
||||
trace_plot_a.add_argument("-d", "--device", type=int, default=None, help=help)
|
||||
help = "path to save the visualization(default: output.html)"
|
||||
trace_plot_a.add_argument("-o", "--output", default="output.html", help=help)
|
||||
if cmd == "trace_plot":
|
||||
help = "visualize change to segments rather than individual allocations"
|
||||
trace_plot_a.add_argument(
|
||||
"-s", "--segments", action="store_true", help=help
|
||||
)
|
||||
|
||||
help = (
|
||||
"filter out allocation-free pairs to only visualize the allocations that are not freed yet;"
|
||||
"useful to reduce the number of events for large traces for debugging OOM"
|
||||
)
|
||||
trace_plot_a.add_argument(
|
||||
"-f", "--filter_freed", action="store_true", help=help
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
def _read(name):
|
||||
if name == "-":
|
||||
data = pickle.load(sys.stdin.buffer)
|
||||
else:
|
||||
with open(name, "rb") as f:
|
||||
data = pickle.load(f)
|
||||
if isinstance(data, list): # segments only...
|
||||
data = {"segments": data, "traces": []}
|
||||
return data
|
||||
|
||||
def _write(name, data):
|
||||
with open(name, "w") as f:
|
||||
f.write(data)
|
||||
|
||||
if args.action == "segments":
|
||||
data = _read(args.input)
|
||||
_write(args.output, segments(data))
|
||||
elif args.action == "memory":
|
||||
data = _read(args.input)
|
||||
_write(args.output, memory(data))
|
||||
elif args.action == "stats":
|
||||
data = _read(args.input)
|
||||
print(segsum(data))
|
||||
elif args.action == "trace":
|
||||
data = _read(args.input)
|
||||
print(trace(data))
|
||||
elif args.action == "compare":
|
||||
before = _read(args.before)
|
||||
after = _read(args.after)
|
||||
_write(args.output, compare(before, after))
|
||||
elif args.action == "trace_plot":
|
||||
data = _read(args.input)
|
||||
_write(
|
||||
args.output,
|
||||
trace_plot(
|
||||
data,
|
||||
device=args.device,
|
||||
plot_segments=args.segments,
|
||||
filter_freed=args.filter_freed,
|
||||
),
|
||||
)
|
||||
elif args.action == "segment_plot":
|
||||
data = _read(args.input)
|
||||
_write(args.output, segment_plot(data, device=args.device))
|
||||
@@ -0,0 +1,25 @@
|
||||
import torch
|
||||
|
||||
|
||||
def pin_memory(data_ptr: int, size: int) -> None:
|
||||
cudart = torch.cuda.cudart()
|
||||
succ = int(
|
||||
cudart.cudaHostRegister(
|
||||
data_ptr,
|
||||
size,
|
||||
1, # lines up with 'cudaHostRegisterPortable'
|
||||
)
|
||||
)
|
||||
|
||||
if succ != 0:
|
||||
raise RuntimeError(
|
||||
f"Registering memory failed with cudaError: {succ}."
|
||||
" It's possible that this is an asynchronous error raised from a previous cuda operation."
|
||||
" Consider launching with CUDA_LAUNCH_BLOCKING=1 to debug."
|
||||
)
|
||||
|
||||
|
||||
def unpin_memory(data_ptr: int) -> None:
|
||||
succ = int(torch.cuda.cudart().cudaHostUnregister(data_ptr))
|
||||
if succ != 0:
|
||||
raise AssertionError(f"Unpinning shared memory failed with error-code: {succ}")
|
||||
@@ -0,0 +1,670 @@
|
||||
# mypy: allow-untyped-defs
|
||||
r"""
|
||||
This module introduces CUDA Sanitizer, a tool for detecting synchronization errors between kernels ran on different streams.
|
||||
|
||||
It stores information on accesses to tensors to determine if they are synchronized
|
||||
or not. When enabled in a python program and a possible data race is detected, a
|
||||
detailed warning will be printed and the program will exit.
|
||||
|
||||
It can be enabled either by importing this module and calling
|
||||
:func:`enable_cuda_sanitizer()` or by exporting the ``TORCH_CUDA_SANITIZER``
|
||||
environment variable.
|
||||
"""
|
||||
|
||||
import enum
|
||||
import functools
|
||||
import inspect
|
||||
import io
|
||||
import logging
|
||||
import re
|
||||
import sys
|
||||
import textwrap
|
||||
import traceback
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, TypeVar
|
||||
|
||||
import torch
|
||||
import torch.cuda._gpu_trace as gpu_trace
|
||||
from torch.utils import _pytree as pytree
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
DEFAULT_STREAM_ID = 0
|
||||
|
||||
TK = TypeVar("TK")
|
||||
TVa = TypeVar("TVa")
|
||||
TVb = TypeVar("TVb")
|
||||
|
||||
DataPtr = int
|
||||
StreamId = int
|
||||
EventId = int
|
||||
SeqNum = int
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Note that this is only factories that take Tensor as input as they are
|
||||
# the ones we care about.
|
||||
FACTORY_FUNCTION_REGEX = re.compile("(new_.*|.*_like)")
|
||||
|
||||
|
||||
class AccessType(enum.Enum):
|
||||
READ = enum.auto()
|
||||
WRITE = enum.auto()
|
||||
|
||||
def __str__(self):
|
||||
return "reading from" if self is AccessType.READ else "writing to"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Access:
|
||||
r"""Stores information about a single access to a tensor by a kernel.
|
||||
|
||||
Args:
|
||||
type: either AccessType.READ or AccessType.Write.
|
||||
seq_num: the sequential number of the kernel performing the access.
|
||||
stream: the stream id of the stream executing the kernel.
|
||||
operator: the schema of the launched kernel, which lists the
|
||||
arguments and return type.
|
||||
aliases: the arguments in the schema this access corresponds to.
|
||||
is_output: Whether the tensor was an output of the kernel.
|
||||
stack_trace: the stack summary object captured during access.
|
||||
"""
|
||||
|
||||
type: AccessType
|
||||
seq_num: SeqNum
|
||||
stream: StreamId
|
||||
operator: str
|
||||
aliases: list[str]
|
||||
is_output: bool
|
||||
stack_trace: traceback.StackSummary
|
||||
|
||||
|
||||
class SynchronizationError(Exception):
|
||||
"""Base class for errors detected by CUDA Sanitizer."""
|
||||
|
||||
|
||||
class UnsynchronizedAccessError(SynchronizationError):
|
||||
"""Stores information about two unsynchronized accesses to one data pointer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_ptr: DataPtr,
|
||||
allocation_stack_trace: traceback.StackSummary | None,
|
||||
current_access: Access,
|
||||
previous_access: Access,
|
||||
):
|
||||
self.data_ptr = data_ptr
|
||||
self.allocation_stack_trace = allocation_stack_trace
|
||||
self.current_access = current_access
|
||||
self.previous_access = previous_access
|
||||
|
||||
def __str__(self):
|
||||
def format_access(access: Access):
|
||||
message.write(f"{access.operator}\n{access.type}")
|
||||
if access.aliases:
|
||||
message.write(" argument(s) " + ", ".join(access.aliases))
|
||||
if access.is_output:
|
||||
message.write(", and to")
|
||||
if access.is_output:
|
||||
message.write(" the output")
|
||||
message.write(
|
||||
f"\nWith stack trace:\n{''.join(access.stack_trace.format())}\n"
|
||||
)
|
||||
|
||||
with io.StringIO() as message:
|
||||
message.write(
|
||||
textwrap.dedent(
|
||||
f"""\
|
||||
============================
|
||||
CSAN detected a possible data race on tensor with data pointer {self.data_ptr}
|
||||
Access by stream {self.current_access.stream} during kernel:
|
||||
"""
|
||||
)
|
||||
)
|
||||
format_access(self.current_access)
|
||||
|
||||
message.write(
|
||||
f"Previous access by stream {self.previous_access.stream} during kernel:\n"
|
||||
)
|
||||
format_access(self.previous_access)
|
||||
|
||||
if self.allocation_stack_trace:
|
||||
message.write(
|
||||
"Tensor was allocated with stack trace:\n"
|
||||
f"{''.join(self.allocation_stack_trace.format())}"
|
||||
)
|
||||
else:
|
||||
message.write("Trace for tensor allocation not found.")
|
||||
return message.getvalue()
|
||||
|
||||
|
||||
class CUDASanitizerErrors(Exception):
|
||||
"""Wrapper class for errors reported by CUDA Sanitizer."""
|
||||
|
||||
def __init__(self, errors: list[SynchronizationError]):
|
||||
self.errors = errors
|
||||
|
||||
def __str__(self):
|
||||
return f"detected {len(self.errors)} errors"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TensorInfo:
|
||||
r"""Stores information about a single tensor and recent accesses to it.
|
||||
|
||||
Args:
|
||||
allocation_stack_trace: the stack summary object captured during tensor
|
||||
allocation. Can be ``None`` if the allocation wasn't caught by CSAN.
|
||||
reads: list of read accesses to the tensor that were performed since
|
||||
the last write.
|
||||
write: the last write access to the tensor.
|
||||
"""
|
||||
|
||||
allocation_stack_trace: traceback.StackSummary | None
|
||||
reads: list[Access] = field(default_factory=list)
|
||||
write: Access | None = None
|
||||
|
||||
|
||||
class _TensorsAccessed:
|
||||
def __init__(self) -> None:
|
||||
self.accesses: dict[DataPtr, TensorInfo] = {}
|
||||
|
||||
def ensure_tensor_exists(self, data_ptr: DataPtr) -> None:
|
||||
if data_ptr not in self.accesses:
|
||||
logger.info(
|
||||
"Found tensor with pointer: %s, but no matching tensor "
|
||||
"allocation in the trace. Backfilling the trace now. "
|
||||
"Perhaps the sanitizer was enabled after some torch operations?",
|
||||
data_ptr,
|
||||
)
|
||||
self.create_tensor(data_ptr, None)
|
||||
|
||||
def ensure_tensor_does_not_exist(self, data_ptr: DataPtr) -> None:
|
||||
if data_ptr in self.accesses:
|
||||
logger.info(
|
||||
"Found duplicate tensor allocation in the trace for tensor with "
|
||||
"pointer: %s. Assuming the trace for tensor deallocation "
|
||||
"wasn't caught and backfilling it now. "
|
||||
"Perhaps the sanitizer was enabled after some torch operations?",
|
||||
data_ptr,
|
||||
)
|
||||
self.delete_tensor(data_ptr)
|
||||
|
||||
def create_tensor(
|
||||
self, data_ptr: DataPtr, stack_trace: traceback.StackSummary | None
|
||||
) -> None:
|
||||
self.accesses[data_ptr] = TensorInfo(stack_trace)
|
||||
|
||||
def delete_tensor(self, data_ptr: DataPtr) -> None:
|
||||
del self.accesses[data_ptr]
|
||||
|
||||
def were_there_reads_since_last_write(self, data_ptr: DataPtr) -> bool:
|
||||
return bool(self.accesses[data_ptr].reads)
|
||||
|
||||
def get_allocation_stack_trace(
|
||||
self, data_ptr: DataPtr
|
||||
) -> traceback.StackSummary | None:
|
||||
return self.accesses[data_ptr].allocation_stack_trace
|
||||
|
||||
def get_write(self, data_ptr: DataPtr) -> Access | None:
|
||||
return self.accesses[data_ptr].write
|
||||
|
||||
def get_reads(self, data_ptr: DataPtr) -> list[Access]:
|
||||
return self.accesses[data_ptr].reads
|
||||
|
||||
def add_read(self, data_ptr: DataPtr, access: Access) -> None:
|
||||
self.accesses[data_ptr].reads.append(access)
|
||||
|
||||
def set_write(self, data_ptr: DataPtr, access: Access) -> None:
|
||||
self.accesses[data_ptr].write = access
|
||||
self.accesses[data_ptr].reads = []
|
||||
|
||||
|
||||
class StreamSynchronizations:
|
||||
def __init__(self) -> None:
|
||||
self.current_sync_states: dict[StreamId, dict[StreamId, SeqNum]] = {}
|
||||
self.recorded_sync_states: dict[EventId, dict[StreamId, SeqNum]] = {}
|
||||
self.host_sync_state: dict[StreamId, SeqNum] = {}
|
||||
self.create_stream(DEFAULT_STREAM_ID)
|
||||
|
||||
def _ensure_stream_exists(self, stream: StreamId) -> None:
|
||||
if stream not in self.current_sync_states:
|
||||
logger.info(
|
||||
"Found Stream with id: %s, but no matching stream "
|
||||
"creation in the trace. Backfilling the trace now. "
|
||||
"Perhaps the sanitizer was enabled after some torch operations?",
|
||||
stream,
|
||||
)
|
||||
self.create_stream(stream)
|
||||
|
||||
def _ensure_event_exists(self, event: EventId) -> None:
|
||||
if event not in self.recorded_sync_states:
|
||||
logger.info(
|
||||
"Found Event with id: %s, but no matching event "
|
||||
"creation in the trace. Backfilling the trace now. "
|
||||
"Perhaps the sanitizer was enabled after some torch operations?",
|
||||
event,
|
||||
)
|
||||
self.create_event(event)
|
||||
|
||||
def _ensure_event_does_not_exist(self, event: EventId) -> None:
|
||||
if event in self.recorded_sync_states:
|
||||
logger.info(
|
||||
"Found duplicate event creation in the trace for event with "
|
||||
"id: %s. Assuming the trace for event deletion wasn't caught "
|
||||
"and backfilling it now. "
|
||||
"Perhaps the sanitizer was enabled after some torch operations?",
|
||||
event,
|
||||
)
|
||||
self.delete_event(event)
|
||||
|
||||
def create_stream(self, stream: StreamId) -> None:
|
||||
if stream in self.current_sync_states:
|
||||
logger.info(
|
||||
"Found duplicate Stream creation in the trace for Stream with "
|
||||
"id: %s. PyTorch Streams are only created once, so this "
|
||||
"trace entry is ignored.",
|
||||
stream,
|
||||
)
|
||||
else:
|
||||
self.host_sync_state[stream] = 0
|
||||
self.current_sync_states[stream] = self.host_sync_state.copy()
|
||||
|
||||
def create_event(self, event: EventId) -> None:
|
||||
self._ensure_event_does_not_exist(event)
|
||||
self.recorded_sync_states[event] = {}
|
||||
|
||||
def delete_event(self, event: EventId) -> None:
|
||||
self._ensure_event_exists(event)
|
||||
del self.recorded_sync_states[event]
|
||||
|
||||
def update_seq_num(self, stream: StreamId, seq_num: SeqNum) -> None:
|
||||
self._ensure_stream_exists(stream)
|
||||
self.current_sync_states[stream][stream] = seq_num
|
||||
|
||||
def record_state(self, event: EventId, stream: StreamId) -> None:
|
||||
self._ensure_event_exists(event)
|
||||
self._ensure_stream_exists(stream)
|
||||
self.recorded_sync_states[event] = self.current_sync_states[stream].copy()
|
||||
|
||||
def _state_wait_for_other(
|
||||
self, state: dict[StreamId, SeqNum], other: dict[StreamId, SeqNum]
|
||||
) -> None:
|
||||
for stream, seq_num in other.items():
|
||||
state[stream] = max(state.get(stream, -1), seq_num)
|
||||
|
||||
def stream_wait_for_event(self, stream: StreamId, event: EventId) -> None:
|
||||
self._ensure_stream_exists(stream)
|
||||
self._ensure_event_exists(event)
|
||||
self._state_wait_for_other(
|
||||
self.current_sync_states[stream], self.recorded_sync_states[event]
|
||||
)
|
||||
|
||||
def all_streams_wait_for_event(self, event: EventId) -> None:
|
||||
self._ensure_event_exists(event)
|
||||
for stream in self.current_sync_states:
|
||||
self.stream_wait_for_event(stream, event)
|
||||
|
||||
self._state_wait_for_other(
|
||||
self.host_sync_state, self.recorded_sync_states[event]
|
||||
)
|
||||
|
||||
def all_streams_wait_for_stream(self, stream: StreamId) -> None:
|
||||
self._ensure_stream_exists(stream)
|
||||
for state in self.current_sync_states.values():
|
||||
self._state_wait_for_other(state, self.current_sync_states[stream])
|
||||
|
||||
self._state_wait_for_other(
|
||||
self.host_sync_state, self.current_sync_states[stream]
|
||||
)
|
||||
|
||||
def sync_all_streams(self) -> None:
|
||||
for stream, state in self.current_sync_states.items():
|
||||
self.host_sync_state[stream] = state[stream]
|
||||
|
||||
for state in self.current_sync_states.values():
|
||||
self._state_wait_for_other(state, self.host_sync_state)
|
||||
|
||||
def is_ordered_after(
|
||||
self, current_stream: StreamId, seq_num: SeqNum, other_stream: StreamId
|
||||
) -> bool:
|
||||
self._ensure_stream_exists(current_stream)
|
||||
self._ensure_stream_exists(other_stream)
|
||||
return seq_num <= self.current_sync_states[current_stream].get(other_stream, -1)
|
||||
|
||||
|
||||
class EventHandler:
|
||||
"""Analyzes CSAN trace for synchronization errors.
|
||||
|
||||
Stores information on each stream's synchronizations with other streams as well
|
||||
as tensor accesses to determine whether a given kernel launch might cause a
|
||||
data race.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.tensors_accessed = _TensorsAccessed()
|
||||
self.syncs = StreamSynchronizations()
|
||||
self.seq_num: SeqNum = 0
|
||||
|
||||
def _handle_kernel_launch(
|
||||
self,
|
||||
stream: StreamId,
|
||||
read_only: set[DataPtr],
|
||||
read_write: set[DataPtr],
|
||||
outputs: set[DataPtr],
|
||||
operator: str,
|
||||
tensor_aliases: dict[int, list[str]],
|
||||
) -> list[SynchronizationError]:
|
||||
def check_conflict(
|
||||
data_ptr: DataPtr, current_access: Access, previous_access: Access | None
|
||||
) -> None:
|
||||
if previous_access is None:
|
||||
return
|
||||
if not self.syncs.is_ordered_after(
|
||||
current_access.stream, previous_access.seq_num, previous_access.stream
|
||||
):
|
||||
error_list.append(
|
||||
UnsynchronizedAccessError(
|
||||
data_ptr,
|
||||
self.tensors_accessed.get_allocation_stack_trace(data_ptr),
|
||||
current_access,
|
||||
previous_access,
|
||||
)
|
||||
)
|
||||
|
||||
error_list: list[SynchronizationError] = []
|
||||
self.seq_num += 1
|
||||
self.syncs.update_seq_num(stream, self.seq_num)
|
||||
stack_trace = traceback.StackSummary.extract(
|
||||
traceback.walk_stack(inspect.currentframe()), lookup_lines=False
|
||||
)
|
||||
# The stack trace generated in this way is in the inverse order, so it must be
|
||||
# reversed.
|
||||
stack_trace.reverse()
|
||||
|
||||
for data_ptr in read_only:
|
||||
self.tensors_accessed.ensure_tensor_exists(data_ptr)
|
||||
current_access = Access(
|
||||
AccessType.READ,
|
||||
self.seq_num,
|
||||
stream,
|
||||
operator,
|
||||
tensor_aliases[data_ptr],
|
||||
data_ptr in outputs,
|
||||
stack_trace,
|
||||
)
|
||||
check_conflict(
|
||||
data_ptr, current_access, self.tensors_accessed.get_write(data_ptr)
|
||||
)
|
||||
self.tensors_accessed.add_read(data_ptr, current_access)
|
||||
|
||||
for data_ptr in read_write:
|
||||
self.tensors_accessed.ensure_tensor_exists(data_ptr)
|
||||
current_access = Access(
|
||||
AccessType.WRITE,
|
||||
self.seq_num,
|
||||
stream,
|
||||
operator,
|
||||
tensor_aliases[data_ptr],
|
||||
data_ptr in outputs,
|
||||
stack_trace,
|
||||
)
|
||||
if self.tensors_accessed.were_there_reads_since_last_write(data_ptr):
|
||||
for previous_access in self.tensors_accessed.get_reads(data_ptr):
|
||||
check_conflict(data_ptr, current_access, previous_access)
|
||||
else:
|
||||
check_conflict(
|
||||
data_ptr, current_access, self.tensors_accessed.get_write(data_ptr)
|
||||
)
|
||||
self.tensors_accessed.set_write(data_ptr, current_access)
|
||||
|
||||
return error_list
|
||||
|
||||
def _handle_event_creation(self, event: EventId) -> None:
|
||||
self.syncs.create_event(event)
|
||||
|
||||
def _handle_event_deletion(self, event: EventId) -> None:
|
||||
self.syncs.delete_event(event)
|
||||
|
||||
def _handle_event_record(self, event: EventId, stream: StreamId) -> None:
|
||||
self.syncs.record_state(event, stream)
|
||||
|
||||
def _handle_event_wait(self, event: EventId, stream: StreamId) -> None:
|
||||
self.syncs.stream_wait_for_event(stream, event)
|
||||
|
||||
def _handle_memory_allocation(self, data_ptr: DataPtr) -> None:
|
||||
self.tensors_accessed.ensure_tensor_does_not_exist(data_ptr)
|
||||
stack_trace = traceback.StackSummary.extract(
|
||||
traceback.walk_stack(inspect.currentframe()), lookup_lines=False
|
||||
)
|
||||
# The stack trace generated in this way is in the inverse order, so it must be
|
||||
# reversed.
|
||||
stack_trace.reverse()
|
||||
self.tensors_accessed.create_tensor(
|
||||
data_ptr,
|
||||
stack_trace,
|
||||
)
|
||||
|
||||
def _handle_memory_deallocation(self, data_ptr: DataPtr) -> None:
|
||||
self.tensors_accessed.ensure_tensor_exists(data_ptr)
|
||||
self.tensors_accessed.delete_tensor(data_ptr)
|
||||
|
||||
def _handle_stream_creation(self, stream: StreamId) -> None:
|
||||
self.syncs.create_stream(stream)
|
||||
|
||||
def _handle_device_synchronization(self) -> None:
|
||||
self.syncs.sync_all_streams()
|
||||
|
||||
def _handle_stream_synchronization(self, stream: StreamId) -> None:
|
||||
self.syncs.all_streams_wait_for_stream(stream)
|
||||
|
||||
def _handle_event_synchronization(self, event: EventId) -> None:
|
||||
self.syncs.all_streams_wait_for_event(event)
|
||||
|
||||
|
||||
def zip_by_key(a: dict[TK, TVa], b: dict[TK, TVb]) -> Iterator[tuple[TK, TVa, TVb]]:
|
||||
for arg, value in a.items():
|
||||
if arg in b:
|
||||
yield arg, value, b[arg]
|
||||
|
||||
|
||||
def zip_arguments(
|
||||
schema: torch.FunctionSchema, args: tuple[Any, ...], kwargs: dict[str, Any]
|
||||
) -> Iterator[tuple[torch.Argument, Any]]:
|
||||
schema_args = schema.arguments[: len(args)]
|
||||
schema_kwargs = {arg.name: arg for arg in schema.arguments[len(args) :]}
|
||||
|
||||
yield from zip(schema_args, args)
|
||||
|
||||
for _, argument, value in zip_by_key(schema_kwargs, kwargs):
|
||||
yield (argument, value)
|
||||
|
||||
|
||||
class ArgumentHandler:
|
||||
def __init__(self) -> None:
|
||||
self.dataptrs_read: set[DataPtr] = set()
|
||||
self.dataptrs_written: set[DataPtr] = set()
|
||||
self.tensor_aliases: dict[DataPtr, list[str]] = {}
|
||||
self.outputs: set[DataPtr] = set()
|
||||
|
||||
def _handle_argument(
|
||||
self,
|
||||
value: Any,
|
||||
is_write: bool,
|
||||
metadata_only: bool,
|
||||
name: str | None = None,
|
||||
is_output: bool = False,
|
||||
) -> None:
|
||||
if isinstance(value, torch.Tensor) and value.is_cuda:
|
||||
# data_ptr() is preferred, but distinguish Tensors with null data_ptr()
|
||||
# otherwise two empty Tensors could incorrectly match as a conflict
|
||||
data_ptr = value.data_ptr() if value.data_ptr() else id(value)
|
||||
if is_write:
|
||||
self.dataptrs_written.add(data_ptr)
|
||||
elif not metadata_only:
|
||||
self.dataptrs_read.add(data_ptr)
|
||||
|
||||
self.tensor_aliases.setdefault(data_ptr, [])
|
||||
if name is not None:
|
||||
self.tensor_aliases[data_ptr].append(name)
|
||||
if is_output:
|
||||
self.outputs.add(data_ptr)
|
||||
|
||||
def parse_inputs(
|
||||
self,
|
||||
schema: torch.FunctionSchema,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: dict[str, Any],
|
||||
*,
|
||||
is_factory: bool,
|
||||
) -> None:
|
||||
for argument, value in zip_arguments(schema, args, kwargs):
|
||||
is_write = argument.alias_info is not None and argument.alias_info.is_write
|
||||
# A change is metadata only if it is a view or a factory function that
|
||||
# reads only metadata
|
||||
metadata_only = is_factory or (
|
||||
argument.alias_info is not None and not argument.alias_info.is_write
|
||||
)
|
||||
pytree.tree_map_(
|
||||
functools.partial(
|
||||
self._handle_argument,
|
||||
is_write=is_write,
|
||||
name=argument.name,
|
||||
metadata_only=metadata_only,
|
||||
),
|
||||
value,
|
||||
)
|
||||
|
||||
def parse_outputs(
|
||||
self, schema: torch.FunctionSchema, outputs: Any, *, is_factory: bool
|
||||
) -> None:
|
||||
for res, value in zip(schema.returns, (outputs,)):
|
||||
metadata_only = is_factory or (
|
||||
res.alias_info is not None and not res.alias_info.is_write
|
||||
)
|
||||
pytree.tree_map_(
|
||||
functools.partial(
|
||||
self._handle_argument,
|
||||
is_write=not metadata_only,
|
||||
is_output=True,
|
||||
metadata_only=metadata_only,
|
||||
),
|
||||
value,
|
||||
)
|
||||
|
||||
|
||||
class CUDASanitizerDispatchMode(TorchDispatchMode):
|
||||
def __init__(self) -> None:
|
||||
self.event_handler = EventHandler()
|
||||
torch._C._activate_gpu_trace()
|
||||
gpu_trace.register_callback_for_event_creation(
|
||||
self.event_handler._handle_event_creation
|
||||
)
|
||||
gpu_trace.register_callback_for_event_deletion(
|
||||
self.event_handler._handle_event_deletion
|
||||
)
|
||||
gpu_trace.register_callback_for_event_record(
|
||||
self.event_handler._handle_event_record
|
||||
)
|
||||
gpu_trace.register_callback_for_event_wait(
|
||||
self.event_handler._handle_event_wait
|
||||
)
|
||||
gpu_trace.register_callback_for_memory_allocation(
|
||||
self.event_handler._handle_memory_allocation
|
||||
)
|
||||
gpu_trace.register_callback_for_memory_deallocation(
|
||||
self.event_handler._handle_memory_deallocation
|
||||
)
|
||||
gpu_trace.register_callback_for_stream_creation(
|
||||
self.event_handler._handle_stream_creation
|
||||
)
|
||||
gpu_trace.register_callback_for_device_synchronization(
|
||||
self.event_handler._handle_device_synchronization
|
||||
)
|
||||
gpu_trace.register_callback_for_stream_synchronization(
|
||||
self.event_handler._handle_stream_synchronization
|
||||
)
|
||||
gpu_trace.register_callback_for_event_synchronization(
|
||||
self.event_handler._handle_event_synchronization
|
||||
)
|
||||
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
# record_stream is not a kernel dispatch, skip it
|
||||
if func is aten.record_stream.default:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
is_factory = bool(FACTORY_FUNCTION_REGEX.match(func._schema.name))
|
||||
|
||||
argument_handler = ArgumentHandler()
|
||||
argument_handler.parse_inputs(func._schema, args, kwargs, is_factory=is_factory)
|
||||
|
||||
outputs = func(*args, **kwargs)
|
||||
|
||||
argument_handler.parse_outputs(func._schema, outputs, is_factory=is_factory)
|
||||
errors = self.event_handler._handle_kernel_launch(
|
||||
torch.cuda.current_stream().cuda_stream,
|
||||
argument_handler.dataptrs_read - argument_handler.dataptrs_written,
|
||||
argument_handler.dataptrs_written,
|
||||
argument_handler.outputs,
|
||||
func._schema,
|
||||
argument_handler.tensor_aliases,
|
||||
)
|
||||
if errors:
|
||||
for error in errors:
|
||||
print(error, file=sys.stderr)
|
||||
raise CUDASanitizerErrors(errors)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
class CUDASanitizer:
|
||||
"""Manages the lifetime of a CUDASanitizer dispatch mode object.
|
||||
|
||||
The CUDASanitizer class wraps the entering/exiting functions of the dispatch mode
|
||||
context manager in the enable function/destructor, respectively. This is to
|
||||
explicitly set the lifetime of the dispatch mode object to that of the application.
|
||||
This approach was deemed more elegant than using the atexit module.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.dispatch = CUDASanitizerDispatchMode()
|
||||
self.enabled = False
|
||||
|
||||
def enable(self):
|
||||
self.dispatch.__enter__()
|
||||
self.enabled = True
|
||||
|
||||
def disable(self):
|
||||
self.dispatch.__exit__(None, None, None)
|
||||
self.enabled = False
|
||||
|
||||
def __del__(self):
|
||||
# Since this object lifetime is linked to the `torch.cuda._sanitizer` python
|
||||
# module, it often gets deleted as part of the overall `torch` module cleanup
|
||||
# At that time, depending on CPython version, the torch.* module might be in
|
||||
# different states of being already cleaned up.
|
||||
# Similarly other imports might already have been cleaned up so `sys` might
|
||||
# be already gone as well.
|
||||
# Skip exiting the mode if it outlived the runtime.
|
||||
if (sys is not None) and (not sys.is_finalizing()) and self.enabled:
|
||||
self.disable()
|
||||
|
||||
|
||||
def enable_cuda_sanitizer():
|
||||
"""Enable CUDA Sanitizer.
|
||||
|
||||
The sanitizer will begin to analyze low-level CUDA calls invoked by torch functions
|
||||
for synchronization errors. All data races found will be printed to the standard
|
||||
error output along with stack traces of suspected causes. For best results, the
|
||||
sanitizer should be enabled at the very beginning of the program.
|
||||
"""
|
||||
cuda_sanitizer.enable()
|
||||
|
||||
|
||||
cuda_sanitizer = CUDASanitizer()
|
||||
@@ -0,0 +1,586 @@
|
||||
import ctypes
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
try:
|
||||
from cuda.bindings import ( # pyrefly: ignore[missing-import]
|
||||
runtime as _cuda_bindings_runtime,
|
||||
)
|
||||
|
||||
_HAS_CUDA_BINDINGS = True
|
||||
except ImportError:
|
||||
_cuda_bindings_runtime = None # type: ignore[assignment]
|
||||
_HAS_CUDA_BINDINGS = False
|
||||
|
||||
# The _get_device_index has been moved to torch.utils._get_device_index
|
||||
from torch._utils import _get_device_index as _torch_get_device_index
|
||||
|
||||
|
||||
def _get_hip_runtime_library() -> ctypes.CDLL:
|
||||
# If ROCm python packages are available, query the OS-independent absolute
|
||||
# path to the library provided by those packages, including any version suffix.
|
||||
# See https://github.com/ROCm/TheRock/blob/main/docs/packaging/python_packaging.md#dynamic-library-resolution
|
||||
try:
|
||||
# pyrefly: ignore [import-error, missing-import]
|
||||
import rocm_sdk
|
||||
|
||||
lib = ctypes.CDLL(str(rocm_sdk.find_libraries("amdhip64")[0]))
|
||||
except (ImportError, IndexError):
|
||||
if sys.platform == "win32":
|
||||
lib = ctypes.CDLL(f"amdhip64_{torch.version.hip[0]}.dll")
|
||||
else: # Unix-based systems
|
||||
lib = ctypes.CDLL("libamdhip64.so")
|
||||
|
||||
lib.cuGetErrorString = lib.hipGetErrorString # type: ignore[attr-defined]
|
||||
lib.cuModuleLoadData = lib.hipModuleLoadData # type: ignore[attr-defined]
|
||||
lib.cuModuleGetFunction = lib.hipModuleGetFunction # type: ignore[attr-defined]
|
||||
lib.cuLaunchKernel = lib.hipModuleLaunchKernel # type: ignore[attr-defined]
|
||||
lib.cuFuncSetAttribute = lib.hipFuncSetAttribute # type: ignore[attr-defined]
|
||||
return lib
|
||||
|
||||
|
||||
def _get_cuda_library() -> ctypes.CDLL:
|
||||
if sys.platform == "win32":
|
||||
return ctypes.CDLL("nvcuda.dll")
|
||||
else: # Unix-based systems
|
||||
return ctypes.CDLL("libcuda.so.1")
|
||||
|
||||
|
||||
# Load GPU driver runtime
|
||||
def _get_gpu_runtime_library() -> ctypes.CDLL:
|
||||
if torch.version.hip:
|
||||
return _get_hip_runtime_library()
|
||||
else:
|
||||
return _get_cuda_library()
|
||||
|
||||
|
||||
# Helper: check CUDA errors
|
||||
def _check_cuda(result: int) -> None:
|
||||
if result == 0:
|
||||
return
|
||||
err_str = ctypes.c_char_p()
|
||||
libcuda = _get_gpu_runtime_library() # Get reference to CUDA library
|
||||
libcuda.cuGetErrorString(result, ctypes.byref(err_str))
|
||||
error_message = (
|
||||
err_str.value.decode() if err_str.value is not None else "Unknown CUDA error"
|
||||
)
|
||||
raise RuntimeError(f"CUDA error: {error_message}")
|
||||
|
||||
|
||||
def _check_cuda_bindings(result: Any) -> Any:
|
||||
"""Check a cuda.bindings (cuda-python) call result for errors.
|
||||
|
||||
All cuda.bindings runtime calls return ``(error, *outputs)``. This
|
||||
helper unpacks the tuple, raises on non-success, and returns the
|
||||
outputs (``None`` for zero outputs, scalar for one, tuple otherwise).
|
||||
"""
|
||||
if not _HAS_CUDA_BINDINGS:
|
||||
raise RuntimeError("cuda.bindings is not available")
|
||||
err, *out = result
|
||||
if (
|
||||
err
|
||||
!= _cuda_bindings_runtime.cudaError_t.cudaSuccess # pyrefly: ignore[missing-attribute]
|
||||
):
|
||||
_, err_str = (
|
||||
_cuda_bindings_runtime.cudaGetErrorString( # pyrefly: ignore[missing-attribute]
|
||||
err
|
||||
)
|
||||
)
|
||||
if isinstance(err_str, bytes):
|
||||
err_str = err_str.decode()
|
||||
raise RuntimeError(f"CUDA error: {err} ({err_str})")
|
||||
if len(out) == 0:
|
||||
return None
|
||||
if len(out) == 1:
|
||||
return out[0]
|
||||
return out
|
||||
|
||||
|
||||
def _get_hiprtc_library() -> ctypes.CDLL:
|
||||
try:
|
||||
# pyrefly: ignore [import-error, missing-import]
|
||||
import rocm_sdk
|
||||
|
||||
lib = ctypes.CDLL(str(rocm_sdk.find_libraries("hiprtc")[0]))
|
||||
except (ImportError, IndexError):
|
||||
if sys.platform == "win32":
|
||||
version_str = "".join(
|
||||
["0", torch.version.hip[0], "0", torch.version.hip[2]]
|
||||
)
|
||||
lib = ctypes.CDLL(f"hiprtc{version_str}.dll")
|
||||
else:
|
||||
lib = ctypes.CDLL("libhiprtc.so")
|
||||
|
||||
# Provide aliases for HIP RTC functions to match NVRTC API
|
||||
lib.nvrtcGetErrorString = lib.hiprtcGetErrorString # type: ignore[attr-defined]
|
||||
lib.nvrtcCreateProgram = lib.hiprtcCreateProgram # type: ignore[attr-defined]
|
||||
lib.nvrtcDestroyProgram = lib.hiprtcDestroyProgram # type: ignore[attr-defined]
|
||||
lib.nvrtcCompileProgram = lib.hiprtcCompileProgram # type: ignore[attr-defined]
|
||||
lib.nvrtcGetCUBINSize = lib.hiprtcGetCodeSize # type: ignore[attr-defined]
|
||||
lib.nvrtcGetCUBIN = lib.hiprtcGetCode # type: ignore[attr-defined]
|
||||
lib.nvrtcGetProgramLogSize = lib.hiprtcGetProgramLogSize # type: ignore[attr-defined]
|
||||
lib.nvrtcGetProgramLog = lib.hiprtcGetProgramLog # type: ignore[attr-defined]
|
||||
lib.nvrtcAddNameExpression = lib.hiprtcAddNameExpression # type: ignore[attr-defined]
|
||||
lib.nvrtcGetLoweredName = lib.hiprtcGetLoweredName # type: ignore[attr-defined]
|
||||
return lib
|
||||
|
||||
|
||||
def _get_nvrtc_library() -> ctypes.CDLL:
|
||||
major_version = int(torch.version.cuda.split(".")[0]) # type: ignore[union-attr]
|
||||
if sys.platform == "win32":
|
||||
nvrtc_libs = [
|
||||
f"nvrtc64_{major_version}0_0.dll",
|
||||
]
|
||||
else:
|
||||
nvrtc_libs = [
|
||||
f"libnvrtc.so.{major_version}",
|
||||
"libnvrtc.so", # Fallback to unversioned
|
||||
]
|
||||
for lib_name in nvrtc_libs:
|
||||
try:
|
||||
return ctypes.CDLL(lib_name)
|
||||
except OSError:
|
||||
continue
|
||||
raise OSError("Could not find any NVRTC library")
|
||||
|
||||
|
||||
def _get_gpu_rtc_library() -> ctypes.CDLL:
|
||||
# Since PyTorch already loads the GPU RTC library, we can use the system library
|
||||
# which should be compatible with PyTorch's version
|
||||
if torch.version.hip:
|
||||
return _get_hiprtc_library()
|
||||
else:
|
||||
return _get_nvrtc_library()
|
||||
|
||||
|
||||
def _get_gpu_rtc_compatible_flags() -> list[str]:
|
||||
"""
|
||||
Get HIPCC/NVCC flags that are compatible with NVRTC compilation.
|
||||
|
||||
Returns:
|
||||
List of HIPCC/NVCC flags that can be safely used with NVRTC.
|
||||
"""
|
||||
from torch.utils.cpp_extension import COMMON_HIPCC_FLAGS, COMMON_NVCC_FLAGS
|
||||
|
||||
nvrtc_unsupported_flags = {
|
||||
"--expt-relaxed-constexpr",
|
||||
}
|
||||
|
||||
# Filter out unsupported flags
|
||||
compatible_flags = [
|
||||
flag for flag in COMMON_NVCC_FLAGS if flag not in nvrtc_unsupported_flags
|
||||
]
|
||||
|
||||
if torch.version.hip:
|
||||
compatible_flags.extend(COMMON_HIPCC_FLAGS)
|
||||
|
||||
return compatible_flags
|
||||
|
||||
|
||||
def _nvrtc_compile(
|
||||
kernel_source: str,
|
||||
kernel_name: str,
|
||||
compute_capability: str | None = None,
|
||||
cuda_include_dirs: list | None = None,
|
||||
nvcc_options: list | None = None,
|
||||
auto_pch: bool = False,
|
||||
) -> tuple[bytes, str]:
|
||||
"""
|
||||
Compiles a CUDA kernel using NVRTC and returns the PTX code.
|
||||
|
||||
Args:
|
||||
kernel_source (str): The CUDA kernel source code as a string
|
||||
kernel_name (str): The name of the kernel function to compile
|
||||
compute_capability (str, None): The compute capability to target (e.g., "86").
|
||||
If None, will detect from current device.
|
||||
cuda_include_dirs (list, None): List of directories containing CUDA headers
|
||||
nvcc_options (list, None): Additional options to pass to NVRTC
|
||||
auto_pch (bool): Enable automatic precompiled headers (CUDA 12.8+)
|
||||
|
||||
Returns:
|
||||
Tuple[bytes, str]: The compiled PTX code and mangled kernel name
|
||||
"""
|
||||
# Ensure CUDA is initialized
|
||||
import torch.cuda
|
||||
|
||||
# Load NVRTC library
|
||||
libnvrtc = _get_gpu_rtc_library()
|
||||
|
||||
# NVRTC constants
|
||||
NVRTC_SUCCESS = 0
|
||||
|
||||
# Helper: check NVRTC errors
|
||||
def check_nvrtc(result: int) -> None:
|
||||
if result != NVRTC_SUCCESS:
|
||||
err_str = ctypes.c_char_p()
|
||||
libnvrtc.nvrtcGetErrorString(result, ctypes.byref(err_str))
|
||||
error_message = (
|
||||
err_str.value.decode()
|
||||
if err_str.value is not None
|
||||
else "Unknown CUDA error"
|
||||
)
|
||||
raise RuntimeError(f"CUDA error: {error_message}")
|
||||
|
||||
# Convert source to bytes
|
||||
source_bytes = kernel_source.encode("utf-8")
|
||||
|
||||
# Get compute capability if not provided
|
||||
if compute_capability is None:
|
||||
props = torch.cuda.get_device_properties(torch.cuda.current_device())
|
||||
if torch.version.hip:
|
||||
compute_capability = f"{props.gcnArchName}"
|
||||
else:
|
||||
compute_capability = f"{props.major}{props.minor}"
|
||||
|
||||
# Prepare compilation options
|
||||
options = []
|
||||
if torch.version.hip:
|
||||
options.append(f"--offload-arch={compute_capability}".encode())
|
||||
else:
|
||||
options.append(f"--gpu-architecture=sm_{compute_capability}".encode())
|
||||
|
||||
# Auto-detect and add CUDA include paths
|
||||
from torch.utils.cpp_extension import include_paths
|
||||
|
||||
cuda_include_paths = include_paths("cuda")
|
||||
for cuda_path in cuda_include_paths:
|
||||
options.append(f"-I{cuda_path}".encode())
|
||||
|
||||
# Add custom include directories
|
||||
if cuda_include_dirs:
|
||||
for directory in cuda_include_dirs:
|
||||
options.append(f"-I{directory}".encode())
|
||||
|
||||
# Enable automatic precompiled headers (CUDA 12.8+)
|
||||
if auto_pch:
|
||||
if str(torch.version.cuda) < "12.8":
|
||||
raise AssertionError(f"PCH requires CUDA 12.8+, got {torch.version.cuda}")
|
||||
if nvcc_options is None:
|
||||
nvcc_options = []
|
||||
nvcc_options.append("--pch")
|
||||
|
||||
# Add custom NVCC options
|
||||
if nvcc_options:
|
||||
for option in nvcc_options:
|
||||
options.append(option.encode("utf-8"))
|
||||
|
||||
nvrtc_compatible_flags = _get_gpu_rtc_compatible_flags()
|
||||
options.extend([flag.encode("utf-8") for flag in nvrtc_compatible_flags])
|
||||
|
||||
# Convert options to C array
|
||||
num_options = len(options)
|
||||
options_array = (ctypes.c_char_p * num_options)(*options)
|
||||
|
||||
# Create program
|
||||
prog = ctypes.c_void_p()
|
||||
check_nvrtc(
|
||||
libnvrtc.nvrtcCreateProgram(
|
||||
ctypes.byref(prog),
|
||||
source_bytes,
|
||||
f"{kernel_name}.cu".encode(),
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
)
|
||||
|
||||
# Add kernel name, which can be a template expression
|
||||
c_kernel_name = kernel_name.encode("utf-8")
|
||||
check_nvrtc(libnvrtc.nvrtcAddNameExpression(prog, c_kernel_name))
|
||||
|
||||
# Compile program
|
||||
res = libnvrtc.nvrtcCompileProgram(prog, num_options, options_array)
|
||||
|
||||
# Handle compilation errors
|
||||
if res != NVRTC_SUCCESS:
|
||||
# Get log
|
||||
log_size = ctypes.c_size_t()
|
||||
libnvrtc.nvrtcGetProgramLogSize(prog, ctypes.byref(log_size))
|
||||
log = ctypes.create_string_buffer(log_size.value)
|
||||
libnvrtc.nvrtcGetProgramLog(prog, log)
|
||||
raise RuntimeError(f"Kernel compilation failed:\n{log.value.decode()}")
|
||||
|
||||
# Get binary
|
||||
binary_size = ctypes.c_size_t()
|
||||
check_nvrtc(libnvrtc.nvrtcGetCUBINSize(prog, ctypes.byref(binary_size)))
|
||||
binary = ctypes.create_string_buffer(binary_size.value)
|
||||
check_nvrtc(libnvrtc.nvrtcGetCUBIN(prog, binary))
|
||||
|
||||
# Get mangled name
|
||||
c_mangled_name = ctypes.c_char_p()
|
||||
check_nvrtc(
|
||||
libnvrtc.nvrtcGetLoweredName(prog, c_kernel_name, ctypes.byref(c_mangled_name))
|
||||
)
|
||||
if c_mangled_name.value is not None:
|
||||
mangled_name = c_mangled_name.value.decode() # make a copy
|
||||
else:
|
||||
mangled_name = ""
|
||||
|
||||
libnvrtc.nvrtcDestroyProgram(ctypes.byref(prog))
|
||||
|
||||
# For some reason, ".value" causes the string to be truncated,
|
||||
# likely due to the presence of '\0' in the string. So we use .raw instead.
|
||||
return binary.raw, mangled_name
|
||||
|
||||
|
||||
class _CudaModule:
|
||||
def __init__(self, module: ctypes.c_void_p) -> None:
|
||||
self._module = module
|
||||
self._kernels: dict[str, _CudaKernel] = {}
|
||||
|
||||
def __getattr__(self, name: str) -> "_CudaKernel":
|
||||
if name in self._kernels:
|
||||
return self._kernels[name]
|
||||
|
||||
# Import the CUDA library inside the method
|
||||
# pyrefly: ignore [missing-module-attribute]
|
||||
from torch.cuda._utils import _get_gpu_runtime_library
|
||||
|
||||
libcuda = _get_gpu_runtime_library()
|
||||
|
||||
func = ctypes.c_void_p()
|
||||
try:
|
||||
_check_cuda(
|
||||
libcuda.cuModuleGetFunction(
|
||||
ctypes.byref(func), self._module, name.encode("utf-8")
|
||||
)
|
||||
)
|
||||
kernel = _CudaKernel(func, self._module)
|
||||
self._kernels[name] = kernel
|
||||
return kernel
|
||||
|
||||
except RuntimeError as err:
|
||||
raise AttributeError(f"No kernel named '{name}' in this module") from err
|
||||
|
||||
|
||||
class _CudaKernel:
|
||||
"""
|
||||
Represents a compiled CUDA kernel that can be called with PyTorch tensors.
|
||||
"""
|
||||
|
||||
def __init__(self, func: ctypes.c_void_p, module: ctypes.c_void_p) -> None:
|
||||
self.func = func
|
||||
self.module = module
|
||||
self._max_shared_mem_bytes = 0
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
grid: tuple[int, int, int] = (1, 1, 1),
|
||||
block: tuple[int, int, int] = (1, 1, 1),
|
||||
args: list | None = None,
|
||||
shared_mem: int = 0,
|
||||
stream: Any | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Call the compiled CUDA kernel
|
||||
|
||||
Args:
|
||||
grid (tuple): Grid dimensions (grid_x, grid_y, grid_z)
|
||||
block (tuple): Block dimensions (block_x, block_y, block_z)
|
||||
args (list): List of arguments to pass to the kernel.
|
||||
PyTorch tensor arguments will be automatically converted to pointers.
|
||||
shared_mem (int): Shared memory size in bytes
|
||||
stream (torch.cuda.Stream): CUDA stream to use. If None, uses current stream.
|
||||
"""
|
||||
import torch
|
||||
|
||||
libcuda = torch.cuda._utils._get_gpu_runtime_library()
|
||||
|
||||
if not args:
|
||||
args = []
|
||||
|
||||
# Process arguments and convert tensors to pointers
|
||||
processed_args: list[ctypes.c_void_p] = []
|
||||
c_args = []
|
||||
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor):
|
||||
if not arg.is_cuda and not (arg.is_cpu and arg.is_pinned()):
|
||||
raise ValueError(
|
||||
"All tensor arguments must be CUDA tensors or pinned CPU tensors"
|
||||
)
|
||||
# Get pointer to tensor data
|
||||
ptr = ctypes.c_void_p(arg.data_ptr())
|
||||
processed_args.append(ptr)
|
||||
c_args.append(ctypes.byref(ptr))
|
||||
elif isinstance(arg, int):
|
||||
# Convert integers to C int
|
||||
c_int = ctypes.c_int(arg)
|
||||
# Store the C int for reference keeping, not in processed_args
|
||||
c_args.append(ctypes.byref(c_int))
|
||||
elif isinstance(arg, float):
|
||||
# Python floats are doubles - use double by default
|
||||
c_double = ctypes.c_double(arg)
|
||||
# Store the C double for reference keeping, not in processed_args
|
||||
c_args.append(ctypes.byref(c_double))
|
||||
else:
|
||||
raise TypeError(f"Unsupported argument type: {type(arg)}")
|
||||
|
||||
# Convert to array of void pointers
|
||||
c_args_array = (ctypes.c_void_p * len(c_args))()
|
||||
for i, arg in enumerate(c_args):
|
||||
c_args_array[i] = ctypes.cast(arg, ctypes.c_void_p)
|
||||
|
||||
# Get the stream
|
||||
if stream is None:
|
||||
# Defer import to avoid circular imports
|
||||
import torch.cuda
|
||||
|
||||
stream = torch.cuda.current_stream()
|
||||
|
||||
# Check if kernel requires large shared memory but hasn't been configured
|
||||
if shared_mem >= 48 * 1024 and (
|
||||
self._max_shared_mem_bytes == 0 or shared_mem > self._max_shared_mem_bytes
|
||||
):
|
||||
configured_msg = (
|
||||
"not configured"
|
||||
if self._max_shared_mem_bytes == 0
|
||||
else f"only {self._max_shared_mem_bytes} bytes configured"
|
||||
)
|
||||
raise RuntimeError(
|
||||
f"Kernel requires {shared_mem} bytes of shared memory (>= 48KB), "
|
||||
f"but {configured_msg}. "
|
||||
"Call kernel.set_shared_memory_config(shared_mem) after compilation "
|
||||
"and before launching the kernel."
|
||||
)
|
||||
|
||||
_check_cuda(
|
||||
libcuda.cuLaunchKernel(
|
||||
self.func,
|
||||
grid[0],
|
||||
grid[1],
|
||||
grid[2],
|
||||
block[0],
|
||||
block[1],
|
||||
block[2],
|
||||
shared_mem,
|
||||
stream._as_parameter_,
|
||||
c_args_array,
|
||||
None,
|
||||
)
|
||||
)
|
||||
|
||||
def set_shared_memory_config(self, shared_mem_bytes: int) -> None:
|
||||
if shared_mem_bytes < 48 * 1024:
|
||||
# No configuration needed for <= 48KB, just update the value
|
||||
self._max_shared_mem_bytes = shared_mem_bytes
|
||||
return
|
||||
|
||||
libcuda = _get_gpu_runtime_library()
|
||||
|
||||
# Get device properties to validate against limits
|
||||
device_props = torch.cuda.get_device_properties()
|
||||
# HIP doesn't have shared_memory_per_block_optin in device properties, so we hard-code it here
|
||||
if torch.version.hip:
|
||||
# navi, CDNA1-CDNA3 allows a max of 64KB shared memory
|
||||
# CDNA4 allows a max of 160KB shared memory
|
||||
max_shared_mem = (
|
||||
65536 if device_props.gcnArchName != "gfx950" else 160 * 1024
|
||||
)
|
||||
else:
|
||||
max_shared_mem = getattr(
|
||||
device_props, "shared_memory_per_block_optin", 49152
|
||||
)
|
||||
|
||||
if shared_mem_bytes > max_shared_mem:
|
||||
raise RuntimeError(
|
||||
f"Requested shared memory ({shared_mem_bytes} bytes) exceeds "
|
||||
f"device limit ({max_shared_mem} bytes). "
|
||||
"Consider reducing block size or shared memory usage."
|
||||
)
|
||||
|
||||
# Set the function attribute once
|
||||
# https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize = 8
|
||||
_check_cuda(
|
||||
libcuda.cuFuncSetAttribute(
|
||||
self.func,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
shared_mem_bytes,
|
||||
)
|
||||
)
|
||||
|
||||
self._max_shared_mem_bytes = shared_mem_bytes
|
||||
|
||||
|
||||
def _cuda_load_module(
|
||||
ptx: str | bytes, kernel_names: list[str] | None = None
|
||||
) -> _CudaModule | dict[str, "_CudaKernel"]:
|
||||
"""
|
||||
Loads a CUDA module from PTX code and returns a module object that can access kernels.
|
||||
|
||||
Args:
|
||||
ptx (bytes or str): The PTX code to load
|
||||
kernel_names (list, optional): List of kernel names to extract from the module.
|
||||
If None, will return a module object with __getattr__.
|
||||
|
||||
Returns:
|
||||
object: If kernel_names is None, returns a module object with __getattr__ to access kernels.
|
||||
If kernel_names is provided, returns a dict mapping kernel names to _CudaKernel objects.
|
||||
"""
|
||||
# Ensure CUDA is initialized
|
||||
import torch.cuda
|
||||
|
||||
# Load CUDA driver library
|
||||
libcuda = _get_gpu_runtime_library()
|
||||
|
||||
# Convert PTX to bytes if it's a string
|
||||
if isinstance(ptx, str):
|
||||
ptx = ptx.encode("utf-8")
|
||||
|
||||
# Load PTX module
|
||||
module = ctypes.c_void_p()
|
||||
# Get the current stream without directly importing torch.cuda at module level
|
||||
stream = torch.cuda.current_stream()
|
||||
with stream:
|
||||
_check_cuda(libcuda.cuModuleLoadData(ctypes.byref(module), ptx))
|
||||
|
||||
if not kernel_names:
|
||||
return _CudaModule(module)
|
||||
|
||||
# Return specific kernels
|
||||
kernels = {}
|
||||
for name in kernel_names:
|
||||
func = ctypes.c_void_p()
|
||||
_check_cuda(
|
||||
libcuda.cuModuleGetFunction(
|
||||
ctypes.byref(func), module, name.encode("utf-8")
|
||||
)
|
||||
)
|
||||
kernels[name] = _CudaKernel(func, module)
|
||||
return kernels
|
||||
|
||||
|
||||
def _get_device_index(
|
||||
device: Any, optional: bool = False, allow_cpu: bool = False
|
||||
) -> int:
|
||||
r"""Get the device index from :attr:`device`, which can be a torch.device object, a Python integer, or ``None``.
|
||||
|
||||
If :attr:`device` is a torch.device object, returns the device index if it
|
||||
is a CUDA device. Note that for a CUDA device without a specified index,
|
||||
i.e., ``torch.device('cuda')``, this will return the current default CUDA
|
||||
device if :attr:`optional` is ``True``. If :attr:`allow_cpu` is ``True``,
|
||||
CPU devices will be accepted and ``-1`` will be returned in this case.
|
||||
|
||||
If :attr:`device` is a Python integer, it is returned as is.
|
||||
|
||||
If :attr:`device` is ``None``, this will return the current default CUDA
|
||||
device if :attr:`optional` is ``True``.
|
||||
"""
|
||||
if isinstance(device, int):
|
||||
return device
|
||||
if isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
if isinstance(device, torch.device):
|
||||
if allow_cpu:
|
||||
if device.type not in ["cuda", "cpu"]:
|
||||
raise ValueError(f"Expected a cuda or cpu device, but got: {device}")
|
||||
elif device.type != "cuda":
|
||||
raise ValueError(f"Expected a cuda device, but got: {device}")
|
||||
if not torch.jit.is_scripting():
|
||||
if isinstance(device, torch.cuda.device):
|
||||
return device.idx
|
||||
return _torch_get_device_index(device, optional, allow_cpu)
|
||||
@@ -0,0 +1,13 @@
|
||||
# pyrefly: ignore [deprecated]
|
||||
from .autocast_mode import autocast, custom_bwd, custom_fwd
|
||||
from .common import amp_definitely_not_available
|
||||
from .grad_scaler import GradScaler
|
||||
|
||||
|
||||
__all__ = [
|
||||
"amp_definitely_not_available",
|
||||
"autocast",
|
||||
"custom_bwd",
|
||||
"custom_fwd",
|
||||
"GradScaler",
|
||||
]
|
||||
@@ -0,0 +1,110 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
import sys
|
||||
from typing import Any
|
||||
from typing_extensions import deprecated
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = ["autocast", "custom_fwd", "custom_bwd"]
|
||||
|
||||
|
||||
@deprecated(
|
||||
"`torch.cuda.amp.autocast(args...)` is deprecated. "
|
||||
"Please use `torch.amp.autocast('cuda', args...)` instead.",
|
||||
category=FutureWarning,
|
||||
)
|
||||
class autocast(torch.amp.autocast_mode.autocast):
|
||||
r"""See :class:`torch.autocast`.
|
||||
|
||||
``torch.cuda.amp.autocast(args...)`` is deprecated. Please use ``torch.amp.autocast("cuda", args...)`` instead.
|
||||
"""
|
||||
|
||||
# TODO: remove this conditional once we stop supporting Python < 3.13
|
||||
# Prior to Python 3.13, inspect.signature could not retrieve the correct
|
||||
# signature information for classes decorated with @deprecated (unless
|
||||
# the __new__ static method was explicitly defined);
|
||||
#
|
||||
# However, this issue has been fixed in Python 3.13 and later versions.
|
||||
if sys.version_info < (3, 13):
|
||||
|
||||
def __new__(
|
||||
cls,
|
||||
enabled: bool = True,
|
||||
dtype: torch.dtype = torch.float16,
|
||||
cache_enabled: bool = True,
|
||||
):
|
||||
return super().__new__(cls)
|
||||
|
||||
def __init_subclass__(cls):
|
||||
pass
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
enabled: bool = True,
|
||||
dtype: torch.dtype = torch.float16,
|
||||
cache_enabled: bool = True,
|
||||
):
|
||||
if torch._jit_internal.is_scripting():
|
||||
self._enabled = enabled
|
||||
self.device = "cuda"
|
||||
self.fast_dtype = dtype
|
||||
return
|
||||
super().__init__(
|
||||
"cuda", enabled=enabled, dtype=dtype, cache_enabled=cache_enabled
|
||||
)
|
||||
|
||||
def __enter__(self):
|
||||
if torch._jit_internal.is_scripting():
|
||||
return self
|
||||
return super().__enter__()
|
||||
|
||||
# TODO: discuss a unified TorchScript-friendly API for autocast
|
||||
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): # type: ignore[override]
|
||||
if torch._jit_internal.is_scripting():
|
||||
return
|
||||
return super().__exit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
def __call__(self, func):
|
||||
if torch._jit_internal.is_scripting():
|
||||
return func
|
||||
return super().__call__(func)
|
||||
|
||||
|
||||
# Preserved only for BC reasons
|
||||
@deprecated(
|
||||
"`torch.cuda.amp.autocast_mode._cast(value, dtype)` is deprecated. "
|
||||
"Please use `torch.amp.autocast_mode._cast(value, 'cuda', dtype)` instead.",
|
||||
category=FutureWarning,
|
||||
)
|
||||
def _cast(value, dtype):
|
||||
return torch.amp.autocast_mode._cast(value, "cuda", dtype)
|
||||
|
||||
|
||||
@deprecated(
|
||||
"`torch.cuda.amp.custom_fwd(args...)` is deprecated. "
|
||||
"Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.",
|
||||
category=FutureWarning,
|
||||
)
|
||||
def custom_fwd(fwd=None, *, cast_inputs=None):
|
||||
"""
|
||||
``torch.cuda.amp.custom_fwd(args...)`` is deprecated. Please use
|
||||
``torch.amp.custom_fwd(args..., device_type='cuda')`` instead.
|
||||
"""
|
||||
return functools.partial(torch.amp.custom_fwd, device_type="cuda")(
|
||||
fwd=fwd, cast_inputs=cast_inputs
|
||||
)
|
||||
|
||||
|
||||
@deprecated(
|
||||
"`torch.cuda.amp.custom_bwd(args...)` is deprecated. "
|
||||
"Please use `torch.amp.custom_bwd(args..., device_type='cuda')` instead.",
|
||||
category=FutureWarning,
|
||||
)
|
||||
def custom_bwd(bwd):
|
||||
"""
|
||||
``torch.cuda.amp.custom_bwd(args...)`` is deprecated. Please use
|
||||
``torch.amp.custom_bwd(args..., device_type='cuda')`` instead.
|
||||
"""
|
||||
return functools.partial(torch.amp.custom_bwd, device_type="cuda")(bwd)
|
||||
@@ -0,0 +1,11 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from importlib.util import find_spec
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = ["amp_definitely_not_available"]
|
||||
|
||||
|
||||
def amp_definitely_not_available():
|
||||
return not (torch.cuda.is_available() or find_spec("torch_xla"))
|
||||
@@ -0,0 +1,38 @@
|
||||
from typing_extensions import deprecated
|
||||
|
||||
import torch
|
||||
|
||||
# We need to keep this unused import for BC reasons
|
||||
from torch.amp.grad_scaler import OptState # noqa: F401
|
||||
|
||||
|
||||
__all__ = ["GradScaler"]
|
||||
|
||||
|
||||
class GradScaler(torch.amp.GradScaler):
|
||||
r"""
|
||||
See :class:`torch.amp.GradScaler`.
|
||||
``torch.cuda.amp.GradScaler(args...)`` is deprecated. Please use ``torch.amp.GradScaler("cuda", args...)`` instead.
|
||||
"""
|
||||
|
||||
@deprecated(
|
||||
"`torch.cuda.amp.GradScaler(args...)` is deprecated. "
|
||||
"Please use `torch.amp.GradScaler('cuda', args...)` instead.",
|
||||
category=FutureWarning,
|
||||
)
|
||||
def __init__(
|
||||
self,
|
||||
init_scale: float = 2.0**16,
|
||||
growth_factor: float = 2.0,
|
||||
backoff_factor: float = 0.5,
|
||||
growth_interval: int = 2000,
|
||||
enabled: bool = True,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
"cuda",
|
||||
init_scale=init_scale,
|
||||
growth_factor=growth_factor,
|
||||
backoff_factor=backoff_factor,
|
||||
growth_interval=growth_interval,
|
||||
enabled=enabled,
|
||||
)
|
||||
@@ -0,0 +1,19 @@
|
||||
# The functions here have been moved to torch.nn.parallel.comm
|
||||
from torch.nn.parallel.comm import (
|
||||
broadcast,
|
||||
broadcast_coalesced,
|
||||
gather,
|
||||
reduce_add,
|
||||
reduce_add_coalesced,
|
||||
scatter,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"broadcast",
|
||||
"broadcast_coalesced",
|
||||
"reduce_add",
|
||||
"reduce_add_coalesced",
|
||||
"scatter",
|
||||
"gather",
|
||||
]
|
||||
@@ -0,0 +1,177 @@
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from torch.types import Storage
|
||||
|
||||
|
||||
__all__: list[str] = [
|
||||
"gds_register_buffer",
|
||||
"gds_deregister_buffer",
|
||||
"GdsFile",
|
||||
]
|
||||
|
||||
|
||||
def _dummy_fn(name: str) -> Callable:
|
||||
def fn(*args, **kwargs): # type: ignore[no-untyped-def]
|
||||
raise RuntimeError(f"torch._C.{name} is not supported on this platform")
|
||||
|
||||
return fn
|
||||
|
||||
|
||||
if not hasattr(torch._C, "_gds_register_buffer"):
|
||||
if hasattr(torch._C, "_gds_deregister_buffer"):
|
||||
raise AssertionError(
|
||||
"_gds_deregister_buffer exists but _gds_register_buffer does not"
|
||||
)
|
||||
if hasattr(torch._C, "_gds_register_handle"):
|
||||
raise AssertionError(
|
||||
"_gds_register_handle exists but _gds_register_buffer does not"
|
||||
)
|
||||
if hasattr(torch._C, "_gds_deregister_handle"):
|
||||
raise AssertionError(
|
||||
"_gds_deregister_handle exists but _gds_register_buffer does not"
|
||||
)
|
||||
if hasattr(torch._C, "_gds_load_storage"):
|
||||
raise AssertionError(
|
||||
"_gds_load_storage exists but _gds_register_buffer does not"
|
||||
)
|
||||
if hasattr(torch._C, "_gds_save_storage"):
|
||||
raise AssertionError(
|
||||
"_gds_save_storage exists but _gds_register_buffer does not"
|
||||
)
|
||||
# Define functions
|
||||
torch._C.__dict__["_gds_register_buffer"] = _dummy_fn("_gds_register_buffer")
|
||||
torch._C.__dict__["_gds_deregister_buffer"] = _dummy_fn("_gds_deregister_buffer")
|
||||
torch._C.__dict__["_gds_register_handle"] = _dummy_fn("_gds_register_handle")
|
||||
torch._C.__dict__["_gds_deregister_handle"] = _dummy_fn("_gds_deregister_handle")
|
||||
torch._C.__dict__["_gds_load_storage"] = _dummy_fn("_gds_load_storage")
|
||||
torch._C.__dict__["_gds_save_storage"] = _dummy_fn("_gds_save_storage")
|
||||
|
||||
|
||||
def gds_register_buffer(s: Storage) -> None:
|
||||
"""Registers a storage on a CUDA device as a cufile buffer.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +SKIP("gds filesystem requirements")
|
||||
>>> src = torch.randn(1024, device="cuda")
|
||||
>>> s = src.untyped_storage()
|
||||
>>> gds_register_buffer(s)
|
||||
|
||||
Args:
|
||||
s (Storage): Buffer to register.
|
||||
"""
|
||||
torch._C._gds_register_buffer(s)
|
||||
|
||||
|
||||
def gds_deregister_buffer(s: Storage) -> None:
|
||||
"""Deregisters a previously registered storage on a CUDA device as a cufile buffer.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +SKIP("gds filesystem requirements")
|
||||
>>> src = torch.randn(1024, device="cuda")
|
||||
>>> s = src.untyped_storage()
|
||||
>>> gds_register_buffer(s)
|
||||
>>> gds_deregister_buffer(s)
|
||||
|
||||
Args:
|
||||
s (Storage): Buffer to register.
|
||||
"""
|
||||
torch._C._gds_deregister_buffer(s)
|
||||
|
||||
|
||||
class GdsFile:
|
||||
r"""Wrapper around cuFile.
|
||||
|
||||
cuFile is a file-like interface to the GPUDirect Storage (GDS) API.
|
||||
|
||||
See the `cufile docs <https://docs.nvidia.com/gpudirect-storage/api-reference-guide/index.html#cufile-io-api>`_
|
||||
for more details.
|
||||
|
||||
Args:
|
||||
filename (str): Name of the file to open.
|
||||
flags (int): Flags to pass to ``os.open`` when opening the file. ``os.O_DIRECT`` will
|
||||
be added automatically.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +SKIP("gds filesystem requirements")
|
||||
>>> src1 = torch.randn(1024, device="cuda")
|
||||
>>> src2 = torch.randn(2, 1024, device="cuda")
|
||||
>>> file = torch.cuda.gds.GdsFile(f, os.O_CREAT | os.O_RDWR)
|
||||
>>> file.save_storage(src1.untyped_storage(), offset=0)
|
||||
>>> file.save_storage(src2.untyped_storage(), offset=src1.nbytes)
|
||||
>>> dest1 = torch.empty(1024, device="cuda")
|
||||
>>> dest2 = torch.empty(2, 1024, device="cuda")
|
||||
>>> file.load_storage(dest1.untyped_storage(), offset=0)
|
||||
>>> file.load_storage(dest2.untyped_storage(), offset=src1.nbytes)
|
||||
>>> torch.equal(src1, dest1)
|
||||
True
|
||||
>>> torch.equal(src2, dest2)
|
||||
True
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, filename: str, flags: int):
|
||||
if sys.platform == "win32":
|
||||
raise RuntimeError("GdsFile is not supported on this platform.")
|
||||
self.filename = filename
|
||||
self.flags = flags
|
||||
self.fd = os.open(filename, flags | os.O_DIRECT) # type: ignore[attr-defined]
|
||||
self.handle: int | None = None
|
||||
self.register_handle()
|
||||
|
||||
def __del__(self) -> None:
|
||||
if self.handle is not None:
|
||||
self.deregister_handle()
|
||||
os.close(self.fd)
|
||||
|
||||
def register_handle(self) -> None:
|
||||
"""Registers file descriptor to cuFile Driver.
|
||||
|
||||
This is a wrapper around ``cuFileHandleRegister``.
|
||||
"""
|
||||
if self.handle is not None:
|
||||
raise AssertionError("Cannot register a handle that is already registered.")
|
||||
self.handle = torch._C._gds_register_handle(self.fd)
|
||||
|
||||
def deregister_handle(self) -> None:
|
||||
"""Deregisters file descriptor from cuFile Driver.
|
||||
|
||||
This is a wrapper around ``cuFileHandleDeregister``.
|
||||
"""
|
||||
if self.handle is None:
|
||||
raise AssertionError("Cannot deregister a handle that is not registered.")
|
||||
torch._C._gds_deregister_handle(self.handle)
|
||||
self.handle = None
|
||||
|
||||
def load_storage(self, storage: Storage, offset: int = 0) -> None:
|
||||
"""Loads data from the file into the storage.
|
||||
|
||||
This is a wrapper around ``cuFileRead``. ``storage.nbytes()`` of data
|
||||
will be loaded from the file at ``offset`` into the storage.
|
||||
|
||||
Args:
|
||||
storage (Storage): Storage to load data into.
|
||||
offset (int, optional): Offset into the file to start loading from. (Default: 0)
|
||||
"""
|
||||
if self.handle is None:
|
||||
raise AssertionError("Cannot load data from a file that is not registered.")
|
||||
torch._C._gds_load_storage(self.handle, storage, offset)
|
||||
|
||||
def save_storage(self, storage: Storage, offset: int = 0) -> None:
|
||||
"""Saves data from the storage into the file.
|
||||
|
||||
This is a wrapper around ``cuFileWrite``. All bytes of the storage
|
||||
will be written to the file at ``offset``.
|
||||
|
||||
Args:
|
||||
storage (Storage): Storage to save data from.
|
||||
offset (int, optional): Offset into the file to start saving to. (Default: 0)
|
||||
"""
|
||||
if self.handle is None:
|
||||
raise AssertionError("Cannot save data to a file that is not registered.")
|
||||
torch._C._gds_save_storage(self.handle, storage, offset)
|
||||
@@ -0,0 +1,645 @@
|
||||
# pylint: disable=useless-parent-delegation
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import typing
|
||||
from collections.abc import Callable
|
||||
from typing import overload, TYPE_CHECKING, TypeAlias, Union
|
||||
from typing_extensions import ParamSpec, Self, TypeVar
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# importing _POOL_HANDLE at runtime toplevel causes an import cycle
|
||||
from torch.cuda import _POOL_HANDLE
|
||||
|
||||
from .._utils import _dummy_type
|
||||
|
||||
|
||||
__all__ = [
|
||||
"is_current_stream_capturing",
|
||||
"graph_pool_handle",
|
||||
"CUDAGraph",
|
||||
"graph",
|
||||
"make_graphed_callables",
|
||||
]
|
||||
|
||||
|
||||
_R = TypeVar("_R")
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
|
||||
if not hasattr(torch._C, "_CudaStreamBase"):
|
||||
# Define dummy base classes
|
||||
torch._C.__dict__["_CUDAGraph"] = _dummy_type("_CUDAGraph")
|
||||
torch._C.__dict__["_graph_pool_handle"] = _dummy_type("_graph_pool_handle")
|
||||
torch._C.__dict__["_cuda_isCurrentStreamCapturing"] = _dummy_type(
|
||||
"_cuda_isCurrentStreamCapturing"
|
||||
)
|
||||
|
||||
from torch._C import _cuda_isCurrentStreamCapturing, _CUDAGraph, _graph_pool_handle
|
||||
|
||||
|
||||
def is_current_stream_capturing() -> bool:
|
||||
r"""Return True if CUDA graph capture is underway on the current CUDA stream, False otherwise.
|
||||
|
||||
If a CUDA context does not exist on the current device, returns False without initializing the context.
|
||||
"""
|
||||
return _cuda_isCurrentStreamCapturing()
|
||||
|
||||
|
||||
# Python shim helps Sphinx process docstrings more reliably.
|
||||
def graph_pool_handle() -> _POOL_HANDLE:
|
||||
r"""Return an opaque token representing the id of a graph memory pool.
|
||||
|
||||
See :ref:`Graph memory management<graph-memory-management>`.
|
||||
|
||||
.. warning::
|
||||
This API is in beta and may change in future releases.
|
||||
"""
|
||||
return torch.cuda._POOL_HANDLE(_graph_pool_handle())
|
||||
|
||||
|
||||
# Python shim helps Sphinx process docstrings more reliably.
|
||||
class CUDAGraph(_CUDAGraph):
|
||||
r"""Wrapper around a CUDA graph.
|
||||
|
||||
Arguments:
|
||||
keep_graph (bool, optional): If ``keep_graph=False``, the
|
||||
cudaGraphExec_t will be instantiated on GPU at the end of
|
||||
``capture_end`` and the underlying cudaGraph_t will be
|
||||
destroyed. Users who want to query or otherwise modify the
|
||||
underlying cudaGraph_t before instantiation can set
|
||||
``keep_graph=True`` and access it via ``raw_cuda_graph`` after
|
||||
``capture_end``. Note that the cudaGraphExec_t will not be
|
||||
instantiated at the end of ``capture_end`` in this
|
||||
case. Instead, it will be instantiated via an explicit called
|
||||
to ``instantiate`` or automatically on the first call to
|
||||
``replay`` if ``instantiate`` was not already called. Calling
|
||||
``instantiate`` manually before ``replay`` is recommended to
|
||||
prevent increased latency on the first call to ``replay``. It
|
||||
is allowed to modify the raw cudaGraph_t after first calling
|
||||
``instantiate``, but the user must call ``instantiate`` again
|
||||
manually to make sure the instantiated graph has these
|
||||
changes. Pytorch has no means of tracking these changes.
|
||||
|
||||
.. warning::
|
||||
This API is in beta and may change in future releases.
|
||||
|
||||
"""
|
||||
|
||||
def __new__(cls, keep_graph: bool = False) -> Self:
|
||||
return super().__new__(cls, keep_graph)
|
||||
|
||||
def capture_begin(
|
||||
self, pool: _POOL_HANDLE | None = None, capture_error_mode: str = "global"
|
||||
) -> None:
|
||||
r"""Begin capturing CUDA work on the current stream.
|
||||
|
||||
Typically, you shouldn't call ``capture_begin`` yourself.
|
||||
Use :class:`~torch.cuda.graph` or :func:`~torch.cuda.make_graphed_callables`,
|
||||
which call ``capture_begin`` internally.
|
||||
|
||||
Arguments:
|
||||
pool (optional): Token (returned by :func:`~torch.cuda.graph_pool_handle` or
|
||||
:meth:`other_Graph_instance.pool()<torch.cuda.CUDAGraph.pool>`) that hints this graph may share memory
|
||||
with the indicated pool. See :ref:`Graph memory management<graph-memory-management>`.
|
||||
capture_error_mode (str, optional): specifies the cudaStreamCaptureMode for the graph capture stream.
|
||||
Can be "global", "thread_local" or "relaxed". During cuda graph capture, some actions, such as cudaMalloc,
|
||||
may be unsafe. "global" will error on actions in other threads, "thread_local" will only error for
|
||||
actions in the current thread, and "relaxed" will not error on these actions. Do NOT change this setting
|
||||
unless you're familiar with `cudaStreamCaptureMode <https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1g9d0535d93a214cbf126835257b16ba85>`_
|
||||
""" # noqa: B950
|
||||
super().capture_begin(pool=pool, capture_error_mode=capture_error_mode)
|
||||
|
||||
def capture_end(self) -> None:
|
||||
r"""End CUDA graph capture on the current stream.
|
||||
|
||||
After ``capture_end``, ``replay`` may be called on this instance.
|
||||
|
||||
Typically, you shouldn't call ``capture_end`` yourself.
|
||||
Use :class:`~torch.cuda.graph` or :func:`~torch.cuda.make_graphed_callables`,
|
||||
which call ``capture_end`` internally.
|
||||
"""
|
||||
super().capture_end()
|
||||
|
||||
def instantiate(self) -> None:
|
||||
r"""Instantiate the CUDA graph. Will be called by
|
||||
``capture_end`` if ``keep_graph=False``, or by ``replay`` if
|
||||
``keep_graph=True`` and ``instantiate`` has not already been
|
||||
explicitly called. Does not destroy the cudaGraph_t returned
|
||||
by ``raw_cuda_graph``.
|
||||
"""
|
||||
super().instantiate()
|
||||
|
||||
def replay(self) -> None:
|
||||
r"""Replay the CUDA work captured by this graph."""
|
||||
super().replay()
|
||||
|
||||
def reset(self) -> None:
|
||||
r"""Delete the graph currently held by this instance."""
|
||||
super().reset()
|
||||
|
||||
def pool(self) -> _POOL_HANDLE:
|
||||
r"""Return an opaque token representing the id of this graph's memory pool.
|
||||
|
||||
This id can optionally be passed to another graph's ``capture_begin``,
|
||||
which hints the other graph may share the same memory pool.
|
||||
"""
|
||||
return super().pool()
|
||||
|
||||
def enable_debug_mode(self) -> None:
|
||||
r"""Enable debugging mode for CUDAGraph.debug_dump."""
|
||||
return super().enable_debug_mode()
|
||||
|
||||
def debug_dump(self, debug_path: str) -> None:
|
||||
r"""
|
||||
Arguments:
|
||||
debug_path (required): Path to dump the graph to.
|
||||
|
||||
Calls a debugging function to dump the graph if the debugging is
|
||||
enabled via CUDAGraph.enable_debug_mode()
|
||||
"""
|
||||
return super().debug_dump(debug_path)
|
||||
|
||||
def raw_cuda_graph(self) -> int:
|
||||
r"""Returns the underlying cudaGraph_t. ``keep_graph`` must be True.
|
||||
|
||||
See the following for APIs for how to manipulate this object: `Graph Managmement <https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__GRAPH.html>`_ and `cuda-python Graph Management bindings <https://nvidia.github.io/cuda-python/cuda-bindings/latest/module/runtime.html#graph-management>`_
|
||||
""" # noqa: B950
|
||||
return super().raw_cuda_graph()
|
||||
|
||||
def raw_cuda_graph_exec(self) -> int:
|
||||
r"""Returns the underlying cudaGraphExec_t. ``instantiate`` must have been called if ``keep_graph`` is True, or ``capture_end`` must have been called if ``keep_graph`` is False. If you call ``instantiate()`` after ``raw_cuda_graph_exec()``, the previously returned cudaGraphExec_t will be destroyed. It is your responsibility not to use this object after destruction.
|
||||
|
||||
See the following for APIs for how to manipulate this object: `Graph Execution <https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__GRAPH__EXEC.html>`_ and `cuda-python Graph Execution bindings <https://nvidia.github.io/cuda-python/cuda-bindings/latest/module/runtime.html#graph-execution>`_
|
||||
""" # noqa: B950
|
||||
return super().raw_cuda_graph_exec()
|
||||
|
||||
|
||||
class graph:
|
||||
r"""Context-manager that captures CUDA work into a :class:`torch.cuda.CUDAGraph` object for later replay.
|
||||
|
||||
See :ref:`CUDA Graphs <cuda-graph-semantics>` for a general introduction,
|
||||
detailed use, and constraints.
|
||||
|
||||
Arguments:
|
||||
cuda_graph (torch.cuda.CUDAGraph): Graph object used for capture.
|
||||
pool (optional): Opaque token (returned by a call to :func:`~torch.cuda.graph_pool_handle()` or
|
||||
:meth:`other_Graph_instance.pool()<torch.cuda.CUDAGraph.pool>`) hinting this graph's capture
|
||||
may share memory from the specified pool. See :ref:`Graph memory management<graph-memory-management>`.
|
||||
stream (torch.cuda.Stream, optional): If supplied, will be set as the current stream in the context.
|
||||
If not supplied, ``graph`` sets its own internal side stream as the current stream in the context.
|
||||
capture_error_mode (str, optional): specifies the cudaStreamCaptureMode for the graph capture stream.
|
||||
Can be "global", "thread_local" or "relaxed". During cuda graph capture, some actions, such as cudaMalloc,
|
||||
may be unsafe. "global" will error on actions in other threads, "thread_local" will only error for
|
||||
actions in the current thread, and "relaxed" will not error on actions. Do NOT change this setting
|
||||
unless you're familiar with `cudaStreamCaptureMode <https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1g9d0535d93a214cbf126835257b16ba85>`_
|
||||
enable_annotations (bool, optional): If ``True``, enables kernel annotation
|
||||
recording on entry and automatically calls
|
||||
:func:`~torch.cuda._graph_annotations.resolve_pending_annotations` before
|
||||
the capture ends. Annotations are **not** cleared on exit so that multiple
|
||||
graphs in the same workload can accumulate annotations.
|
||||
Requires ``cuda.bindings`` package and cuda-compat >= 13.1 or CUDA driver >= 13.1.
|
||||
|
||||
.. note::
|
||||
For effective memory sharing, if you pass a ``pool`` used by a previous capture and the previous capture
|
||||
used an explicit ``stream`` argument, you should pass the same ``stream`` argument to this capture.
|
||||
|
||||
.. warning::
|
||||
This API is in beta and may change in future releases.
|
||||
|
||||
.. _cudaStreamCaptureMode:
|
||||
https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1g9d0535d93a214cbf126835257b16ba85
|
||||
""" # noqa: B950
|
||||
|
||||
default_capture_stream: torch.cuda.Stream | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cuda_graph: CUDAGraph,
|
||||
pool: _POOL_HANDLE | None = None,
|
||||
stream: torch.cuda.Stream | None = None,
|
||||
capture_error_mode: str = "global",
|
||||
enable_annotations: bool = False,
|
||||
):
|
||||
# Lazy-init of default_capture_stream helps avoid circular-import errors.
|
||||
# Not thread safe, but graphs already have the general (explicitly documented)
|
||||
# restriction that only one capture may be underway at a time in the process.
|
||||
if stream is None and self.__class__.default_capture_stream is None:
|
||||
self.__class__.default_capture_stream = torch.cuda.Stream()
|
||||
|
||||
self.pool: tuple[()] | tuple[_POOL_HANDLE] = () if pool is None else (pool,)
|
||||
self.capture_stream = (
|
||||
stream if stream is not None else self.__class__.default_capture_stream
|
||||
)
|
||||
if self.capture_stream is None:
|
||||
raise AssertionError("capture_stream must not be None")
|
||||
self.stream_ctx = torch.cuda.stream(self.capture_stream)
|
||||
self.cuda_graph = cuda_graph
|
||||
self.capture_error_mode = capture_error_mode
|
||||
self._enable_annotations = enable_annotations
|
||||
|
||||
def __enter__(self) -> None:
|
||||
# Free as much memory as we can for the graph
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if torch.compiler.config.force_cudagraph_gc:
|
||||
# Originally we unconditionally garbage collected here. On one hand
|
||||
# that's nice because we have a chance to collect more memory, but
|
||||
# on the other hand it is REALLY expensive, especially for doing
|
||||
# multiple cudagraph captures in a row. In theory it will only help
|
||||
# when a dead python cycle is holding onto CUDA memory.
|
||||
gc.collect()
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
torch._C._host_emptyCache()
|
||||
|
||||
if self._enable_annotations:
|
||||
from torch.cuda._graph_annotations import enable_annotations as _enable_ann
|
||||
|
||||
_enable_ann()
|
||||
|
||||
# Stackoverflow seems comfortable with this pattern
|
||||
# https://stackoverflow.com/questions/26635684/calling-enter-and-exit-manually#39172487
|
||||
self.stream_ctx.__enter__()
|
||||
|
||||
self.cuda_graph.capture_begin(
|
||||
# type: ignore[misc]
|
||||
*self.pool,
|
||||
# pyrefly: ignore [bad-keyword-argument]
|
||||
capture_error_mode=self.capture_error_mode,
|
||||
)
|
||||
|
||||
def __exit__(self, *args: object) -> None:
|
||||
if self._enable_annotations:
|
||||
from torch.cuda._graph_annotations import resolve_pending_annotations
|
||||
|
||||
resolve_pending_annotations()
|
||||
|
||||
self.cuda_graph.capture_end()
|
||||
self.stream_ctx.__exit__(*args)
|
||||
|
||||
if self._enable_annotations:
|
||||
from torch.cuda._graph_annotations import remap_to_exec_graph
|
||||
|
||||
remap_to_exec_graph(self.cuda_graph)
|
||||
# returning None should propagate exceptions from either capture_end or stream_ctx.__exit__()
|
||||
|
||||
|
||||
_ModuleOrCallable: TypeAlias = Union["torch.nn.Module", Callable[..., object]]
|
||||
|
||||
|
||||
@overload
|
||||
def make_graphed_callables(
|
||||
callables: _ModuleOrCallable,
|
||||
sample_args: tuple[Tensor, ...],
|
||||
num_warmup_iters: int = 3,
|
||||
allow_unused_input: bool = False,
|
||||
pool: _POOL_HANDLE | None = None,
|
||||
) -> _ModuleOrCallable: ...
|
||||
|
||||
|
||||
@overload
|
||||
def make_graphed_callables(
|
||||
callables: tuple[_ModuleOrCallable, ...],
|
||||
sample_args: tuple[tuple[Tensor, ...], ...],
|
||||
num_warmup_iters: int = 3,
|
||||
allow_unused_input: bool = False,
|
||||
pool: _POOL_HANDLE | None = None,
|
||||
) -> tuple[_ModuleOrCallable, ...]: ...
|
||||
|
||||
|
||||
def make_graphed_callables(
|
||||
callables: _ModuleOrCallable | tuple[_ModuleOrCallable, ...],
|
||||
sample_args: tuple[Tensor, ...] | tuple[tuple[Tensor, ...], ...],
|
||||
num_warmup_iters: int = 3,
|
||||
allow_unused_input: bool = False,
|
||||
pool: _POOL_HANDLE | None = None,
|
||||
) -> _ModuleOrCallable | tuple[_ModuleOrCallable, ...]:
|
||||
r"""Accept callables (functions or :class:`nn.Module<torch.nn.Module>`\ s) and returns graphed versions.
|
||||
|
||||
Each graphed callable's forward pass runs its source callable's
|
||||
forward CUDA work as a CUDA graph inside a single autograd node.
|
||||
|
||||
The graphed callable's forward pass also appends
|
||||
a backward node to the autograd graph. During backward, this node runs the
|
||||
callable's backward work as a CUDA graph.
|
||||
|
||||
Therefore, each graphed callable should be a drop-in replacement for its source callable
|
||||
in an autograd-enabled training loop.
|
||||
|
||||
See :ref:`Partial-network capture<partial-network-capture>` for detailed use and constraints.
|
||||
|
||||
If you pass a tuple of several callables, their captures will use the same memory pool.
|
||||
See :ref:`Graph memory management<graph-memory-management>` for when this is appropriate.
|
||||
|
||||
Arguments:
|
||||
callables (torch.nn.Module or Python function, or tuple of these): Callable or callables to graph.
|
||||
See :ref:`Graph memory management<graph-memory-management>` for when passing a tuple of callables
|
||||
is appropriate. If you pass a tuple of callables, their order in the tuple must be the same order
|
||||
they'll run in the live workload.
|
||||
sample_args (tuple of Tensors, or tuple of tuples of Tensors): Samples args for each callable.
|
||||
If a single callable was passed, ``sample_args`` must be a single tuple of argument Tensors.
|
||||
If a tuple of callables was passed, ``sample_args`` must be tuple of tuples of argument Tensors.
|
||||
num_warmup_iters (int): The number of warmup iterations. Currently, ``DataDistributedParallel`` needs
|
||||
11 iterations for warm up. Default: ``3``.
|
||||
allow_unused_input (bool): If False, specifying inputs that were not used when computing outputs
|
||||
(and therefore their grad is always zero) is an error. Defaults to False.
|
||||
pool (optional): Token (returned by :func:`~torch.cuda.graph_pool_handle` or
|
||||
:meth:`other_Graph_instance.pool()<torch.cuda.CUDAGraph.pool>`) that hints this graph may share memory
|
||||
with the indicated pool. See :ref:`Graph memory management<graph-memory-management>`.
|
||||
.. note::
|
||||
The ``requires_grad`` state of each Tensor in ``sample_args`` must match the state
|
||||
that's expected for the corresponding real input in the training loop.
|
||||
|
||||
.. warning::
|
||||
This API is in beta and may change in future releases.
|
||||
|
||||
.. warning::
|
||||
``sample_args`` for each callable must contain only Tensors. Other types are not allowed.
|
||||
|
||||
.. warning::
|
||||
Returned callables do not support higher order differentiation (e.g., double backward).
|
||||
|
||||
.. warning::
|
||||
In any :class:`~torch.nn.Module` passed to :func:`~make_graphed_callables`, only parameters
|
||||
may be trainable. Buffers must have ``requires_grad=False``.
|
||||
|
||||
.. warning::
|
||||
After you pass a :class:`torch.nn.Module` through :func:`~make_graphed_callables`,
|
||||
you may not add or remove any of that Module's parameters or buffers.
|
||||
|
||||
.. warning::
|
||||
:class:`torch.nn.Module`\s passed to :func:`~torch.cuda.make_graphed_callables` must not have module hooks
|
||||
registered on them at the time they are passed. However, registering hooks on modules *after* passing them
|
||||
through :func:`~torch.cuda.make_graphed_callables` is allowed.
|
||||
|
||||
.. warning::
|
||||
When running a graphed callable, you must pass its arguments in the same order and format
|
||||
they appeared in that callable's ``sample_args``.
|
||||
|
||||
.. warning::
|
||||
The automatic mixed precision is supported in :func:`~torch.cuda.make_graphed_callables` only with disabled
|
||||
caching. The context manager `torch.cuda.amp.autocast()` must have `cache_enabled=False`.
|
||||
"""
|
||||
if torch.is_autocast_enabled() and torch.is_autocast_cache_enabled():
|
||||
raise RuntimeError(
|
||||
"make_graphed_callables does not support the autocast caching. Please set `cache_enabled=False`."
|
||||
)
|
||||
|
||||
just_one_callable = False
|
||||
|
||||
_sample_args: tuple[tuple[Tensor, ...], ...]
|
||||
if not isinstance(callables, tuple):
|
||||
just_one_callable = True
|
||||
callables = (callables,)
|
||||
_sample_args = (typing.cast(tuple[Tensor, ...], sample_args),)
|
||||
else:
|
||||
_sample_args = typing.cast(tuple[tuple[Tensor, ...], ...], sample_args)
|
||||
|
||||
flatten_sample_args = []
|
||||
|
||||
for c, args in zip(callables, _sample_args):
|
||||
if isinstance(c, torch.nn.Module):
|
||||
if not (
|
||||
len(c._backward_hooks) == 0
|
||||
and len(c._forward_hooks) == 0
|
||||
and len(c._forward_pre_hooks) == 0
|
||||
):
|
||||
raise AssertionError(
|
||||
"Modules must not have hooks registered at the time they are passed. However, registering hooks "
|
||||
+ "on modules after passing them through make_graphed_callables is allowed."
|
||||
)
|
||||
if not all(b.requires_grad is False for b in c.buffers()):
|
||||
raise AssertionError(
|
||||
"In any :class:`~torch.nn.Module` passed to "
|
||||
+ ":func:`~make_graphed_callables`, only parameters may be trainable. All buffers must have "
|
||||
+ "``requires_grad=False``."
|
||||
)
|
||||
flatten_arg = torch.utils._pytree.arg_tree_leaves(*args)
|
||||
flatten_sample_args.append(tuple(flatten_arg))
|
||||
if not all(isinstance(arg, torch.Tensor) for arg in flatten_arg):
|
||||
raise AssertionError(
|
||||
"In the beta API, sample_args "
|
||||
+ "for each callable must contain only Tensors. Other types are not allowed."
|
||||
)
|
||||
|
||||
# If a callable is an nn.Module, its graph's full input surface is the args the user explicitly
|
||||
# passes to forward (ie, its sample_args) AND the module's parameter attributes.
|
||||
per_callable_len_user_args = [len(args) for args in flatten_sample_args]
|
||||
per_callable_module_params = [
|
||||
tuple(c.parameters()) if isinstance(c, torch.nn.Module) else ()
|
||||
for c in callables
|
||||
]
|
||||
per_callable_static_input_surfaces = [
|
||||
flatten_sample_args[i] + per_callable_module_params[i]
|
||||
for i in range(len(callables))
|
||||
]
|
||||
|
||||
fwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(callables))]
|
||||
bwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(callables))]
|
||||
|
||||
mempool = graph_pool_handle() if pool is None else pool
|
||||
|
||||
# Warmup
|
||||
# Hopefully prevents cudnn benchmarking and other lazy-initialization cuda work
|
||||
# from ending up in any captures.
|
||||
torch.cuda.synchronize()
|
||||
with torch.cuda.stream(torch.cuda.Stream()):
|
||||
for func, args, static_input_surface in zip(
|
||||
callables, _sample_args, per_callable_static_input_surfaces
|
||||
):
|
||||
grad_inputs, outputs, outputs_grad = None, None, None
|
||||
for _ in range(num_warmup_iters):
|
||||
outputs = torch.utils._pytree.tree_leaves(func(*args))
|
||||
outputs_grad = tuple(o for o in outputs if o.requires_grad)
|
||||
if len(outputs_grad) > 0:
|
||||
grad_inputs = torch.autograd.grad(
|
||||
outputs=outputs_grad,
|
||||
inputs=tuple(
|
||||
i for i in static_input_surface if i.requires_grad
|
||||
),
|
||||
grad_outputs=tuple(
|
||||
torch.empty_like(o) for o in outputs if o.requires_grad
|
||||
),
|
||||
only_inputs=True,
|
||||
allow_unused=allow_unused_input,
|
||||
)
|
||||
for v in [outputs, outputs_grad, grad_inputs]:
|
||||
del v
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# All captures here share a mempool. To avoid replays corrupting each other's memory,
|
||||
# the safest approach is to capture all passes in the same order they'll run:
|
||||
# fwd 1, fwd 2, ... fwd N, then bwd N, bwd N-1, ... bwd 1.
|
||||
|
||||
# Capture forward graphs
|
||||
per_callable_static_outputs = []
|
||||
per_callable_output_unflatten_spec = []
|
||||
for func, args, fwd_graph in zip(callables, _sample_args, fwd_graphs):
|
||||
with torch.cuda.graph(fwd_graph, pool=mempool):
|
||||
func_outputs = func(*args)
|
||||
|
||||
flatten_outputs, spec = torch.utils._pytree.tree_flatten(func_outputs)
|
||||
per_callable_static_outputs.append(tuple(flatten_outputs))
|
||||
per_callable_output_unflatten_spec.append(spec)
|
||||
|
||||
# Capture backward graphs in reverse order
|
||||
per_callable_static_grad_outputs = []
|
||||
per_callable_static_grad_inputs = []
|
||||
for static_input_surface, static_outputs, bwd_graph in zip(
|
||||
reversed(per_callable_static_input_surfaces),
|
||||
reversed(per_callable_static_outputs),
|
||||
reversed(bwd_graphs),
|
||||
):
|
||||
# For now, assumes all static_outputs require grad
|
||||
# assert all(o.requires_grad for o in static_outputs), "Outputs of graphed callables must require grad."
|
||||
static_grad_outputs = tuple(
|
||||
torch.empty_like(o) if o.requires_grad else None for o in static_outputs
|
||||
)
|
||||
|
||||
outputs_grad = tuple(o for o in static_outputs if o.requires_grad)
|
||||
grad_inputs = None
|
||||
if len(outputs_grad) > 0:
|
||||
with torch.cuda.graph(bwd_graph, pool=mempool):
|
||||
grad_inputs = torch.autograd.grad(
|
||||
outputs=outputs_grad,
|
||||
inputs=tuple(i for i in static_input_surface if i.requires_grad),
|
||||
grad_outputs=tuple(o for o in static_grad_outputs if o is not None),
|
||||
only_inputs=True,
|
||||
allow_unused=allow_unused_input,
|
||||
)
|
||||
|
||||
# Constructs a tuple suitable for returning from Graphed.backward:
|
||||
# Pads out the actually-needed grads with Nones in gradient slots for inputs that don't require grad.
|
||||
# I couldn't think of a slick one-liner for this pattern.
|
||||
static_grad_inputs = []
|
||||
grad_idx = 0
|
||||
for arg in static_input_surface:
|
||||
if arg.requires_grad and grad_inputs is not None:
|
||||
static_grad_inputs.append(grad_inputs[grad_idx])
|
||||
grad_idx += 1
|
||||
else:
|
||||
static_grad_inputs.append(None) # type: ignore[arg-type]
|
||||
static_grad_inputs = tuple(static_grad_inputs) # type: ignore[assignment]
|
||||
|
||||
per_callable_static_grad_outputs.append(static_grad_outputs)
|
||||
per_callable_static_grad_inputs.append(static_grad_inputs)
|
||||
|
||||
# Reverses the most recent two lists
|
||||
per_callable_static_grad_outputs.reverse()
|
||||
per_callable_static_grad_inputs.reverse()
|
||||
# Now for every per_callable list, per_callable_*[i] holds the stuff for the ith callable.
|
||||
|
||||
def make_graphed_autograd_function(
|
||||
fwd_graph: CUDAGraph,
|
||||
bwd_graph: CUDAGraph,
|
||||
module_params: tuple[torch.nn.Parameter, ...],
|
||||
len_user_args: int,
|
||||
output_unflatten_spec: torch.utils._pytree.TreeSpec,
|
||||
static_input_surface: tuple[Tensor, ...],
|
||||
static_outputs: tuple[Tensor, ...],
|
||||
static_grad_outputs: tuple[Tensor | None, ...],
|
||||
static_grad_inputs: tuple[Tensor, ...],
|
||||
) -> Callable[..., object]:
|
||||
class Graphed(torch.autograd.Function):
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def forward(ctx: object, *inputs: Tensor) -> tuple[Tensor, ...]:
|
||||
# At this stage, only the user args may (potentially) be new tensors.
|
||||
for i in range(len_user_args):
|
||||
if static_input_surface[i].data_ptr() != inputs[i].data_ptr():
|
||||
static_input_surface[i].copy_(inputs[i])
|
||||
fwd_graph.replay()
|
||||
if not isinstance(static_outputs, tuple):
|
||||
raise AssertionError(
|
||||
f"static_outputs must be tuple, got {type(static_outputs)}"
|
||||
)
|
||||
return tuple(o.detach() for o in static_outputs)
|
||||
|
||||
@staticmethod
|
||||
@torch.autograd.function.once_differentiable
|
||||
# pyrefly: ignore [bad-override]
|
||||
def backward(ctx: object, *grads: Tensor) -> tuple[Tensor, ...]:
|
||||
if len(grads) != len(static_grad_outputs):
|
||||
raise AssertionError(
|
||||
f"len(grads)={len(grads)} != len(static_grad_outputs)={len(static_grad_outputs)}"
|
||||
)
|
||||
for g, grad in zip(static_grad_outputs, grads):
|
||||
if g is not None:
|
||||
# don't copy if autograd gods have been kind and the
|
||||
# incoming grad is already in the right place
|
||||
if g.data_ptr() != grad.data_ptr():
|
||||
g.copy_(grad)
|
||||
bwd_graph.replay()
|
||||
|
||||
# Input args that didn't require grad expect a None gradient.
|
||||
if not isinstance(static_grad_inputs, tuple):
|
||||
raise AssertionError(
|
||||
f"static_grad_inputs must be tuple, got {type(static_grad_inputs)}"
|
||||
)
|
||||
return tuple(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
b.detach() if b is not None else b
|
||||
for b in static_grad_inputs
|
||||
)
|
||||
|
||||
def functionalized(*user_args: object) -> object:
|
||||
# Runs the autograd function with inputs == all inputs to the graph that might require grad
|
||||
# (explicit user args + module parameters)
|
||||
# Assumes module params didn't change since capture.
|
||||
flatten_user_args = torch.utils._pytree.arg_tree_leaves(*user_args)
|
||||
out = Graphed.apply(*(tuple(flatten_user_args) + module_params))
|
||||
return torch.utils._pytree.tree_unflatten(out, output_unflatten_spec)
|
||||
|
||||
return functionalized
|
||||
|
||||
# Put together the final graphed callables
|
||||
ret: list[_ModuleOrCallable] = []
|
||||
for i, func in enumerate(callables):
|
||||
graphed = make_graphed_autograd_function(
|
||||
fwd_graphs[i],
|
||||
bwd_graphs[i],
|
||||
per_callable_module_params[i],
|
||||
per_callable_len_user_args[i],
|
||||
per_callable_output_unflatten_spec[i],
|
||||
per_callable_static_input_surfaces[i],
|
||||
per_callable_static_outputs[i],
|
||||
per_callable_static_grad_outputs[i],
|
||||
per_callable_static_grad_inputs[i],
|
||||
)
|
||||
|
||||
if isinstance(func, torch.nn.Module):
|
||||
|
||||
def make_graphed_forward(
|
||||
func: torch.nn.Module,
|
||||
graph_training_state: bool,
|
||||
graphed: Callable[_P, _R],
|
||||
orig_fwd: Callable[_P, _R],
|
||||
) -> Callable[_P, _R]:
|
||||
def new_fwd(*user_args: _P.args, **user_kwargs: _P.kwargs) -> _R:
|
||||
# If the module's training-or-eval state matches what we graphed,
|
||||
# run the graph, otherwise run the original forward method
|
||||
if func.training == graph_training_state:
|
||||
return graphed(*user_args, **user_kwargs)
|
||||
else:
|
||||
return orig_fwd(*user_args, **user_kwargs)
|
||||
|
||||
return new_fwd
|
||||
|
||||
func.forward = make_graphed_forward(
|
||||
func, func.training, graphed, func.forward
|
||||
)
|
||||
ret.append(func)
|
||||
else:
|
||||
ret.append(graphed)
|
||||
|
||||
if just_one_callable:
|
||||
return ret[0]
|
||||
|
||||
return tuple(ret)
|
||||
@@ -0,0 +1,92 @@
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GreenContext",
|
||||
]
|
||||
|
||||
_GreenContext = object
|
||||
SUPPORTED = False
|
||||
|
||||
if hasattr(torch._C, "_CUDAGreenContext"):
|
||||
_GreenContext = torch._C._CUDAGreenContext # type: ignore[misc]
|
||||
SUPPORTED = True
|
||||
|
||||
|
||||
# Python shim helps Sphinx process docstrings more reliably.
|
||||
# pyrefly: ignore [invalid-inheritance]
|
||||
class GreenContext(_GreenContext):
|
||||
r"""Wrapper around a CUDA green context.
|
||||
|
||||
.. warning::
|
||||
This API is in beta and may change in future releases.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
*,
|
||||
num_sms: int | None = None,
|
||||
workqueue_scope: str | None = None,
|
||||
workqueue_concurrency_limit: int | None = None,
|
||||
device_id: int | None = None,
|
||||
) -> _GreenContext:
|
||||
r"""Create a CUDA green context.
|
||||
|
||||
At least one of ``num_sms`` or ``workqueue_scope`` must be specified.
|
||||
Both can be combined to partition SMs and configure workqueues in the
|
||||
same green context.
|
||||
|
||||
Arguments:
|
||||
num_sms (int, optional): The number of SMs to use in the green
|
||||
context. When ``None``, SMs are not partitioned.
|
||||
workqueue_scope (str, optional): Workqueue sharing scope. One of
|
||||
``"device_ctx"`` (shared across all contexts, default driver
|
||||
behaviour) or ``"balanced"`` (non-overlapping workqueues with
|
||||
other balanced green contexts). When ``None``, no workqueue
|
||||
configuration is applied.
|
||||
workqueue_concurrency_limit (int, optional): Maximum number of
|
||||
concurrent stream-ordered workloads for the workqueue. Requires
|
||||
``workqueue_scope`` to be set.
|
||||
device_id (int, optional): The device index of green context.
|
||||
When ``None``, the current device is used.
|
||||
"""
|
||||
if not SUPPORTED:
|
||||
raise RuntimeError("PyTorch was not built with Green Context support!")
|
||||
return _GreenContext.create( # type: ignore[attr-defined]
|
||||
device_id=device_id,
|
||||
num_sms=num_sms,
|
||||
workqueue_scope=workqueue_scope,
|
||||
workqueue_concurrency_limit=workqueue_concurrency_limit,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def max_workqueue_concurrency(device_id: int | None = None) -> int:
|
||||
r"""Return the maximum workqueue concurrency limit for the device.
|
||||
|
||||
This queries the device for the default number of concurrent
|
||||
stream-ordered workloads supported by workqueue configuration
|
||||
resources.
|
||||
|
||||
Arguments:
|
||||
device_id (int, optional): The device index to query. When
|
||||
``None``, the current device is used.
|
||||
"""
|
||||
if not SUPPORTED:
|
||||
raise RuntimeError("PyTorch was not built with Green Context support!")
|
||||
return _GreenContext.max_workqueue_concurrency(device_id=device_id) # type: ignore[attr-defined]
|
||||
|
||||
# Note that these functions are bypassed but we define them here
|
||||
# for Sphinx documentation purposes
|
||||
def set_context(self) -> None: # pylint: disable=useless-parent-delegation
|
||||
r"""Make the green context the current context."""
|
||||
return super().set_context() # type: ignore[misc]
|
||||
|
||||
def pop_context(self) -> None: # pylint: disable=useless-parent-delegation
|
||||
r"""Assuming the green context is the current context, pop it from the
|
||||
context stack and restore the previous context.
|
||||
"""
|
||||
return super().pop_context() # type: ignore[misc]
|
||||
|
||||
def Stream(self) -> "torch.cuda.Stream":
|
||||
r"""Return the CUDA Stream used by the green context."""
|
||||
return super().Stream()
|
||||
@@ -0,0 +1,192 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
__all__: list[str] = []
|
||||
|
||||
|
||||
class _CodeParser:
|
||||
def __init__(self, code_string: str):
|
||||
optional_ws = r"\s*"
|
||||
required_ws = r"\s+"
|
||||
template_params = r"(?P<template_params>\<.+\>)"
|
||||
return_type = r"(?P<return_type>\w+)"
|
||||
function_name = r"(?P<function_name>\w+)"
|
||||
function_params = r"(?P<function_params>\(.+\))"
|
||||
function_body = r"(?P<function_body>\{.+\})"
|
||||
|
||||
pattern = (
|
||||
optional_ws
|
||||
+ "template"
|
||||
+ optional_ws
|
||||
+ template_params
|
||||
+ optional_ws
|
||||
+ return_type
|
||||
+ required_ws
|
||||
+ function_name
|
||||
+ optional_ws
|
||||
+ function_params
|
||||
+ optional_ws
|
||||
+ function_body
|
||||
+ optional_ws
|
||||
)
|
||||
|
||||
result = re.match(
|
||||
pattern, code_string, re.DOTALL
|
||||
) # DOTALL for matching multiline
|
||||
|
||||
if result is None:
|
||||
raise Exception( # noqa: TRY002
|
||||
f"Couldn't parse code, please check correctness:\n {code_string}"
|
||||
)
|
||||
|
||||
self.template_params = result["template_params"]
|
||||
self.return_type = result["return_type"]
|
||||
self.function_name = result["function_name"]
|
||||
self.function_params = result["function_params"]
|
||||
self.function_body = result["function_body"]
|
||||
|
||||
|
||||
class _JittedFunction:
|
||||
def __init__(
|
||||
self, code_string: str, return_by_ref: bool, num_outputs: int, **kwargs
|
||||
):
|
||||
self.code_string = code_string
|
||||
|
||||
if not (return_by_ref or num_outputs == 1):
|
||||
raise AssertionError("Return by value only works for single output.")
|
||||
self.return_by_ref = return_by_ref
|
||||
self.num_outputs = num_outputs
|
||||
|
||||
parsed_code = _CodeParser(code_string)
|
||||
self.kernel_name = parsed_code.function_name
|
||||
|
||||
self.kwargs_dict = kwargs
|
||||
self.is_cuda_available = torch.cuda.is_available()
|
||||
|
||||
def __call__(self, *tensors: Tensor, **kwargs):
|
||||
# Jiterator follow torch.cuda's lazy initialization behavior
|
||||
# Defer checking cuda's availability at the function invocation time
|
||||
if not self.is_cuda_available:
|
||||
raise AssertionError(
|
||||
"Jiterator is only supported on CUDA and ROCm GPUs, none are available."
|
||||
)
|
||||
|
||||
if len(tensors) > 8:
|
||||
raise AssertionError(
|
||||
f"jiterator only supports up to 8 tensor inputs, got {len(tensors)}"
|
||||
)
|
||||
|
||||
expanded_kwargs = self.kwargs_dict.copy()
|
||||
for key, value in kwargs.items():
|
||||
if key in self.kwargs_dict:
|
||||
expanded_kwargs[key] = value
|
||||
else:
|
||||
raise KeyError(f"{key} is not declared in function definition")
|
||||
|
||||
return torch._C._cuda_jiterator_compile_and_launch_kernel(
|
||||
self.code_string,
|
||||
self.kernel_name,
|
||||
self.return_by_ref,
|
||||
self.num_outputs,
|
||||
tensors,
|
||||
expanded_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _create_jit_fn(code_string: str, **kwargs) -> Callable:
|
||||
"""
|
||||
Create a jiterator-generated cuda kernel for an elementwise op.
|
||||
|
||||
The code string has to be a valid CUDA function that describes the computation for a single element. The code
|
||||
string has to follow the c++ template pattern, as shown in the example below. This function will be inlined
|
||||
into elementwise kernel template, and compiled on the fly. Compiled kernel will be cached in memory, as well as
|
||||
local temp dir.
|
||||
|
||||
Jiterator-generated kernels accepts noncontiguous tensors, and supports broadcasting and type promotion.
|
||||
|
||||
Args:
|
||||
code_string (str): CUDA code string to be compiled by jiterator. The entry functor must return by value.
|
||||
kwargs (Dict, optional): Keyword arguments for generated function
|
||||
|
||||
Example::
|
||||
|
||||
code_string = "template <typename T> T my_kernel(T x, T y, T alpha) { return -x + alpha * y; }"
|
||||
jitted_fn = create_jit_fn(code_string, alpha=1.0)
|
||||
a = torch.rand(3, device="cuda")
|
||||
b = torch.rand(3, device="cuda")
|
||||
# invoke jitted function like a regular python function
|
||||
result = jitted_fn(a, b, alpha=3.14)
|
||||
|
||||
code_string also allows multiple function definitions, and the last function will be treated as the entry function.
|
||||
|
||||
Example::
|
||||
|
||||
code_string = (
|
||||
"template <typename T> T util_fn(T x, T y) { return ::sin(x) + ::cos(y); }"
|
||||
)
|
||||
code_string += "template <typename T> T my_kernel(T x, T y, T val) { return ::min(val, util_fn(x, y)); }"
|
||||
jitted_fn = create_jit_fn(code_string, val=0.0)
|
||||
a = torch.rand(3, device="cuda")
|
||||
b = torch.rand(3, device="cuda")
|
||||
# invoke jitted function like a regular python function
|
||||
result = jitted_fn(a, b) # using default val=0.0
|
||||
|
||||
Jiterator can be used together with python registration to override an operator's cuda kernel.
|
||||
Following example is overriding gelu's cuda kernel with relu.
|
||||
|
||||
Example::
|
||||
|
||||
code_string = "template <typename T> T my_gelu(T a) { return a > 0 ? a : 0; }"
|
||||
my_gelu = create_jit_fn(code_string)
|
||||
my_lib = torch.library.Library("aten", "IMPL")
|
||||
my_lib.impl("aten::gelu", my_gelu, "CUDA")
|
||||
# torch.nn.GELU and torch.nn.function.gelu are now overridden
|
||||
a = torch.rand(3, device="cuda")
|
||||
torch.allclose(torch.nn.functional.gelu(a), torch.nn.functional.relu(a))
|
||||
|
||||
.. warning::
|
||||
This API is in beta and may change in future releases.
|
||||
|
||||
.. warning::
|
||||
This API only supports up to 8 inputs and 1 output
|
||||
|
||||
.. warning::
|
||||
All input tensors must live in CUDA device
|
||||
"""
|
||||
return _JittedFunction(code_string, return_by_ref=False, num_outputs=1, **kwargs)
|
||||
|
||||
|
||||
def _create_multi_output_jit_fn(
|
||||
code_string: str, num_outputs: int, **kwargs
|
||||
) -> Callable:
|
||||
"""
|
||||
Create a jiterator-generated cuda kernel for an elementwise op that supports returning one or more outputs.
|
||||
|
||||
Args:
|
||||
code_string (str): CUDA code string to be compiled by jiterator. The entry functor must return value by reference.
|
||||
num_outputs(int): number of outputs return by the kernel
|
||||
kwargs (Dict, optional): Keyword arguments for generated function
|
||||
|
||||
Example::
|
||||
|
||||
code_string = "template <typename T> void my_kernel(T x, T y, T alpha, T& out) { out = -x + alpha * y; }"
|
||||
jitted_fn = create_jit_fn(code_string, alpha=1.0)
|
||||
a = torch.rand(3, device="cuda")
|
||||
b = torch.rand(3, device="cuda")
|
||||
# invoke jitted function like a regular python function
|
||||
result = jitted_fn(a, b, alpha=3.14)
|
||||
|
||||
.. warning::
|
||||
This API is in beta and may change in future releases.
|
||||
|
||||
.. warning::
|
||||
This API only supports up to 8 inputs and 8 outputs
|
||||
"""
|
||||
return _JittedFunction(
|
||||
code_string, return_by_ref=True, num_outputs=num_outputs, **kwargs
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,151 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import collections
|
||||
import warnings
|
||||
from collections.abc import Sequence
|
||||
|
||||
import torch.cuda
|
||||
|
||||
|
||||
__all__ = ["all_reduce", "reduce", "broadcast", "all_gather", "reduce_scatter"]
|
||||
|
||||
SUM = 0 # ncclRedOp_t
|
||||
|
||||
|
||||
def is_available(tensors):
|
||||
if not hasattr(torch._C, "_nccl_all_reduce"):
|
||||
warnings.warn("PyTorch is not compiled with NCCL support", stacklevel=2)
|
||||
return False
|
||||
|
||||
devices = set()
|
||||
for tensor in tensors:
|
||||
if tensor.is_sparse:
|
||||
return False
|
||||
if not tensor.is_contiguous():
|
||||
return False
|
||||
if not tensor.is_cuda:
|
||||
return False
|
||||
device = tensor.get_device()
|
||||
if device in devices:
|
||||
return False
|
||||
devices.add(device)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def version():
|
||||
"""
|
||||
Returns the version of the NCCL.
|
||||
|
||||
|
||||
This function returns a tuple containing the major, minor, and patch version numbers of the NCCL.
|
||||
The suffix is also included in the tuple if a version suffix exists.
|
||||
Returns:
|
||||
tuple: The version information of the NCCL.
|
||||
"""
|
||||
ver = torch._C._nccl_version()
|
||||
major = ver >> 32
|
||||
minor = (ver >> 16) & 65535
|
||||
patch = ver & 65535
|
||||
suffix = torch._C._nccl_version_suffix().decode("utf-8")
|
||||
if suffix == "":
|
||||
return (major, minor, patch)
|
||||
else:
|
||||
return (major, minor, patch, suffix)
|
||||
|
||||
|
||||
def unique_id():
|
||||
return torch._C._nccl_unique_id()
|
||||
|
||||
|
||||
def init_rank(num_ranks, uid, rank):
|
||||
return torch._C._nccl_init_rank(num_ranks, uid, rank)
|
||||
|
||||
|
||||
def _check_sequence_type(inputs: torch.Tensor | Sequence[torch.Tensor]) -> None:
|
||||
if not isinstance(inputs, collections.abc.Container) or isinstance(
|
||||
inputs, torch.Tensor
|
||||
):
|
||||
raise TypeError("Inputs should be a collection of tensors")
|
||||
|
||||
|
||||
def all_reduce(inputs, outputs=None, op=SUM, streams=None, comms=None):
|
||||
_check_sequence_type(inputs)
|
||||
if outputs is None:
|
||||
outputs = inputs
|
||||
_check_sequence_type(outputs)
|
||||
torch._C._nccl_all_reduce(inputs, outputs, op, streams, comms)
|
||||
|
||||
|
||||
# `output` used to be `outputs`, taking in a list of tensors. So we have two
|
||||
# arguments for BC reasons.
|
||||
def reduce(
|
||||
inputs: Sequence[torch.Tensor],
|
||||
output: torch.Tensor | Sequence[torch.Tensor] | None = None,
|
||||
root: int = 0,
|
||||
op: int = SUM,
|
||||
streams: Sequence[torch.cuda.Stream] | None = None,
|
||||
comms=None,
|
||||
*,
|
||||
outputs: Sequence[torch.Tensor] | None = None,
|
||||
) -> None:
|
||||
_check_sequence_type(inputs)
|
||||
_output: torch.Tensor
|
||||
if outputs is not None:
|
||||
if output is not None:
|
||||
raise ValueError(
|
||||
"'output' and 'outputs' can not be both specified. 'outputs' is deprecated in "
|
||||
"favor of 'output', taking in a single output tensor. The signature of reduce is: "
|
||||
"reduce(inputs, output=None, root=0, op=SUM, streams=None, comms=None)."
|
||||
)
|
||||
else:
|
||||
warnings.warn(
|
||||
"`nccl.reduce` with an output tensor list is deprecated. "
|
||||
"Please specify a single output tensor with argument 'output' instead instead.",
|
||||
FutureWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
_output = outputs[root]
|
||||
elif not isinstance(output, torch.Tensor) and isinstance(
|
||||
output, collections.abc.Sequence
|
||||
):
|
||||
# User called old API with positional arguments of list of output tensors.
|
||||
warnings.warn(
|
||||
"nccl.reduce with an output tensor list is deprecated. "
|
||||
"Please specify a single output tensor.",
|
||||
FutureWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
_output = output[root]
|
||||
else:
|
||||
_output = inputs[root] if output is None else output
|
||||
torch._C._nccl_reduce(inputs, _output, root, op, streams, comms)
|
||||
|
||||
|
||||
def broadcast(
|
||||
inputs: Sequence[torch.Tensor], root: int = 0, streams=None, comms=None
|
||||
) -> None:
|
||||
_check_sequence_type(inputs)
|
||||
torch._C._nccl_broadcast(inputs, root, streams, comms)
|
||||
|
||||
|
||||
def all_gather(
|
||||
inputs: Sequence[torch.Tensor],
|
||||
outputs: Sequence[torch.Tensor],
|
||||
streams=None,
|
||||
comms=None,
|
||||
) -> None:
|
||||
_check_sequence_type(inputs)
|
||||
_check_sequence_type(outputs)
|
||||
torch._C._nccl_all_gather(inputs, outputs, streams, comms)
|
||||
|
||||
|
||||
def reduce_scatter(
|
||||
inputs: Sequence[torch.Tensor],
|
||||
outputs: Sequence[torch.Tensor],
|
||||
op: int = SUM,
|
||||
streams=None,
|
||||
comms=None,
|
||||
) -> None:
|
||||
_check_sequence_type(inputs)
|
||||
_check_sequence_type(outputs)
|
||||
torch._C._nccl_reduce_scatter(inputs, outputs, op, streams, comms)
|
||||
@@ -0,0 +1,129 @@
|
||||
# mypy: allow-untyped-defs
|
||||
r"""This package adds support for NVIDIA Tools Extension (NVTX) used in profiling."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
|
||||
|
||||
try:
|
||||
from torch._C import _nvtx
|
||||
except ImportError:
|
||||
|
||||
class _NVTXStub:
|
||||
@staticmethod
|
||||
def _fail(*args, **kwargs):
|
||||
raise RuntimeError(
|
||||
"NVTX functions not installed. Are you sure you have a CUDA build?"
|
||||
)
|
||||
|
||||
rangePushA = _fail
|
||||
rangePop = _fail
|
||||
markA = _fail
|
||||
|
||||
_nvtx = _NVTXStub() # type: ignore[assignment]
|
||||
|
||||
__all__ = ["range_push", "range_pop", "range_start", "range_end", "mark", "range"]
|
||||
|
||||
|
||||
def range_push(msg):
|
||||
"""
|
||||
Push a range onto a stack of nested range span. Returns zero-based depth of the range that is started.
|
||||
|
||||
Args:
|
||||
msg (str): ASCII message to associate with range
|
||||
"""
|
||||
return _nvtx.rangePushA(msg)
|
||||
|
||||
|
||||
def range_pop():
|
||||
"""Pop a range off of a stack of nested range spans. Returns the zero-based depth of the range that is ended."""
|
||||
return _nvtx.rangePop()
|
||||
|
||||
|
||||
def range_start(msg) -> int:
|
||||
"""
|
||||
Mark the start of a range with string message. It returns an unique handle
|
||||
for this range to pass to the corresponding call to rangeEnd().
|
||||
|
||||
A key difference between this and range_push/range_pop is that the
|
||||
range_start/range_end version supports range across threads (start on one
|
||||
thread and end on another thread).
|
||||
|
||||
Returns: A range handle (uint64_t) that can be passed to range_end().
|
||||
|
||||
Args:
|
||||
msg (str): ASCII message to associate with the range.
|
||||
"""
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return _nvtx.rangeStartA(msg)
|
||||
|
||||
|
||||
def range_end(range_id) -> None:
|
||||
"""
|
||||
Mark the end of a range for a given range_id.
|
||||
|
||||
Args:
|
||||
range_id (int): an unique handle for the start range.
|
||||
"""
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
_nvtx.rangeEnd(range_id)
|
||||
|
||||
|
||||
def _device_range_start(msg: str, stream: int = 0) -> object:
|
||||
"""
|
||||
Marks the start of a range with string message.
|
||||
It returns an opaque heap-allocated handle for this range
|
||||
to pass to the corresponding call to device_range_end().
|
||||
|
||||
A key difference between this and range_start is that the
|
||||
range_start marks the range right away, while _device_range_start
|
||||
marks the start of the range as soon as all the tasks on the
|
||||
CUDA stream are completed.
|
||||
|
||||
Returns: An opaque heap-allocated handle that should be passed to _device_range_end().
|
||||
|
||||
Args:
|
||||
msg (str): ASCII message to associate with the range.
|
||||
stream (int): CUDA stream id.
|
||||
"""
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return _nvtx.deviceRangeStart(msg, stream)
|
||||
|
||||
|
||||
def _device_range_end(range_handle: object, stream: int = 0) -> None:
|
||||
"""
|
||||
Mark the end of a range for a given range_handle as soon as all the tasks
|
||||
on the CUDA stream are completed.
|
||||
|
||||
Args:
|
||||
range_handle: an unique handle for the start range.
|
||||
stream (int): CUDA stream id.
|
||||
"""
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
_nvtx.deviceRangeEnd(range_handle, stream)
|
||||
|
||||
|
||||
def mark(msg):
|
||||
"""
|
||||
Describe an instantaneous event that occurred at some point.
|
||||
|
||||
Args:
|
||||
msg (str): ASCII message to associate with the event.
|
||||
"""
|
||||
return _nvtx.markA(msg)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def range(msg, *args, **kwargs):
|
||||
"""
|
||||
Context manager / decorator that pushes an NVTX range at the beginning
|
||||
of its scope, and pops it at the end. If extra arguments are given,
|
||||
they are passed as arguments to msg.format().
|
||||
|
||||
Args:
|
||||
msg (str): message to associate with the range
|
||||
"""
|
||||
range_push(msg.format(*args, **kwargs))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
range_pop()
|
||||
@@ -0,0 +1,56 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import contextlib
|
||||
|
||||
from . import check_error, cudart
|
||||
|
||||
|
||||
__all__ = ["start", "stop", "profile"]
|
||||
|
||||
DEFAULT_FLAGS = [
|
||||
"gpustarttimestamp",
|
||||
"gpuendtimestamp",
|
||||
"gridsize3d",
|
||||
"threadblocksize",
|
||||
"streamid",
|
||||
"enableonstart 0",
|
||||
"conckerneltrace",
|
||||
]
|
||||
|
||||
|
||||
def start():
|
||||
r"""Starts cuda profiler data collection.
|
||||
|
||||
.. warning::
|
||||
Raises CudaError in case of it is unable to start the profiler.
|
||||
"""
|
||||
check_error(cudart().cudaProfilerStart())
|
||||
|
||||
|
||||
def stop():
|
||||
r"""Stops cuda profiler data collection.
|
||||
|
||||
.. warning::
|
||||
Raises CudaError in case of it is unable to stop the profiler.
|
||||
"""
|
||||
check_error(cudart().cudaProfilerStop())
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def profile():
|
||||
"""
|
||||
Enable profiling.
|
||||
|
||||
Context Manager to enabling profile collection by the active profiling tool from CUDA backend.
|
||||
Example:
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
|
||||
>>> import torch
|
||||
>>> model = torch.nn.Linear(20, 30).cuda()
|
||||
>>> inputs = torch.randn(128, 20).cuda()
|
||||
>>> with torch.cuda.profiler.profile() as prof:
|
||||
... model(inputs)
|
||||
"""
|
||||
try:
|
||||
start()
|
||||
yield
|
||||
finally:
|
||||
stop()
|
||||
@@ -0,0 +1,181 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from . import _lazy_call, _lazy_init, current_device, device_count, is_initialized
|
||||
|
||||
|
||||
__all__ = [
|
||||
"get_rng_state",
|
||||
"get_rng_state_all",
|
||||
"set_rng_state",
|
||||
"set_rng_state_all",
|
||||
"manual_seed",
|
||||
"manual_seed_all",
|
||||
"seed",
|
||||
"seed_all",
|
||||
"initial_seed",
|
||||
]
|
||||
|
||||
|
||||
def get_rng_state(device: int | str | torch.device = "cuda") -> Tensor:
|
||||
r"""Return the random number generator state of the specified GPU as a ByteTensor.
|
||||
|
||||
Args:
|
||||
device (torch.device or int, optional): The device to return the RNG state of.
|
||||
Default: ``'cuda'`` (i.e., ``torch.device('cuda')``, the current CUDA device).
|
||||
|
||||
.. warning::
|
||||
This function eagerly initializes CUDA.
|
||||
"""
|
||||
_lazy_init()
|
||||
if isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
elif isinstance(device, int):
|
||||
device = torch.device("cuda", device)
|
||||
idx = device.index
|
||||
if idx is None:
|
||||
idx = current_device()
|
||||
default_generator = torch.cuda.default_generators[idx]
|
||||
return default_generator.get_state()
|
||||
|
||||
|
||||
def get_rng_state_all() -> list[Tensor]:
|
||||
r"""Return a list of ByteTensor representing the random number states of all devices."""
|
||||
results = [get_rng_state(i) for i in range(device_count())]
|
||||
return results
|
||||
|
||||
|
||||
def set_rng_state(new_state: Tensor, device: int | str | torch.device = "cuda") -> None:
|
||||
r"""Set the random number generator state of the specified GPU.
|
||||
|
||||
Args:
|
||||
new_state (torch.ByteTensor): The desired state
|
||||
device (torch.device or int, optional): The device to set the RNG state.
|
||||
Default: ``'cuda'`` (i.e., ``torch.device('cuda')``, the current CUDA device).
|
||||
"""
|
||||
if not is_initialized():
|
||||
with torch._C._DisableFuncTorch():
|
||||
# Clone the state because the callback will be triggered
|
||||
# later when CUDA is lazy initialized.
|
||||
new_state = new_state.clone(memory_format=torch.contiguous_format)
|
||||
if isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
elif isinstance(device, int):
|
||||
device = torch.device("cuda", device)
|
||||
|
||||
def cb():
|
||||
idx = device.index
|
||||
if idx is None:
|
||||
idx = current_device()
|
||||
default_generator = torch.cuda.default_generators[idx]
|
||||
default_generator.set_state(new_state)
|
||||
|
||||
_lazy_call(cb)
|
||||
|
||||
|
||||
def set_rng_state_all(new_states: Iterable[Tensor]) -> None:
|
||||
r"""Set the random number generator state of all devices.
|
||||
|
||||
Args:
|
||||
new_states (Iterable of torch.ByteTensor): The desired state for each device.
|
||||
"""
|
||||
for i, state in enumerate(new_states):
|
||||
set_rng_state(state, i)
|
||||
|
||||
|
||||
def manual_seed(seed: int) -> None:
|
||||
r"""Set the seed for generating random numbers for the current GPU.
|
||||
|
||||
It's safe to call this function if CUDA is not available; in that
|
||||
case, it is silently ignored.
|
||||
|
||||
Args:
|
||||
seed (int): The desired seed.
|
||||
|
||||
.. warning::
|
||||
If you are working with a multi-GPU model, this function is insufficient
|
||||
to get determinism. To seed all GPUs, use :func:`manual_seed_all`.
|
||||
"""
|
||||
seed = int(seed)
|
||||
|
||||
def cb():
|
||||
idx = current_device()
|
||||
default_generator = torch.cuda.default_generators[idx]
|
||||
default_generator.manual_seed(seed)
|
||||
|
||||
_lazy_call(cb, seed=True)
|
||||
|
||||
|
||||
def manual_seed_all(seed: int) -> None:
|
||||
r"""Set the seed for generating random numbers on all GPUs.
|
||||
|
||||
It's safe to call this function if CUDA is not available; in that
|
||||
case, it is silently ignored.
|
||||
|
||||
Args:
|
||||
seed (int): The desired seed.
|
||||
"""
|
||||
seed = int(seed)
|
||||
|
||||
def cb():
|
||||
for i in range(device_count()):
|
||||
default_generator = torch.cuda.default_generators[i]
|
||||
default_generator.manual_seed(seed)
|
||||
|
||||
_lazy_call(cb, seed_all=True)
|
||||
|
||||
|
||||
def seed() -> None:
|
||||
r"""Set the seed for generating random numbers to a random number for the current GPU.
|
||||
|
||||
It's safe to call this function if CUDA is not available; in that
|
||||
case, it is silently ignored.
|
||||
|
||||
.. warning::
|
||||
If you are working with a multi-GPU model, this function will only initialize
|
||||
the seed on one GPU. To initialize all GPUs, use :func:`seed_all`.
|
||||
"""
|
||||
|
||||
def cb():
|
||||
idx = current_device()
|
||||
default_generator = torch.cuda.default_generators[idx]
|
||||
default_generator.seed()
|
||||
|
||||
_lazy_call(cb)
|
||||
|
||||
|
||||
def seed_all() -> None:
|
||||
r"""Set the seed for generating random numbers to a random number on all GPUs.
|
||||
|
||||
It's safe to call this function if CUDA is not available; in that
|
||||
case, it is silently ignored.
|
||||
"""
|
||||
|
||||
def cb():
|
||||
random_seed = 0
|
||||
seeded = False
|
||||
for i in range(device_count()):
|
||||
default_generator = torch.cuda.default_generators[i]
|
||||
if not seeded:
|
||||
default_generator.seed()
|
||||
random_seed = default_generator.initial_seed()
|
||||
seeded = True
|
||||
else:
|
||||
default_generator.manual_seed(random_seed)
|
||||
|
||||
_lazy_call(cb)
|
||||
|
||||
|
||||
def initial_seed() -> int:
|
||||
r"""Return the current random seed of the current GPU.
|
||||
|
||||
.. warning::
|
||||
This function eagerly initializes CUDA.
|
||||
"""
|
||||
_lazy_init()
|
||||
idx = current_device()
|
||||
default_generator = torch.cuda.default_generators[idx]
|
||||
return default_generator.initial_seed()
|
||||
@@ -0,0 +1 @@
|
||||
# The Tensor classes are added to this module by python_tensor.cpp
|
||||
@@ -0,0 +1,271 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# pylint: disable=useless-parent-delegation
|
||||
from __future__ import annotations
|
||||
|
||||
import ctypes
|
||||
|
||||
import torch
|
||||
from torch._utils import _dummy_type
|
||||
|
||||
|
||||
if not hasattr(torch._C, "_CudaStreamBase"):
|
||||
# Define dummy base classes
|
||||
torch._C.__dict__["_CudaStreamBase"] = _dummy_type("_CudaStreamBase")
|
||||
torch._C.__dict__["_CudaEventBase"] = _dummy_type("_CudaEventBase")
|
||||
|
||||
|
||||
class Stream(torch._C._CudaStreamBase):
|
||||
r"""Wrapper around a CUDA stream.
|
||||
|
||||
A CUDA stream is a linear sequence of execution that belongs to a specific
|
||||
device, independent from other streams. It supports with statement as a
|
||||
context manager to ensure the operators within the with block are running
|
||||
on the corresponding stream. See :ref:`cuda-semantics` for details.
|
||||
|
||||
Args:
|
||||
device(torch.device or int, optional): a device on which to allocate
|
||||
the stream. If :attr:`device` is ``None`` (default) or a negative
|
||||
integer, this will use the current device.
|
||||
priority(int, optional): priority of the stream, which can be positive, 0, or negative.
|
||||
A lower number indicates a higher priority. By default, the priority is set to 0.
|
||||
If the value falls outside of the allowed priority range, it will automatically be
|
||||
mapped to the nearest valid priority (lowest for large positive numbers or
|
||||
highest for large negative numbers).
|
||||
|
||||
"""
|
||||
|
||||
def __new__(cls, device=None, priority=0, **kwargs):
|
||||
# Check CUDA availability
|
||||
if not torch.backends.cuda.is_built():
|
||||
raise RuntimeError("torch.cuda.Stream requires CUDA support")
|
||||
# setting device manager is expensive, so we avoid it unless necessary
|
||||
if device is None or ("stream_id" in kwargs and "device_index" in kwargs):
|
||||
return super().__new__(cls, priority=priority, **kwargs)
|
||||
else:
|
||||
with torch.cuda.device(device):
|
||||
return super().__new__(cls, priority=priority, **kwargs)
|
||||
|
||||
def wait_event(self, event: Event | torch.Event) -> None:
|
||||
r"""Make all future work submitted to the stream wait for an event.
|
||||
|
||||
Args:
|
||||
event (Event, torch.Event): an event to wait for.
|
||||
|
||||
.. note:: This is a wrapper around ``cudaStreamWaitEvent()``: see
|
||||
`CUDA Stream documentation`_ for more info.
|
||||
|
||||
This function returns without waiting for :attr:`event`: only future
|
||||
operations are affected.
|
||||
|
||||
.. _CUDA Stream documentation:
|
||||
https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html
|
||||
"""
|
||||
event.wait(self)
|
||||
|
||||
def wait_stream(self, stream: Stream | torch.Stream) -> None:
|
||||
r"""Synchronize with another stream.
|
||||
|
||||
All future work submitted to this stream will wait until all kernels
|
||||
submitted to a given stream at the time of call complete.
|
||||
|
||||
Args:
|
||||
stream (Stream, torch.Stream): a stream to synchronize.
|
||||
|
||||
.. note:: This function returns without waiting for currently enqueued
|
||||
kernels in :attr:`stream`: only future operations are affected.
|
||||
"""
|
||||
self.wait_event(stream.record_event())
|
||||
|
||||
def record_event(self, event: Event | torch.Event | None = None):
|
||||
r"""Record an event.
|
||||
|
||||
Args:
|
||||
event (Event, torch.Event, optional): event to record. If not given, a new one
|
||||
will be allocated.
|
||||
|
||||
Returns:
|
||||
Recorded event.
|
||||
"""
|
||||
if event is None:
|
||||
event = Event()
|
||||
event.record(self)
|
||||
return event
|
||||
|
||||
def query(self) -> bool:
|
||||
r"""Check if all the work submitted has been completed.
|
||||
|
||||
Returns:
|
||||
A boolean indicating if all kernels in this stream are completed.
|
||||
"""
|
||||
return super().query()
|
||||
|
||||
def synchronize(self) -> None:
|
||||
r"""Wait for all the kernels in this stream to complete.
|
||||
|
||||
.. note:: This is a wrapper around ``cudaStreamSynchronize()``: see
|
||||
`CUDA Stream documentation`_ for more info.
|
||||
"""
|
||||
super().synchronize()
|
||||
|
||||
@property
|
||||
def _as_parameter_(self):
|
||||
return ctypes.c_void_p(self.cuda_stream)
|
||||
|
||||
def __eq__(self, o) -> bool:
|
||||
if isinstance(o, Stream):
|
||||
return super().__eq__(o)
|
||||
return False
|
||||
|
||||
def __hash__(self):
|
||||
return hash((self.cuda_stream, self.device))
|
||||
|
||||
def __repr__(self):
|
||||
return f"<torch.cuda.Stream device={self.device} cuda_stream={self.cuda_stream:#x}>"
|
||||
|
||||
def __cuda_stream__(self):
|
||||
"""Implements the CUDA Stream Protocol:
|
||||
https://nvidia.github.io/cuda-python/cuda-core/latest/interoperability.html#cuda-stream-protocol
|
||||
|
||||
Returns:
|
||||
tuple: A 2-tuple of (version, handle) where version is the protocol version
|
||||
and handle is the address of cudaStream_t (CUDA) or hipStream_t (ROCm) as a Python int.
|
||||
"""
|
||||
return (0, self.cuda_stream)
|
||||
|
||||
|
||||
class ExternalStream(Stream):
|
||||
r"""Wrapper around an externally allocated CUDA stream.
|
||||
|
||||
This class is used to wrap streams allocated in other libraries in order
|
||||
to facilitate data exchange and multi-library interactions.
|
||||
|
||||
.. note:: This class doesn't manage the stream life-cycle, it is the user
|
||||
responsibility to keep the referenced stream alive while this class is
|
||||
being used.
|
||||
|
||||
Args:
|
||||
stream_ptr(int): Integer representation of the `cudaStream_t` value.
|
||||
allocated externally.
|
||||
device(torch.device or int, optional): the device where the stream
|
||||
was originally allocated. If device is specified incorrectly,
|
||||
subsequent launches using this stream may fail.
|
||||
"""
|
||||
|
||||
def __new__(cls, stream_ptr, device=None, **kwargs):
|
||||
with torch.cuda.device(device):
|
||||
return super().__new__(cls, stream_ptr=stream_ptr, **kwargs)
|
||||
|
||||
|
||||
class Event(torch._C._CudaEventBase):
|
||||
r"""Wrapper around a CUDA event.
|
||||
|
||||
CUDA events are synchronization markers that can be used to monitor the
|
||||
device's progress, to accurately measure timing, and to synchronize CUDA
|
||||
streams.
|
||||
|
||||
The underlying CUDA events are lazily initialized when the event is first
|
||||
recorded or exported to another process. After creation, only streams on the
|
||||
same device may record the event. However, streams on any device can wait on
|
||||
the event.
|
||||
|
||||
Args:
|
||||
enable_timing (bool, optional): indicates if the event should measure time
|
||||
(default: ``False``)
|
||||
blocking (bool, optional): if ``True``, :meth:`wait` will be blocking (default: ``False``)
|
||||
interprocess (bool): if ``True``, the event can be shared between processes
|
||||
(default: ``False``)
|
||||
external (bool, optional): indicates whether this event should create event record and event wait nodes, or create an internal cross-stream dependency, when captured in a cuda graph. See `cross-stream dependencies <https://docs.nvidia.com/cuda/archive/12.9.0/cuda-c-programming-guide/index.html#cross-stream-dependencies-and-events>`_, `cudaEventRecordExternal <https://docs.nvidia.com/cuda/archive/12.9.0/cuda-runtime-api/group__CUDART__TYPES.html#group__CUDART__TYPES_1g3457b81d1d32c6a00f6132fbc2693d47>`_, and `cudaEventWaitExternal <https://docs.nvidia.com/cuda/archive/12.9.0/cuda-runtime-api/group__CUDART__TYPES.html#group__CUDART__TYPES_1g0c23426b7252eaa9cef695859991304e>`_ for more information about internal vs. external events. (default: ``False``)
|
||||
|
||||
.. _CUDA Event Documentation:
|
||||
https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__EVENT.html
|
||||
""" # noqa: B950
|
||||
|
||||
def __new__(
|
||||
cls, enable_timing=False, blocking=False, interprocess=False, external=False
|
||||
):
|
||||
return super().__new__(
|
||||
cls,
|
||||
enable_timing=enable_timing,
|
||||
blocking=blocking,
|
||||
interprocess=interprocess,
|
||||
external=external,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_ipc_handle(cls, device, handle):
|
||||
r"""Reconstruct an event from an IPC handle on the given device."""
|
||||
return super().from_ipc_handle(device, handle)
|
||||
|
||||
def record(self, stream: Stream | torch.Stream | None = None):
|
||||
r"""Record the event in a given stream.
|
||||
|
||||
Args:
|
||||
stream (Stream, torch.Stream, optional): Uses ``torch.cuda.current_stream()`` if no stream is specified.
|
||||
The stream's device must match the event's device.
|
||||
"""
|
||||
if stream is None:
|
||||
stream = torch.cuda.current_stream()
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
super().record(stream)
|
||||
|
||||
def wait(self, stream: Stream | torch.Stream | None = None) -> None:
|
||||
r"""Make all future work submitted to the given stream wait for this event.
|
||||
|
||||
Args:
|
||||
stream (Stream, torch.Stream, optional): Uses ``torch.cuda.current_stream()`` if no stream is specified.
|
||||
|
||||
.. note:: This is a wrapper around ``cudaStreamWaitEvent()``: see
|
||||
`CUDA Event documentation`_ for more info.
|
||||
"""
|
||||
if stream is None:
|
||||
stream = torch.cuda.current_stream()
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
super().wait(stream)
|
||||
|
||||
def query(self):
|
||||
r"""Check if all work currently captured by event has completed.
|
||||
|
||||
Returns:
|
||||
A boolean indicating if all work currently captured by event has
|
||||
completed.
|
||||
"""
|
||||
return super().query()
|
||||
|
||||
def elapsed_time(self, end_event: Event):
|
||||
r"""Return the time elapsed.
|
||||
|
||||
Time reported in milliseconds after the event was recorded and
|
||||
before the end_event was recorded.
|
||||
|
||||
Args:
|
||||
end_event (Event): the end event.
|
||||
"""
|
||||
return super().elapsed_time(end_event)
|
||||
|
||||
def synchronize(self) -> None:
|
||||
r"""Wait for the event to complete.
|
||||
|
||||
Waits until the completion of all work currently captured in this event.
|
||||
This prevents the CPU thread from proceeding until the event completes.
|
||||
|
||||
.. note:: This is a wrapper around ``cudaEventSynchronize()``: see
|
||||
`CUDA Event documentation`_ for more info.
|
||||
"""
|
||||
super().synchronize()
|
||||
|
||||
def ipc_handle(self):
|
||||
r"""Return an IPC handle of this event.
|
||||
|
||||
If not recorded yet, the event will use the current device.
|
||||
"""
|
||||
return super().ipc_handle()
|
||||
|
||||
@property
|
||||
def _as_parameter_(self):
|
||||
return ctypes.c_void_p(self.cuda_event)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
if self.cuda_event:
|
||||
return f"<torch.cuda.Event {self._as_parameter_.value:#x}>"
|
||||
else:
|
||||
return "<torch.cuda.Event uninitialized>"
|
||||
@@ -0,0 +1,834 @@
|
||||
r"""
|
||||
This module exposes a TunableOp interface.
|
||||
|
||||
Some operations, such as GEMMs, could be implemented using more than one library
|
||||
or more than one technique. For example, a GEMM could be implemented for CUDA or
|
||||
ROCm using either the blas or blasLt libraries. Further, ROCm's rocblas and
|
||||
hipblaslt libraries allow the user to query for all possible algorithms and then
|
||||
choose one. How does one know which implementation is the fastest and should be
|
||||
chosen? That's what TunableOp provides.
|
||||
|
||||
Enabling TunableOp and Tuning Separately
|
||||
========================================
|
||||
|
||||
The TunableOp feature is enabled separately from enabling the tuning phase
|
||||
itself. Enabling TunableOp means that PyTorch will replace any standard
|
||||
operators with their Tunable implementations. Any call to a TunableOp first
|
||||
checks whether it has already been tuned for the given operator inputs. If so,
|
||||
it will immediately call the tuned operation; no further tuning will take place
|
||||
even when the tuning setting is enabled. Instead if no tuning result is found,
|
||||
and tuning is enabled, the TunableOp will benchmark every registered
|
||||
implementation of that operator for the given set of inputs and select the
|
||||
fastest.
|
||||
|
||||
File Input and Output
|
||||
=====================
|
||||
|
||||
The first time any TunableOp is invoked, the internal database of tuned
|
||||
operations will be prepared by attempting to read the results from the given
|
||||
file. The default filename is 'tunableop_results.csv'. To support tuning when
|
||||
multiple GPUs are used across multiple processes, the GPU device ordinal is
|
||||
automatically inserted into the filename to avoid multiple processes overwriting
|
||||
the same file.
|
||||
|
||||
If tuning is enabled and new tunings are discovered during the course of your
|
||||
workload, it will also write out to this same filename with all tunings, both
|
||||
the ones it read in at startup as well as the new ones found at runtime. This
|
||||
can be used, for example, to build up a tunings file across many workloads by
|
||||
reusing the same file. The output file is automatically created when the
|
||||
application terminates. This behavior can be controlled by the C++ and Python
|
||||
APIs but not the environment variables.
|
||||
|
||||
Assuming you specified a filename, you'll end up with a CSV file with contents
|
||||
like so::
|
||||
|
||||
Validator,PT_VERSION,2.2.0
|
||||
Validator,ROCM_VERSION,6.0.0.0-12969-1544e39
|
||||
Validator,HIPBLASLT_VERSION,0.6.0-a9c5cc7
|
||||
Validator,ROCBLAS_VERSION,4.0.0-72e57364-dirty
|
||||
GemmTunableOp_float_NT,nt_25088_4096_64,Gemm_Hipblaslt_1219,1.262
|
||||
GemmTunableOp_float_NT,nt_4096_4096_64,Gemm_Rocblas_1216,0.033
|
||||
|
||||
Note the "Validator" lines. If you change a library version, or ROCm version, or
|
||||
PyTorch version, TunableOp will detect this and reject the tunings file because
|
||||
the prior tunings are likely affected by other software changes.
|
||||
|
||||
The remaining lines are the tuned solutions for each TunableOp encountered
|
||||
during your execution. Each line consists of 4 comma-separated fields: operator
|
||||
name, operator parameters, solution name, and average execution time. The
|
||||
execution time is an optional field. The CSV file can be edited, but with
|
||||
caution. For example, the solution name (field 3) can be changed to "Default"
|
||||
and it will fall back to the original PyTorch untuned implementation. Or, in the
|
||||
case of ROCm's hipBLAS or hipBLASLt libraries, if you know the specific solution
|
||||
index you can override the solution that TunableOp selected by replacing the
|
||||
value. The operator name and parameters (fields 1 and 2) are internally named
|
||||
and should not be modified. In the case of GemmTunableOp, field 1 indicates the
|
||||
datatype and whether the inputs are transposed (T) or not (N) and field 2
|
||||
indicates the M, N, K input shapes.
|
||||
|
||||
There is an option to enable verbose output but it is only recommended for
|
||||
debugging purposes. This will produce a lot of diagnostic messages but may be
|
||||
useful to see if TunableOp is being used at all. Otherwise, TunableOp is
|
||||
completely silent, besides file output, unless there is a warning or error
|
||||
during its use. The verbose option is only available by setting the environment
|
||||
variable PYTORCH_TUNABLEOP_VEROBSE=1.
|
||||
|
||||
A Note on Tuning Behavior, Warmup, and Cache Effects
|
||||
====================================================
|
||||
|
||||
Tuning an operator consists of iterating through the list or registered
|
||||
implementations and profiling each one. The profile is established by running a
|
||||
single implementation in a loop multiple times and taking the average execution
|
||||
time. There is also an optional warmup phase prior to tuning that can help with
|
||||
reaching stable power states by the hardware. During tuning of a workload the
|
||||
various hardware caches will more likely produce hits than when not tuning.
|
||||
There are options for flushing the instruction cache and rotate the input tensors
|
||||
which might help produce a more faithful profile of the tuned operator as if the
|
||||
operator were run within a larger workload instead of in a tight, repetitive loop.
|
||||
|
||||
By default, each possible solution for a given operator will be run for either
|
||||
100 iterations or as many iterations that can be run within 30ms, whichever is
|
||||
smaller, and its average execution will be calculated. The fastest solution
|
||||
among all that were successfully profiled will be chosen. A profile might fail
|
||||
if the given solution doesn't achieve the same accuracy as the default
|
||||
implementation or if the solution returns an error code.
|
||||
|
||||
Current Tunable Operators
|
||||
=========================
|
||||
|
||||
TunableGemm for ROCm
|
||||
--------------------
|
||||
|
||||
Currently only a TunableGemm for ROCm is implemented. Note that CUDA builds of
|
||||
PyTorch will function correctly when using TunableOp but the only solution
|
||||
available to CUDA builds is the 'Default' implementation i.e. the original
|
||||
cuBLAS default, now called through TunableOp. Any call to at::cuda::blas::gemm()
|
||||
or ::bgemm() will be routed through TunableOp when enabled. Calling gemm() for a
|
||||
given set of input arguments (transa, transb, m, n, k) will attempt to use the
|
||||
fastest available implementation across both rocblas and hipblaslt.
|
||||
|
||||
Offline Tuning
|
||||
==============
|
||||
|
||||
Motivation
|
||||
----------
|
||||
There are several use cases for offline tuning.
|
||||
|
||||
One use case involves a workload with a high-memory utilization, where regular tuning might lead to running out of memory.
|
||||
|
||||
Another use case is for compute-intensive workloads. In such cases, it is more resource-efficient to collect
|
||||
the GEMMs for the workload once and then tune repeatedly with different tuning parameters or libraries.
|
||||
|
||||
Workflow
|
||||
--------
|
||||
There are basically two steps:
|
||||
1) Set the environment variables to collect the untuned GEMM and this will generate ``tunableop_untuned0.csv``:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
export PYTORCH_TUNABLEOP_ENABLED=1
|
||||
export PYTORCH_TUNABLEOP_TUNING=0
|
||||
export PYTORCH_TUNABLEOP_RECORD_UNTUNED=1
|
||||
...
|
||||
|
||||
2) Run a Python script that reads the ``tunableop_untuned0.csv`` and generates the ``tunableop_results0.csv``, like this:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch.cuda.tunable as tunable
|
||||
import os
|
||||
|
||||
os.putenv("PYTORCH_TUNABLEOP_ENABLED", "1")
|
||||
os.putenv("PYTORCH_TUNABLEOP_TUNING", "1")
|
||||
os.putenv("PYTORCH_TUNABLEOP_RECORD_UNTUNED", "0")
|
||||
tunable.tune_gemm_in_file("tunableop_untuned0.csv")
|
||||
|
||||
|
||||
It is also possible to take multiple untuned files and distribute the GEMMs for tuning to multiple GPUs
|
||||
within a single node. In the first step, the GEMMs are first gathered and duplicate GEMMs are eliminated.
|
||||
Next, the GEMMs are distributed to different GPUs for tuning. After all GEMMs are tuned, the results from
|
||||
all the GPUs are then gathered into a single file whose base filename has ``_full0`` appended to it
|
||||
(for example ``tunableop_results_full0.csv``). Finally, this new file, containing the gathered results, will be
|
||||
duplicated N times, once for each GPU as convenience to the user will run the workload with the tuned
|
||||
configuration on N GPUs.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
if __name__ == "__main__":
|
||||
num_gpus = 8 # number of GPUs that will be used during the tuning process
|
||||
tunable.mgpu_tune_gemm_in_file("tunableop_untuned?.csv", num_gpus)
|
||||
|
||||
Note that the usage of the ``mgpu_tune_gemm_in_file`` API is different from its single GPU counterpart
|
||||
(``tune_gemm_in_file``). The body of the Python script that calls the API must be wrapped in ``main()`` as shown
|
||||
due to the use of concurrent futures module. The argument to ``mgpu_tune_gemm_in_file`` must contain a wild card
|
||||
expression (``?`` or ``*``) to generate the list of untuned files containing the GEMMs to be processed. The ``num_gpus``
|
||||
must between 1 and the total number of GPUs available.
|
||||
|
||||
Tuning Context
|
||||
==============
|
||||
|
||||
The behavior of TunableOp is currently manipulated through environment
|
||||
variables, the C++ interface of at::cuda::tunable::getTuningContext(), or the
|
||||
torch.cuda.tunable python interfaces. The environment variables take precedence
|
||||
over any setting you manipulate using the C++ or Python APIs.
|
||||
|
||||
Environment Variable Interface
|
||||
------------------------------
|
||||
Environment variables are cached the first time they are read. You cannot use the
|
||||
environment variable interface programmatically since the settings become fixed.
|
||||
Use the C++ or Python APIs instead.
|
||||
|
||||
"""
|
||||
|
||||
import concurrent.futures
|
||||
import glob
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import shutil
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = [
|
||||
"enable",
|
||||
"is_enabled",
|
||||
"tuning_enable",
|
||||
"tuning_is_enabled",
|
||||
"record_untuned_enable",
|
||||
"record_untuned_is_enabled",
|
||||
"set_max_tuning_duration",
|
||||
"get_max_tuning_duration",
|
||||
"set_max_tuning_iterations",
|
||||
"get_max_tuning_iterations",
|
||||
"set_filename",
|
||||
"get_filename",
|
||||
"get_results",
|
||||
"get_validators",
|
||||
"read_file",
|
||||
"tune_gemm_in_file",
|
||||
"mgpu_tune_gemm_in_file",
|
||||
"set_rotating_buffer_size",
|
||||
"get_rotating_buffer_size",
|
||||
"set_numerical_check_tolerances",
|
||||
]
|
||||
|
||||
|
||||
def enable(val: bool = True) -> None:
|
||||
r"""This is the big on/off switch for all TunableOp implementations."""
|
||||
torch._C._cuda_tunableop_enable(val) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def is_enabled() -> bool:
|
||||
r"""Returns whether the TunableOp feature is enabled."""
|
||||
return torch._C._cuda_tunableop_is_enabled() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def tuning_enable(val: bool = True) -> None:
|
||||
r"""Enable tuning of TunableOp implementations.
|
||||
|
||||
When enabled, if a tuned entry isn't found, run the tuning step and record
|
||||
the entry.
|
||||
"""
|
||||
torch._C._cuda_tunableop_tuning_enable(val) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def tuning_is_enabled() -> bool:
|
||||
r"""Returns whether TunableOp implementations can be tuned."""
|
||||
return torch._C._cuda_tunableop_tuning_is_enabled() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def record_untuned_enable(val: bool = True) -> None:
|
||||
r"""Enable recording untuned of TunableOp perations for offline tuning.
|
||||
|
||||
When enabled, if a tuned entry isn't found, write it to the untuned file.
|
||||
"""
|
||||
torch._C._cuda_record_untuned_enable(val) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def record_untuned_is_enabled() -> bool:
|
||||
r"""Returns whether TunableOp operations are recorded for offline tuning."""
|
||||
return torch._C._cuda_record_untuned_is_enabled() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def set_max_tuning_duration(duration: int) -> None:
|
||||
r"""Set max time in milliseconds to spend tuning a given solution.
|
||||
|
||||
If both max tuning duration and iterations are set, the smaller of the two
|
||||
will be honored. At minimum 1 tuning iteration will always be run.
|
||||
"""
|
||||
torch._C._cuda_tunableop_set_max_tuning_duration(duration) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def get_max_tuning_duration() -> int:
|
||||
r"""Get max time to spend tuning a given solution."""
|
||||
return torch._C._cuda_tunableop_get_max_tuning_duration() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def set_max_tuning_iterations(iterations: int) -> None:
|
||||
r"""Set max number of iterations to spend tuning a given solution.
|
||||
|
||||
If both max tuning duration and iterations are set, the smaller of the two
|
||||
will be honored. At minimum 1 tuning iteration will always be run.
|
||||
"""
|
||||
torch._C._cuda_tunableop_set_max_tuning_iterations(iterations) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def get_max_tuning_iterations() -> int:
|
||||
r"""Get max iterations to spend tuning a given solution."""
|
||||
return torch._C._cuda_tunableop_get_max_tuning_iterations() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def set_filename(filename: str, insert_device_ordinal: bool = False) -> None:
|
||||
r"""Set the filename to use for input/output of tuning results.
|
||||
|
||||
If :attr:`insert_device_ordinal` is ``True`` then the current device ordinal
|
||||
will be added to the given filename automatically. This can be used in a
|
||||
1-process-per-gpu scenario to ensure all processes write to a separate file.
|
||||
"""
|
||||
torch._C._cuda_tunableop_set_filename(filename, insert_device_ordinal) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def get_filename() -> str:
|
||||
r"""Get the results filename."""
|
||||
return torch._C._cuda_tunableop_get_filename() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def get_results() -> tuple[str, str, str, float]:
|
||||
r"""Return all TunableOp results."""
|
||||
return torch._C._cuda_tunableop_get_results() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def get_validators() -> tuple[str, str]:
|
||||
r"""Return the TunableOp validators."""
|
||||
return torch._C._cuda_tunableop_get_validators() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def read_file(filename: str | None = None) -> bool:
|
||||
r"""Read results from a TunableOp CSV file.
|
||||
|
||||
If :attr:`filename` is not given, ``get_filename()`` is called.
|
||||
"""
|
||||
if filename is None:
|
||||
filename = get_filename()
|
||||
return torch._C._cuda_tunableop_read_file(filename) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def set_rotating_buffer_size(buffer_size: int) -> None:
|
||||
r"""Set rotating buffer size to this value in MB, if the buffer size is greater than zero.
|
||||
|
||||
If less than zero, query L2 cache size. If equal to zero, means deactivate rotating buffer.
|
||||
"""
|
||||
return torch._C._cuda_tunableop_set_rotating_buffer_size(buffer_size) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def get_rotating_buffer_size() -> int:
|
||||
r"""Get the rotating buffer size in kilobytes."""
|
||||
return torch._C._cuda_tunableop_get_rotating_buffer_size() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def set_numerical_check_tolerances(
|
||||
enable: bool, atol: float = 1e-5, rtol: float = 1e-5
|
||||
) -> None:
|
||||
r"""Set the atol and rtol values in numeric check"""
|
||||
return torch._C._cuda_tunableop_set_numerical_check_tolerances(enable, atol, rtol) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def tune_gemm_in_file(filename: str) -> None:
|
||||
r"""tune GEMM in file."""
|
||||
|
||||
if not is_enabled():
|
||||
raise AssertionError("TunableOp is not enabled")
|
||||
if not tuning_is_enabled():
|
||||
raise AssertionError("Tuning is not enabled")
|
||||
|
||||
deviceid = torch.cuda.current_device()
|
||||
|
||||
with open(filename) as file:
|
||||
for line in file:
|
||||
if line.startswith(("Gemm", "ScaledGemm")):
|
||||
_process_single_offline_gemm(line, deviceid)
|
||||
|
||||
|
||||
def _gather_unique_untuned_gemm_from_files(filename_pattern: str) -> set[str]:
|
||||
r"""Process multiple untuned results file and return a set with duplicates removed."""
|
||||
unique_gemm_entries = set() # set will avoid duplicates
|
||||
|
||||
for file_path in glob.glob(filename_pattern):
|
||||
with open(file_path) as file:
|
||||
for line in file:
|
||||
if line.startswith(("Gemm", "ScaledGemm")):
|
||||
unique_gemm_entries.add(line)
|
||||
|
||||
return unique_gemm_entries
|
||||
|
||||
|
||||
def _gather_tunableop_results() -> None:
|
||||
r"""Gather results from multiple tunableop results file and create a single file."""
|
||||
gemm_lines = set()
|
||||
validator_lines = []
|
||||
|
||||
# Need to allow for the possibility that results filename was
|
||||
# set with the Python API instead of with environment variable.
|
||||
# Also possible that results filename was not set at all.
|
||||
# There are several test cases to check, but ultimately we
|
||||
# need a glob-able expression
|
||||
results_filename = get_filename() # Note empty string could be returned here
|
||||
|
||||
if (
|
||||
results_filename is not None and results_filename != ""
|
||||
): # Case were the Python API was used to set the filename
|
||||
dot_pos = results_filename.find(".")
|
||||
if dot_pos != -1 and dot_pos > 0:
|
||||
# Replace the character just to the left of the dot
|
||||
filename_pattern = (
|
||||
results_filename[: dot_pos - 1] + "?" + results_filename[dot_pos:]
|
||||
)
|
||||
else:
|
||||
filename_pattern = "" # Needed to make linter happy
|
||||
else: # Case where the environment variable was used to set the filename.
|
||||
results_filename_env = os.getenv("PYTORCH_TUNABLEOP_FILENAME")
|
||||
if results_filename_env is None or results_filename_env == "":
|
||||
filename_pattern = "tunableop_results?.csv"
|
||||
elif "%d" in results_filename_env:
|
||||
filename_pattern = results_filename_env.replace("%d", "?")
|
||||
else:
|
||||
filename_pattern = results_filename_env.replace(".", "?.")
|
||||
|
||||
if "?" not in filename_pattern:
|
||||
raise AssertionError(
|
||||
f"filename_pattern must contain '?', got {filename_pattern!r}"
|
||||
)
|
||||
|
||||
FirstFile = False
|
||||
matching_files = glob.glob(filename_pattern)
|
||||
num_matching_files = len(matching_files)
|
||||
for file_path in matching_files:
|
||||
with open(file_path) as file:
|
||||
for line in file:
|
||||
if line.startswith("Validator"):
|
||||
if not (FirstFile):
|
||||
# Only read Validator from first file
|
||||
validator_lines.append(line)
|
||||
else:
|
||||
gemm_lines.add(line)
|
||||
|
||||
FirstFile = True
|
||||
|
||||
output_file = filename_pattern.replace("?", "_full0")
|
||||
|
||||
with open(output_file, "w") as out_file:
|
||||
for line in validator_lines:
|
||||
out_file.write(line)
|
||||
for line in gemm_lines:
|
||||
out_file.write(line)
|
||||
|
||||
# Create num_matching_copies of the results file
|
||||
for i in range(1, num_matching_files):
|
||||
duplicate_file = output_file.replace("0", str(i))
|
||||
shutil.copy(output_file, duplicate_file)
|
||||
|
||||
|
||||
def _create_matrices(
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
lda: int,
|
||||
ldb: int,
|
||||
ldc: int,
|
||||
transA: bool,
|
||||
transB: bool,
|
||||
dtypeA: torch.dtype,
|
||||
deviceid: str,
|
||||
dtypeB: torch.dtype | None = None,
|
||||
randn: bool = True,
|
||||
subMatrix: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
r"""Helper function for _process_single_offline_gemm.
|
||||
Creates matrices that are then consumed by one of the Torch GEMM APIs.
|
||||
"""
|
||||
# Fill parameters set for use with ScaledGEMM
|
||||
fillA = 0.25
|
||||
fillB = 0.75
|
||||
|
||||
if dtypeB is None:
|
||||
dtypeB = dtypeA
|
||||
|
||||
if subMatrix:
|
||||
# User reference for understanding leading dimension:
|
||||
# https://github.com/Reference-LAPACK/lapack/blob/master/BLAS/SRC/dgemm.f
|
||||
# TO DO: According to lines 108 - 133, there is no lower bound on rowsA,
|
||||
# but there is a restriction on rowsB. Using this formula for now as it
|
||||
# seems to work for all UTs.
|
||||
rowsA = rowsB = max(ldc, k)
|
||||
|
||||
if randn:
|
||||
matA = torch.randn(rowsA, lda, dtype=dtypeA, device=deviceid)
|
||||
matB = torch.randn(rowsB, ldb, dtype=dtypeA, device=deviceid)
|
||||
else:
|
||||
matA = torch.full((rowsA, lda), fillA, dtype=dtypeB, device=deviceid)
|
||||
matB = torch.full((rowsB, ldb), fillB, dtype=dtypeB, device=deviceid)
|
||||
|
||||
subA = matA[:k, :m].t() if transA else matA[:m, :k]
|
||||
subB = matB[:n, :k].t() if transB else matB[:k, :n]
|
||||
return subA, subB
|
||||
else:
|
||||
if randn:
|
||||
matA = (
|
||||
torch.rand(k, m, dtype=dtypeA, device=deviceid).t()
|
||||
if transA
|
||||
else torch.rand(m, k, dtype=dtypeA, device=deviceid)
|
||||
)
|
||||
matB = (
|
||||
torch.rand(n, k, dtype=dtypeB, device=deviceid).t()
|
||||
if transB
|
||||
else torch.rand(k, n, dtype=dtypeB, device=deviceid)
|
||||
)
|
||||
else:
|
||||
matA = (
|
||||
torch.full((k, m), fillA, dtype=dtypeA, device=deviceid).t()
|
||||
if transA
|
||||
else torch.full((m, k), fillA, dtype=dtypeA, device=deviceid)
|
||||
)
|
||||
matB = (
|
||||
torch.full((n, k), fillB, dtype=dtypeB, device=deviceid).t()
|
||||
if transB
|
||||
else torch.full((k, n), fillB, dtype=dtypeB, device=deviceid)
|
||||
)
|
||||
return matA, matB
|
||||
|
||||
|
||||
def _create_batch_matrices(
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
b: int,
|
||||
lda: int,
|
||||
ldb: int,
|
||||
ldc: int,
|
||||
transA: bool,
|
||||
transB: bool,
|
||||
dtype: torch.dtype,
|
||||
deviceid: str,
|
||||
subMatrix: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
r"""Helper function for _process_single_offline_gemm.
|
||||
Creates batch matrices that are then consumed by one of the Torch GEMM APIs.
|
||||
Similar to _create_matrices but for 3D batch matrices.
|
||||
"""
|
||||
if subMatrix:
|
||||
# User reference for understanding leading dimension:
|
||||
# https://github.com/Reference-LAPACK/lapack/blob/master/BLAS/SRC/dgemm.f
|
||||
# TO DO: According to lines 108 - 133, there is no lower bound on rowsA,
|
||||
# but there is a restriction on rowsB. Using this formula for now as it
|
||||
# seems to work for all UTs.
|
||||
rowsA = rowsB = max(ldc, k)
|
||||
|
||||
matA = torch.randn(b, rowsA, lda, dtype=dtype, device=deviceid)
|
||||
matB = torch.randn(b, rowsB, ldb, dtype=dtype, device=deviceid)
|
||||
|
||||
subA = matA[:b, :k, :m].transpose(1, 2) if transA else matA[:b, :m, :k]
|
||||
subB = matB[:b, :n, :k].transpose(1, 2) if transB else matB[:b, :k, :n]
|
||||
return subA, subB
|
||||
else:
|
||||
matA = (
|
||||
torch.rand(b, k, m, dtype=dtype, device=deviceid)
|
||||
if transA
|
||||
else torch.rand(b, m, k, dtype=dtype, device=deviceid)
|
||||
)
|
||||
matB = (
|
||||
torch.rand(b, n, k, dtype=dtype, device=deviceid)
|
||||
if transB
|
||||
else torch.rand(b, k, n, dtype=dtype, device=deviceid)
|
||||
)
|
||||
matA = matA.transpose(1, 2) if transA else matA
|
||||
matB = matB.transpose(1, 2) if transB else matB
|
||||
return matA, matB
|
||||
|
||||
|
||||
def _process_single_offline_gemm(untuned_gemm_line: str, gpu_id: int) -> None:
|
||||
r"""Process a single untuned GEMM."""
|
||||
|
||||
deviceid = "cuda:" + str(gpu_id)
|
||||
|
||||
dtype_dict = {
|
||||
"float": torch.float32,
|
||||
"tf32": torch.float32,
|
||||
"double": torch.float64,
|
||||
"BFloat16": torch.bfloat16,
|
||||
"Half": torch.half,
|
||||
"c10::complex<double>": torch.complex128,
|
||||
"c10::complex<float>": torch.complex64,
|
||||
"Float8_e4m3fn": torch.float8_e4m3fn,
|
||||
"Float8_e5m2": torch.float8_e5m2,
|
||||
"Float8_e4m3fnuz": torch.float8_e4m3fnuz,
|
||||
"Float8_e5m2fnuz": torch.float8_e5m2fnuz,
|
||||
}
|
||||
|
||||
untuned_gemm = untuned_gemm_line.strip().split(",")[:]
|
||||
|
||||
underscore_count = untuned_gemm[0].count("_")
|
||||
|
||||
# Initialize dtype to make linter happy
|
||||
dtype = None
|
||||
dtypeA = None
|
||||
dtypeB = None
|
||||
dtypeC = None
|
||||
|
||||
# Extract BLAS parameters
|
||||
if underscore_count == 2:
|
||||
[op_sig, data_type, layout] = untuned_gemm[0].split("_")
|
||||
transB = layout[0] == "T"
|
||||
transA = layout[1] == "T"
|
||||
dtype = dtype_dict.get(data_type)
|
||||
if data_type == "tf32":
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
else:
|
||||
torch.backends.cuda.matmul.allow_tf32 = False
|
||||
|
||||
else: # ScaledGEMM
|
||||
count = untuned_gemm[0].count("_")
|
||||
if count not in [6, 7]:
|
||||
raise AssertionError(f"count must be 6 or 7, got {count}")
|
||||
untuned_gemm_temp = untuned_gemm[0].split("_")
|
||||
# dtypeC = might not be FP8 type, keep track
|
||||
# of the number of underscores
|
||||
op_sig = untuned_gemm_temp[0]
|
||||
data_typeA = untuned_gemm_temp[1] + "_" + untuned_gemm_temp[2]
|
||||
data_typeB = untuned_gemm_temp[3] + "_" + untuned_gemm_temp[4]
|
||||
if count == 7:
|
||||
data_typeC = untuned_gemm_temp[5] + "_" + untuned_gemm_temp[6]
|
||||
else:
|
||||
data_typeC = untuned_gemm_temp[5]
|
||||
transB = untuned_gemm_temp[count][0] == "T"
|
||||
transA = untuned_gemm_temp[count][1] == "T"
|
||||
dtypeA = dtype_dict.get(data_typeA)
|
||||
dtypeB = dtype_dict.get(data_typeB)
|
||||
dtypeC = dtype_dict.get(data_typeC)
|
||||
|
||||
untuned_gemm_temp = untuned_gemm[1].split("_")
|
||||
[n, m, k] = [int(g) for g in untuned_gemm_temp[1:4]]
|
||||
if op_sig == "GemmStridedBatchedTunableOp":
|
||||
if untuned_gemm_temp[6] != "ld":
|
||||
raise AssertionError(
|
||||
f"expected 'ld' at index 6, got {untuned_gemm_temp[6]!r}"
|
||||
)
|
||||
[ldb, lda, ldc] = [int(g) for g in untuned_gemm_temp[7:10]]
|
||||
else:
|
||||
if untuned_gemm_temp[4] != "ld":
|
||||
raise AssertionError(
|
||||
f"expected 'ld' at index 4, got {untuned_gemm_temp[4]!r}"
|
||||
)
|
||||
[ldb, lda, ldc] = [int(g) for g in untuned_gemm_temp[5:8]]
|
||||
|
||||
# Detect subMatrix case
|
||||
if all(item in [n, m, k] for item in [lda, ldb, ldc]):
|
||||
subMatrix = False
|
||||
else:
|
||||
subMatrix = True
|
||||
|
||||
if op_sig == "GemmTunableOp":
|
||||
# Warnings for unsupported cases:
|
||||
if m == 1 or n == 1 or k == 1:
|
||||
if (not transA) and (not transB):
|
||||
pass # case is supported
|
||||
elif transA and n == 1:
|
||||
pass # case is supported
|
||||
else:
|
||||
warnings.warn(
|
||||
"Offline tuning is not supported for this GEMM. Use online tuning instead. "
|
||||
+ f"Skipped tuning for: {untuned_gemm[1]}",
|
||||
stacklevel=2,
|
||||
)
|
||||
return
|
||||
|
||||
# Resolve linter issue
|
||||
if dtype is None or not isinstance(dtype, torch.dtype):
|
||||
raise TypeError(f"dtype must be a torch.dtype, but got {dtype}")
|
||||
|
||||
matA, matB = _create_matrices(
|
||||
m, n, k, lda, ldb, ldc, transA, transB, dtype, deviceid, subMatrix=subMatrix
|
||||
)
|
||||
torch.mm(matA, matB)
|
||||
|
||||
elif op_sig == "GemmStridedBatchedTunableOp":
|
||||
# Warnings for unsupported cases:
|
||||
if m == 1 or n == 1 or k == 1:
|
||||
warnings.warn(
|
||||
"Offline tuning is not support for this GEMM. Use online tuning instead. "
|
||||
+ f"Skipped tuning for: {untuned_gemm[1]}",
|
||||
stacklevel=2,
|
||||
)
|
||||
return
|
||||
|
||||
[b] = [int(g) for g in untuned_gemm_temp[5:6]]
|
||||
|
||||
# Resolve linter issue
|
||||
if dtype is None or not isinstance(dtype, torch.dtype):
|
||||
raise TypeError(f"dtype must be a torch.dtype, but got {dtype}")
|
||||
|
||||
matA, matB = _create_batch_matrices(
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
b,
|
||||
lda,
|
||||
ldb,
|
||||
ldc,
|
||||
transA,
|
||||
transB,
|
||||
dtype,
|
||||
deviceid,
|
||||
subMatrix=subMatrix,
|
||||
)
|
||||
torch.bmm(matA, matB)
|
||||
elif op_sig == "ScaledGemmTunableOp":
|
||||
# Only combination supported by PyTorch
|
||||
if transB is not True:
|
||||
raise AssertionError(
|
||||
f"transB must be True for ScaledGemmTunableOp, got {transB}"
|
||||
)
|
||||
if transA is not False:
|
||||
raise AssertionError(
|
||||
f"transA must be False for ScaledGemmTunableOp, got {transA}"
|
||||
)
|
||||
|
||||
# Resolve linter issue
|
||||
if dtypeA is None or not isinstance(dtypeA, torch.dtype):
|
||||
raise TypeError(f"dtype must be a torch.dtype, but got {dtypeA}")
|
||||
|
||||
matA, matB = _create_matrices(
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
lda,
|
||||
ldb,
|
||||
ldc,
|
||||
transA,
|
||||
transB,
|
||||
dtypeA,
|
||||
deviceid,
|
||||
dtypeB=dtypeB,
|
||||
randn=False,
|
||||
subMatrix=subMatrix,
|
||||
)
|
||||
|
||||
if untuned_gemm_temp[8] != "rw":
|
||||
raise AssertionError(
|
||||
f"expected 'rw' at index 8, got {untuned_gemm_temp[8]!r}"
|
||||
)
|
||||
if untuned_gemm_temp[9] == "1":
|
||||
rowwise = True
|
||||
else:
|
||||
rowwise = False
|
||||
if rowwise:
|
||||
scaleA = (
|
||||
torch.ones((1, m), device=deviceid)
|
||||
if transA
|
||||
else torch.ones((m, 1), device=deviceid)
|
||||
)
|
||||
scaleB = (
|
||||
torch.ones((1, n), device=deviceid)
|
||||
if transB
|
||||
else torch.ones((n, 1), device=deviceid)
|
||||
)
|
||||
else:
|
||||
scaleA = torch.tensor(0.8, device=deviceid)
|
||||
scaleB = torch.tensor(0.9, device=deviceid)
|
||||
|
||||
if untuned_gemm_temp[10] != "bias":
|
||||
raise AssertionError(
|
||||
f"expected 'bias' at index 10, got {untuned_gemm_temp[10]!r}"
|
||||
)
|
||||
if untuned_gemm_temp[11] == "None": # no bias vector
|
||||
torch._scaled_mm(
|
||||
matA, matB, scale_a=scaleA, scale_b=scaleB, out_dtype=dtypeC
|
||||
)
|
||||
else: # bias vector present
|
||||
fillbias = 0.10
|
||||
bias_dtype = dtype_dict.get(untuned_gemm_temp[11])
|
||||
bias = (
|
||||
torch.full((n,), fillbias, dtype=bias_dtype, device=deviceid)
|
||||
if transB
|
||||
else torch.full((m,), fillbias, dtype=bias_dtype, device=deviceid)
|
||||
)
|
||||
torch._scaled_mm(
|
||||
matA, matB, scale_a=scaleA, scale_b=scaleB, out_dtype=dtypeC, bias=bias
|
||||
)
|
||||
|
||||
elif op_sig == "GemmAndBiasTunableOp":
|
||||
# y = x*A^T + b
|
||||
if transA == transB:
|
||||
raise AssertionError(
|
||||
f"transA and transB must differ for GemmAndBiasTunableOp, got transA={transA}, transB={transB}"
|
||||
)
|
||||
|
||||
# Resolve linter issue
|
||||
if dtype is None or not isinstance(dtype, torch.dtype):
|
||||
raise TypeError(f"dtype must be a torch.dtype, but got {dtype}")
|
||||
|
||||
bias = torch.rand(n, dtype=dtype, device=deviceid)
|
||||
|
||||
X, matA = _create_matrices(
|
||||
m, n, k, lda, ldb, ldc, transA, transB, dtype, deviceid, subMatrix=subMatrix
|
||||
)
|
||||
matA = matA.t()
|
||||
torch.nn.functional.linear(X, matA, bias)
|
||||
else:
|
||||
warnings.warn(f"error: unknown op {op_sig}", stacklevel=2)
|
||||
|
||||
|
||||
def _check_tuning_assertions() -> None:
|
||||
r"""Helper function for multi-GPU tuning case. Need to check that TunableOp feature
|
||||
is enabled and that tuning is enabled.
|
||||
"""
|
||||
|
||||
if is_enabled() is False:
|
||||
warnings.warn("TunableOp was disabled. Trying to enable now.", stacklevel=2)
|
||||
enable(True)
|
||||
if is_enabled() is not True:
|
||||
raise AssertionError("is_enabled() must be True")
|
||||
if tuning_is_enabled() is not True:
|
||||
raise AssertionError("tuning_is_enabled() must be True")
|
||||
if record_untuned_is_enabled() is not False:
|
||||
raise AssertionError("record_untuned_is_enabled() must be False")
|
||||
|
||||
|
||||
def mgpu_tune_gemm_in_file(filename_pattern: str, num_gpus: int) -> None:
|
||||
r"""Process one or more files and distribute work over one or more GPUs."""
|
||||
unique_gemm_entries = _gather_unique_untuned_gemm_from_files(filename_pattern)
|
||||
|
||||
total_gpus = torch.cuda.device_count()
|
||||
|
||||
if not (1 <= num_gpus <= total_gpus):
|
||||
raise AssertionError(
|
||||
f"num_gpus must be between 1 and {total_gpus}, got {num_gpus}"
|
||||
)
|
||||
|
||||
mp_context = mp.get_context("spawn")
|
||||
|
||||
futures = [] # empty list to hold futures
|
||||
|
||||
# GEMM are assigned to GPUs in a round robin manner
|
||||
h = 0
|
||||
with concurrent.futures.ProcessPoolExecutor(
|
||||
max_workers=num_gpus,
|
||||
mp_context=mp_context,
|
||||
initializer=_check_tuning_assertions,
|
||||
) as executor:
|
||||
# The workers are a separate process. TunableOp will be
|
||||
# enabled in the child processes if PYTORCH_TUNABLEOP_ENABLED=1
|
||||
# In the initializer, we also try to enable TunableOP if th
|
||||
# environment variable was NOT set.
|
||||
|
||||
for line in unique_gemm_entries:
|
||||
future = executor.submit(_process_single_offline_gemm, line, h)
|
||||
futures.append(future)
|
||||
h = (h + 1) % num_gpus
|
||||
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
future.result()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
_gather_tunableop_results()
|
||||
Reference in New Issue
Block a user