Интеграция профилировщика

Обзор архитектуры

Такое разделение позволяет PyTorch оставаться независимым от устройства: Python управляет сеансом, ProfilerStubs преобразуют запросы профилировщика в вызовы среды выполнения бэкенда, а среда выполнения взаимодействует с ускорителем.

Основные соглашения

• Хуки записи: record() должны получить (при наличии) индекс устройства, создать событие бэкенда, при необходимости сохранить временную метку CPU и поставить событие в очередь активного потока.
• Время выполнения: elapsed() отвечает за синхронизацию отдельных событий и возврат длительности в микросекундах.
• Хуки синхронизации: synchronize() и onEachDevice() обеспечивают согласование переходов между фазами (например, разогрев → активная фаза) на всех устройствах.
• Аннотации: mark, rangePush и rangePop можно реализовать для обогащения трассировок; в противном случае их можно оставить пустыми.

Заглушки профилировщика (C++)

torch::profiler::impl::OpenRegMethods наследуется от ProfilerStubs и связывает описанные выше хуки:

Конструктор однократно регистрирует методы с помощью registerPrivateUse1Methods(&methods);, делая их доступными, когда профилировщик включен с параметром use_device="openreg".

Уровень управления Python

На стороне Python новая точка входа не требуется — разработчики используют стандартный профилировщик autograd:

from torch.autograd.profiler import profile as autograd_profile
from torch.profiler import record_function

with autograd_profile(use_device="openreg", record_shapes=True) as prof:
    with record_function("matmul"):
        x = torch.randn(512, 512, device="openreg")
        y = torch.randn(512, 512, device="openreg")
        z = x @ y

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
prof.export_chrome_trace("openreg_trace.json")

Поток сбора данных

1. Пользовательский код входит в autograd_profile(use_device="openreg").
2. Профилировщик переходит в состояние ProfilerState.KINETO_PRIVATEUSE1_FALLBACK.
3. Профилировщик запрашивает у активного бэкенда запись события с помощью record().
4. Заглушки OpenReg создают объекты orEvent, связывают их с текущим потоком и сохраняют временные метки CPU.
5. По завершении событий профилировщик вызывает elapsed() для вычисления длительности.

Библиотека устройства: отслеживание корреляции

Библиотека устройства предоставляет openreg::profiler::OpenRegTracer (third_party/openreg/csrc/tracer.h/.cpp) — синглтон со стеком идентификаторов корреляции, локальным для потока, и атомарным флагом включения/отключения, с помощью которого сеанс профилировщика управляет окном записи.

Kineto помещает идентификаторы корреляции в стек и извлекает их из него через сеанс. Сеанс вызывает API активности в стиле C из openreg.h (по аналогии с CUPTI):

• orActivityEnableTracing() / orActivityDisableTracing() — управляют окном записи
• orActivityPushExternalCorrelationId() / orActivityPopExternalCorrelationId() — управляют стеком корреляции

Эквивалентом у реального поставщика будет его SDK для трассировки устройства (например, CUPTI для CUDA).

Интеграция с PyTorch: IActivityProfiler и IActivityProfilerSession

Реализуйте два интерфейса Kineto из third_party/kineto/libkineto/include/IActivityProfiler.h. В OpenReg они находятся в torch_openreg/csrc/profiler/ — уровне интеграции расширения бэкенда.

• IActivityProfiler — фабрика без состояния. Обе перегрузки configure() создают и возвращают сеанс:

  • configure(activity_types, config) — синхронная перегрузка, обязательная для интерфейса. Заглушка OpenReg реализует ее как основной путь создания сеанса.
  • configure(ts_ms, duration_ms, activity_types, config) — дочерний профилировщик Kineto вызывает эту перегрузку для всех трассировок (включая запускаемые по запросу), передавая profileStartTime() эпохи в миллисекундах и profileDuration() в миллисекундах. Заглушка OpenReg игнорирует планирование и делегирует вызов первой перегрузке; поставщики используют ts_ms/duration_ms, чтобы отложить активацию SDK устройства.

• IActivityProfilerSession — сеанс для отдельной трассировки. start()/stop() управляют окном профилирования и переключают трассировку активности через orActivityEnableTracing()/orActivityDisableTracing(); getTraceBuffer() возвращает буфер в Kineto.

  • Эталонная заглушка: processTrace() только задает интервал трассировки (traceBuffer_.span = TraceSpan(startTs_, endTs_, "openreg")); записи ядер не создаются.
  • Расширение поставщика: замените processTrace(), чтобы сбрасывать записи из SDK трассировки устройства и создавать элементы GenericTraceActivity с временными метками (мкс), идентификаторами корреляции и связями потоков (flow.id = correlationId, flow.type = kLinkAsyncCpuGpu, flow.start = 0).

Регистрация и сборка

Для регистрации достаточно одной строки: REGISTER_PRIVATEUSE1_PROFILER(OpenRegActivityProfiler). Макрос (определенный в torch/csrc/profiler/standalone/privateuse1_profiler.h) создает статический объект регистрации, который передает фабрику в Kineto при инициализации профилировщика.

По умолчанию Kineto включен в PyTorch; специальные флаги сборки не нужны, если только он явно не отключен с помощью USE_KINETO=0. В расширении бэкенда find_package(Torch) задает kineto_LIBRARY; свяжите код с kineto и torch_cpu_library и защитите код Kineto с помощью #ifdef USE_KINETO. Без Kineto плагин компилируется как пустая реализация, и доступен только способ с устаревшими заглушками.

Использование

import torch
from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1]) as prof:
    x = torch.randn(512, 512, device="openreg")
    y = torch.randn(512, 512, device="openreg")
    z = x @ y

prof.export_chrome_trace("kernel_trace.json")