Спецификация IR torch.export

call_function

Узел call_function представляет вызов оператора.

Определения

• Функциональный: вызываемый объект считается «функциональным», если он удовлетворяет всем следующим требованиям:

  • Не изменяет данные: оператор не изменяет значение своих входных данных (для тензоров это относится как к метаданным, так и к данным).
  • Не имеет побочных эффектов: оператор не изменяет состояния, видимые извне, например значения параметров модуля.
• Оператор: функциональный вызываемый объект с заранее определенной схемой. Примерами таких операторов являются функциональные операторы ATen.

Представление в FX

%name = call_function[target = operator](args = (%x, %y, …), kwargs = {})

Отличия от стандартного call_function в FX

1. В графе FX call_function может ссылаться на любой вызываемый объект, тогда как в Export IR он ограничен выбранным набором операторов ATen, пользовательских операторов и операторов потока управления.
2. В Export IR аргументы-константы встраиваются в граф.
3. В графе FX узел get_attr может представлять чтение любого атрибута, хранящегося в модуле графа. В Export IR чтение ограничено только вложенными модулями, поскольку все параметры и буферы передаются в модуль графа как входные данные.

placeholder

Placeholder представляет вход графа. Его семантика полностью совпадает с семантикой в FX. Узлы placeholder должны быть первыми N узлами в списке узлов графа. N может быть равен нулю.

Представление в FX

%name = placeholder[target = name](args = ())

Поле target представляет собой строку с именем входных данных.

args, если оно не пустое, должно иметь размер 1 и представлять значение входных данных по умолчанию.

Метаданные

Узлы placeholder также содержат meta[‘val’], как и узлы call_function. В этом случае поле val описывает форму и тип данных входного параметра, которые граф ожидает получить.

output

Вызов output представляет оператор return в функции и тем самым завершает текущий граф. Узел output существует в единственном экземпляре и всегда является последним узлом графа.

Представление в FX

output[](args = (%something, …))

Его семантика полностью совпадает с семантикой в torch.fx. args обозначает возвращаемый узел.

Метаданные

Узел output содержит те же метаданные, что и узлы call_function.

get_attr

Узлы get_attr представляют чтение вложенного модуля из содержащего его torch.fx.GraphModule. В отличие от стандартного графа FX, созданного с помощью torch.fx.symbolic_trace(), где узлы get_attr используются для чтения таких атрибутов, как параметры и буферы, из верхнеуровневого torch.fx.GraphModule, в Export IR параметры и буферы передаются в модуль графа как входные данные и хранятся в верхнеуровневом torch.export.ExportedProgram.

Представление в FX

%name = get_attr[target = name](args = ())

Пример

Рассмотрим следующую модель:

from functorch.experimental.control_flow import cond

def true_fn(x):
    return x.sin()

def false_fn(x):
    return x.cos()

def f(x, y):
    return cond(y, true_fn, false_fn, [x])

Граф:

graph():
    %x_1 : [num_users=1] = placeholder[target=x_1]
    %y_1 : [num_users=1] = placeholder[target=y_1]
    %true_graph_0 : [num_users=1] = get_attr[target=true_graph_0]
    %false_graph_0 : [num_users=1] = get_attr[target=false_graph_0]
    %conditional : [num_users=1] = call_function[target=torch.ops.higher_order.cond](args = (%y_1, %true_graph_0, %false_graph_0, [%x_1]), kwargs = {})
    return conditional

Строка %true_graph_0 : [num_users=1] = get_attr[target=true_graph_0] считывает вложенный модуль true_graph_0, содержащий оператор sin.

SymInt

SymInt — это объект, который может быть либо целочисленным литералом, либо символом, представляющим целое число (в Python он представлен классом sympy.Symbol). Если SymInt является символом, он описывает переменную целочисленного типа, неизвестную графу во время компиляции; ее значение становится известно только во время выполнения.

FakeTensor

FakeTensor — это объект, содержащий метаданные тензора. Его метаданные можно представить следующим образом.

class FakeTensor:
  size: List[SymInt]
  dtype: torch.dtype
  device: torch.device
  dim_order: List[int]  # This doesn't exist yet

Поле size объекта FakeTensor — это список целых чисел или SymInt. Если в нем присутствуют SymInt, это означает, что тензор имеет динамическую форму. Если присутствуют целые числа, предполагается, что тензор будет иметь именно такую статическую форму. Ранг TensorMeta не может быть динамическим. Поле dtype задает тип данных результата этого узла. В Edge IR нет неявного преобразования типов. FakeTensor не содержит сведений о шагах.

Иными словами:

• Если оператор в node.target возвращает Tensor, то node.meta['val'] — это FakeTensor, описывающий этот тензор.
• Если оператор в node.target возвращает n-кортеж тензоров, то node.meta['val'] — это n-кортеж объектов FakeTensor, описывающих каждый тензор.
• Если оператор в node.target возвращает int/float/скаляр, значение которого известно во время компиляции, то node.meta['val'] равно None.
• Если оператор в node.target возвращает int/float/скаляр, значение которого неизвестно во время компиляции, то node.meta['val'] имеет тип SymInt.

Например:

• aten::add возвращает Tensor, поэтому его спецификация будет объектом FakeTensor с dtype и size возвращенного этим оператором тензора.
• aten::sym_size возвращает целое число, поэтому его val будет иметь тип SymInt, поскольку значение доступно только во время выполнения.
• max_pool2d_with_indexes возвращает кортеж (Tensor, Tensor), поэтому спецификация также будет состоять из 2-кортежа объектов FakeTensor; первый TensorMeta описывает первый элемент возвращаемого значения и т. д.

Код Python:

def add_one(x):
  return torch.ops.aten(x, 1)

Граф:

graph():
  %ph_0 : [#users=1] = placeholder[target=ph_0]
  %add_tensor : [#users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%ph_0, 1), kwargs = {})
  return [add_tensor]

FakeTensor:

FakeTensor(dtype=torch.int, size=[2,], device=CPU)

Типы, совместимые с Pytree

Тип называется «совместимым с Pytree», если он является либо листовым типом, либо типом-контейнером, содержащим другие типы, совместимые с Pytree.

Примечание:

Понятие pytree совпадает с описанным здесь для JAX:

Следующие типы определены как листовые типы:

Следующие типы определены как типы-контейнеры: