Обзор стражей

С точки зрения пользовательского интерфейса TorchDynamo очень прост в использовании. Пользователь вызывает torchdynamo.optimize в качестве аннотации:

@torchdynamo.optimize(my_compiler)
def fn_foo(bar):

Где полный пример выглядит так:

from typing import List
import torch
from torch import _dynamo as torchdynamo

def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    print("my_compiler() called with FX graph:")
    gm.graph.print_tabular()
    return gm.forward  # return a python callable

@torchdynamo.optimize(my_compiler)
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    if b.sum() < 0:
        b = b * -1
    return x * b

for _ in range(100):
    toy_example(torch.randn(10), torch.randn(10))

Это позволяет TorchDynamo захватывать интерпретированные фреймы Python, извлекать всю необходимую информацию и ускорять работу там, где это возможно. Ускорение происходит по нескольким причинам и может сильно зависеть от используемого бэкэнда (my_compiler в примере выше), но важным ускоряющим фактором в данном разделе является кеширование. Само по себе кеширование не является прямым ускоряющим фактором, но является критически важным механизмом, предотвращающим перекомпиляцию. Мы выкапываем яму с помощью dynamo, а кеширование позволяет нам из неё выбраться. Оно позволяет сохранять нейтральность производительности, одновременно позволяя бэкендам — истинному источнику нашего ускорения.

Даже с предоставленным бэкэндом, выполняющим проход без операций:

def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    return gm.forward

Мы можем увидеть, как TorchDynamo ускоряет выполнение Python, даже на обычном Python, а не только на PyTorch.

Обзор кеширования и стражей

TorchDynamo работает через кеширование преобразованного (TorchDynamo) байткода пользователя. Когда TorchDynamo получает фрейм для вычисления, он проверяет, изменились ли какие-либо объекты, на которые ссылается фрейм, и если нет, TorchDynamo считывает ранее преобразованный байт-код пользователя для его вычисления. В этом разделе мы сосредоточимся на том, как мы можем определить, изменились ли объекты, на которые ссылается фрейм. Это критически важная функциональность в TorchDynamo, так как она управляет всем жизненным циклом проверки. Эта функциональность называется стражами.

На очень высоком уровне поток можно обобщить так:

1. TorchDynamo получает фрейм Python.
2. Он преобразует фрейм (1), пропуская его через перевод инструкций.

3. Для объектов, захваченных в (2), TorchDynamo создаёт объекты отслеживания, которые являются:

   • отслеживаемыми в графе вывода, который представляет собой внутреннюю специализацию torch.fx.Tracer
   • стражами
4. TorchDynamo обрабатывает объекты стражей, созданные в (3), преобразуя их в генерируемую функцию Python, check_fn, связанную с фрагментом кода.
5. check_fn вычисляется всякий раз, когда мы сталкиваемся с этим кодом в последующее время — если check_fn проходит и вычисляет True, TorchDynamo идентифицирует код в кэше и код, встреченный здесь, как одинаковые, и может безопасно его использовать. Если проверка завершается неудачно и вычисляется False, TorchDynamo определяет, что код в кэше недействителен, и может его выбросить в пользу новой записи, через перекомпиляцию или разрыв графа.

Вычисление фрейма Python и PEP 523

Функциональность TorchDynamo основана на PEP 523.

TorchDynamo устанавливает функцию вычисления фреймов в Python, используя _PyInterpreterState_SetEvalFrameFunc. TorchDynamo имеет обработчик, где Python может вернуть управление нам во время вычисления.

Функция, которую мы установили, это convert_frame или convert_frame_assert в случае nopython=True, но сейчас давайте обойдём этот нюанс и взглянем на convert_frame_assert, так как convert_frame делегирует к ней.

Мы можем найти её на строке 222 convert_frame.py,

с подписью следующим образом:

def  convert_frame_assert(compiler_fn: Callable, one_graph=True):

Эта функция оборачивает точку входа, где Python вызывает TorchDynamo с фреймом:

def  _convert_frame_assert(frame: types.FrameType, cache_size: int):

Вот что делает эта функция:

1. Проверяет, видела ли она этот code(см.: f_code здесь) ранее, и выходит рано, если это так.
2. Проверяет, не является ли код недопустимым случаем.
3. Проверяет, не превышает ли cache_size (второй аргумент выше) ограничения, определённого в конфигурации, cache_size_limit. Если превышает, функция отбрасывает фрейм и регистрирует предупреждения. Это помогает избежать постоянной перекомпиляции фрейма, так как, как правило, это означает, что фрейм активен неожиданным образом, и кеширование его приводит к ненужной нагрузке, поскольку он, скорее всего, будет вытеснен в следующий раз.

4. Передаёт фрейм, а также функцию, создающую InstructionTranslator через преобразование байткода, посредством transform_code_object. Здесь происходит несколько важных событий:

   1. Новый код генерируется посредством transform_code_object.
   2. Производится трекер FX с именем output посредством InstructionTranslator. Это может быть немного запутанно, так как InstructionTranslator не является трекером fx, но он хранится в переменной с именем tracer, и его выходные данные являются трекером fx.
   3. Функция создаёт стражи и сохраняет их в output выше.
   4. Функция создаёт output_instructions и сохраняет их в output выше.
   5. Функция сопоставляет новый созданный преобразованный код с исходным кодом, который она прочитала из фрейма. Это сопоставление стоит запомнить, мы будем ссылаться на него позже, в разделе об ошибках стражей.
5. Используя преобразованный код из 4.1 и стражи из 4.3, функция создаёт GuardedCode.

Теперь, когда мы узнали о вычислении фреймов, давайте рассмотрим InstructionTranslator, и посмотрим, как она превращает фрейм, который мы передали ей, во внутренние типы TorchDynamo.

InstructionTranslator

InstructionTranslator делает много! Мы не будем углубляться в детали всего, что она делает, но самое важное для этого документа — это то, что она генерирует сопоставление symbolic_locals, которое поддерживает сопоставление f_locals фрейма с внутренними объектами переменных TorchDynamo (подробнее об этом чуть позже. symbolic_locals заполняется путём обхода локальных переменных фрейма:

self.symbolic_locals = collections.OrderedDict(
    (k, VariableBuilder(self, LocalSource(k))(f_locals[k]))
    for k in vars
    if k in f_locals
)

Важной составляющей является вызов функции VariableBuilder. Реализация вызова VariableBuilder перенаправляется в функцию под названием _wrap, которая в свою очередь как создаёт экземпляры VariableTracker, так и вызывает make_guards на них. Подробнее об этом позже.

Это сопоставление, в свою очередь, является критически важным, так как каждая переменная имеет связанные стражи, которые затем передаются self.output, экземпляру OutputGraph, трекеру fx, упомянутому в 4.2 раздела выше. Если вы помните, этот OutputGraph, сохранённый в переменной с именем output, — это место, где хранятся наши стражи, прежде чем они будут переданы, чтобы стать GuardedCode

Как InstructionTranslator делает это? В основе всего лежит цикл, который перекачивает функцию step.

step — это всего лишь один шаг обработки, принимающий ровно одну инструкцию и делающий что-то с ней.

Для примера выше, вот фрагмент того, как могут выглядеть несколько Instruction:

Instruction(opcode=124, opname='LOAD_FAST', arg=0, argval='b', offset=32, starts_line=8, is_jump_target=True, target=None)
Instruction(opcode=100, opname='LOAD_CONST', arg=3, argval=-1, offset=34, starts_line=None, is_jump_target=False, target=None)
Instruction(opcode=20, opname='BINARY_MULTIPLY', arg=None, argval=None, offset=36, starts_line=None, is_jump_target=False, target=None)

Это основная функциональность этой функции. Посмотрите на opname, а затем на этот небольшой фрагмент изнутри step;

if not hasattr(self, inst.opname):
    unimplemented(f"missing: {inst.opname}")
getattr(self, inst.opname)(inst)

Как мы видим, функция проверяет, имеет ли текущий класс, InstructionTranslator, установленный атрибут, соответствующий имени оператора (например, LOAD_CONST). Если да, функция вызывает его, передавая весь объект инструкции. Если нет, функция отбрасывает фрейм как нереализованный.

Для примера LOAD_CONST мы видим, что мы действительно его поддерживаем, с относительно простой реализацией:

def LOAD_CONST(self, inst):
    self.push(ConstantVariable(value=inst.argval))

Мы видим, что эта функция создаёт новый экземпляр класса ConstantVariable, со значением, в нашем примере, -1, и затем помещает его в стек.

Существует десятки таких методов — см. symbolic_convert.py для всех из них. В целом, мы реализуем как можно больше соответствующих методов для инструкций байткода Python.

Как по логике ниже step, так и по логике вызова VariableBuilder — у нас теперь много VariableTracker и, конечно же, мы довольно много говорили о создании стражей. Давайте углубимся в то, что такое переменные, и немного ближе узнаем, что такое стражи.

Переменные

ConstantVariable — это экземпляр VariableTracker. VariableTracker представляет отслеживаемое локальное значение Python или значение стека.

Когда дело доходит до представления объекта внутри TorchDynamo, VariableTracker делает именно то, что говорит — отслеживает заданную переменную. Это очень гибкий класс, но есть несколько моментов, которые следует учитывать:

• Он управляет guard отношением вокруг базового объекта посредством:

  • make_guard
  • replace_guards
  • add_guard(s)
  • propagate - propagate(*vars: List[List["VariableTracker"]]) - Возможно, самый важный из них, так как он объединяет стражи из всех переданных VariableTracker экземпляров. Он посещает стражи и объединяет стражи из этих стражей на себя.

• Он действует как прокси от имени базового объекта, реализуя методы для получения информации об отслеживаемом объекте остальной частью TorchDynamo:

  • call_method
  • call_function
  • python_type
  • as_proxy
  • is/as_python_proxy
• Он хранит переменную source типа Source, из torchdynamo/source.py. Этот тип источника — относительно самодостаточный класс, который помогает нам организовать и контролировать, откуда пришёл исходный источник, а также предоставляет удобные методы для получения имени и, что важно для нас, для создания стражей.

И этот класс (VariableTracker) построен на основе наследования, где-то между полным классом абстрактного базового класса и полностью сформированным классом — он оставляет многие методы, вызывающие NotImplementedError — со ссылкой на подклассы. Смотрите torchdynamo/variables/ для всех подклассов, чтобы выполнить контракты и пользовательские поведения.

Зная всё это, мы можем увидеть пример того, как инструкция из dis, BUILD_TUPLE:

BUILD_TUPLE(count) Создаёт кортеж, потребляя count элементов из стека, и помещает полученный кортеж в стек.

В нашем случае наша подпись будет немного отличаться из-за способа создания объектов Instruction, но суть останется той же. Вместо передачи count, мы передаем объект с дополнительной информацией, и, конечно же, мы занимаемся преобразованием обычных объектов Python в понятия TorchDynamo:

def BUILD_TUPLE(self, inst):
    items = self.popn(inst.argval)
    options = VariableTracker.propagate(items)
    self.push(TupleVariable(items, **options))

Вот что делает этот код:

1. Функция считывает argval, что в данном случае аналогично counts в pydoc для эквивалентной инструкции.
2. Функция popn элементы, в данном случае подпись — def  popn(self, n: int) -> List[TensorVariable]:, что намекает на лежащий в основе договор — мы возвращаем TensorVariables. Если мы более внимательно посмотрим на symbolic_convert.py и InstructionTranslatorBase/InstructionTranslator, мы увидим, что единственное, что добавляется на наш стек и удаляется с него, — это VariableTracker.

3. Функция вызывает VariableTracker.propagate. Это берет охранники из каждого отдельного элемента, извлеченного со стека в 2, рекурсивно обходит его и объединяет все охранники в options: py  return {      "guards": guards,  }
4. Затем функция создает новый экземпляр VariableTracker, TupleVariable из items и options. Это позволяет нам установить все соответствующие охранники из items , которые составляют новый TupleVariable

После всего этого происходит перевод байткода, и мы приблизились к созданию GuardedCode. Теперь мы понимаем, как локальные переменные становятся VariableTracker , как обрабатываются инструкции и где вызываются охранники для создания. Прежде чем мы сможем увидеть, как код и охранники объединяются в объект GuardedCode, нам нужно немного углубиться в эти make_guard и source.make_guard вызовы выше. Тогда мы сможем понять, что происходило, когда мы создавали охранников вместе с экземплярами VariableTracker.

Создание охранников

Охранники — это просто объекты Python класса Guard. Давайте рассмотрим их подробнее.

Рассмотрев определение класса данных (и, следовательно, подпись конструктора), мы видим, что у него есть имя, источник и функция создания.

@dataclasses.dataclass
class Guard:
    name: str
    source: GuardSource
    create_fn: Callable

Имя должно соответствовать имени переменной.

Здесь источник — это перечисление, указывающее на то, к какому типу источника относится охранник.

create_fn предоставляет основную функциональность для перехода от простого класса данных к фактическому созданию действительного кода Python, который будет вызываться для определения того, изменились ли данные между вызовами и можно ли безопасно считывать из кэша кода или нет.

Наиболее распространенные пути получения экземпляра охранника — через make_guards в VariableTracker. make_guards -> source.make_guard -> return Guard(self.name(), self.guard_source(), fn)

Или, в конкретном примере:

...
elif istype(value, range):
    guards = self.make_guards(GuardBuilder.EQUALS_MATCH)
    return RangeVariable(value=value, guards=guards)

Поскольку source было установлено во время создания этого VariableTracker, все, что здесь требовалось, — это передать fn, GuardBuilder.EQUALS_MATCH в поле create_fn.

Этот create_fn должен быть методом GuardBuilder. Причина этого становится ясной на нашем следующем шаге. После создания всех охранников для фрейма мы переходим к CheckFunctionManager и compile_check_fn.

Прежде чем функция convert_frame сможет создать GuardedCode, ей нужно выполнить CheckFunctionManager, со всеми охранниками, чтобы создать check_fn, которое затем будет передано вместе с кодом в GuardedCode. Это тот же check_fn, который мы храним в записи кэша, и тот же, который мы выполняем, чтобы узнать, нужно ли извлекать сохраненный рядом код. Для справки, вот этот код:

static CacheEntry *create_cache_entry(CacheEntry *next,
                                      PyObject *guarded_code) {
  CacheEntry *e = (CacheEntry *)malloc(sizeof(CacheEntry));
  DEBUG_NULL_CHECK(e);
  e->check_fn = PyObject_GetAttrString(guarded_code, "check_fn");
  NULL_CHECK(e->check_fn);
  e->code = (PyCodeObject *)PyObject_GetAttrString(guarded_code, "code");
  NULL_CHECK(e->code);
  e->next = next;
  return e;
}

Теперь мы знаем, как используется функция check_fn, кто ее создает и из чего она состоит, но мы еще не знаем, как. Как список объектов Guard превращается в функцию, которую можно выполнить позже?

Сначала мы итерируем эти охранники:

for guard in sorted(guards or [], key=Guard.sort_key):
    if not config.guard_nn_modules and guard.is_nn_module():
        continue
    guard.create(local_builder, global_builder)

Вызов guard.create выполняет функцию create_fn, которую мы установили для класса Guard выше (не путайте его с check_fn, который мы создаем, имена похожи, поэтому может быть немного запутанно). В нашем примере выше наш create_fn — это GuardBuilder.EQUALS_MATCH. Таким образом, мы вызываем его, передавая self, сам охранник, внутрь.

Подпись: def EQUALS_MATCH(self, guard: Guard):

И внутри этой функции мы можем использовать name на охраннике, чтобы получить наш исходный объект, запросить данные и информацию о типе, что, в свою очередь, приводит нас к самой важной части: добавлению кода.

В самом простом случае EQUALS_MATCH добавляет всего одну строку кода: self.code.append(f"{ref} == {val!r}"). Где ref — имя переменной, а val — значение. Это может привести к коду такого вида:

y == 2

Это базовый пример. Но если мы добавим несколько других функций GuardBuilder и затем объединим их все с and между каждой инструкцией (как мы это делаем), мы можем получить что-то вроде этого:

___guarded_code.valid and ___check_type_id(y, 94367738391392) and y == 2 and ___check_tensors(x)

Вот что выполняет этот код:

1. Проверка .valid
2. Проверка идентификатора типа
3. Проверка значения
4. Проверка тензора

Это становится ядром кода нашего check_fn, который в свою очередь оценивается в следующий раз, когда мы встречаем этот код. Затем он проверит:

1. Является ли этот код по-прежнему действительным?
2. Если (1), имеет ли y тип 94367738391392?
3. Если (2), равно ли y 2?
4. Если (3), проверим, изменился ли тензор x какими-то определенными способами.

Если все это по-прежнему верно, мы можем использовать кэшированный код вместе с этим check_fn.

Если нет, мы можем перейти к повторной компиляции кода и сохранению этого кода в кэше вместе с этим кодом и совершенно новым check_fn, снова для проверки в последующем фрейме.

Существует множество других таких функций в GuardBuilder, которые объединяются в, порой, огромные строки, которые затем оцениваются как код Python и сохраняются в check_fn. Приведенный выше пример иллюстрирует простой случай. Чтобы лучше понять эту функциональность, прочитайте другие функции в GuardBuilder, или еще лучше, выведите переменную code в compile_check_fn , чтобы увидеть, что создается, особенно в больших реальных моделях.

Заключение

В этом разделе мы рассмотрели:

• Роль .valid и проверки действительности по слабым ссылкам (и, возможно, в ближайшее время проверку действительности NN Module).
• Как работает C++ часть функций охранников (___check_type_id, ___check_tensors, и т. д.).
• Что происходит, когда охранники терпят неудачу.
• Что происходит, если мы создаем недействительный код охранника.

Мы рассмотрели, как пользовательский код, заключенный в контекст TorchDynamo, прослеживается и отслеживается внутри, организованный в VariableTracker Source и, впоследствии, Guard, и как эти Guards , в свою очередь, управляют выбором записи кэша и проверкой действительности при обработке кода Python.