Символические выражения и карты

Символическое выражение ( SymbolicExpr ) — это система математической абстракции, позволяющая выполнять символические тензорные вычисления. Символическая карта ( SymbolicMap ) — это набор символических выражений, которые математически представляют тензорные отображения и преобразования в конвейере компиляции. Они выступают в качестве «математического моста» между высокоуровневой операцией HLO и фактическими адресами памяти, к которым обращается графический процессор/центральный процессор.

SymbolicExpr и SymbolicMap — это собственные реализации XLA, заменяющие устаревшие mlir::AffineExpr и mlir::AffineMap .

SymbolicExpr

SymbolicExpr представляет собой узел в абстрактном синтаксическом дереве (AST). В отличие от стандартного аффинного выражения MLIR, он поддерживает более широкий спектр операций, необходимых для современной обработки изображений на графических процессорах.

Пример SymbolicExpr : d0 + s0 * 8

Поддерживаемые типы

Константы: фиксированные целочисленные значения ( int64 ).

Переменные: Символические переменные с размерами (d) и символами (s) . Все переменные (размеры и символы) рассматриваются как VariableID , и интерпретация переменной зависит от ее контекста в рамках Символической карты.

Операции: add , mul mod , floorDiv , ceilDiv , min , max .

Поддерживаемые операторы: + , - , * , / ( floorDiv ), % ( mod )

Пример использования

v0 = CreateSymbolicVariable(0, context); // 0 is the var_id
v1 = CreateSymbolicVariable(1, context);

SymbolicExpr expr = (((v0 + 42) * v1.min(2).max(0)) / 2).ceilDiv(2);

int64_t result = expr.Evaluate({5, 1}); // Result: 12

Основные характеристики

• Неизменяемость : Символические выражения представляют собой указатели на внутреннее хранилище, управляемое mlir::MLIRContext . Они автоматически дедуплицируются для обеспечения уникальности.

• Канонизация ( Canonicalize() ) : Символические выражения алгебраически упрощаются (с помощью свертывания констант, устранения тождественности, ассоциативного свойства, дистрибутивного свойства и т. д.), чтобы гарантировать представление выражений в стандартной, минимальной форме. Например, (d0 + 1) - 1 будет упрощено до d0 . Это крайне важно для предсказуемых проверок равенства ( operator== ) в выражениях, которые математически эквивалентны, но структурно различны.

SymbolicMap

SymbolicMap представляет собой математическое отображение преобразования между системами координат, как правило, между входными и выходными тензорами.

Пример SymbolicMap : (d0, d1)[s0, s1] -> (d0 + s0, d1 * s1)

Пример использования

SymbolicMap map = SymbolicMap::Get(
    context,
    2,                   // number of dimensions
    1,                   // number of symbols
    {d0 + s0, d1 * s1}); // SymbolicExprs

// Access components
int64_t num_dims = map.GetNumDims();
int64_t num_symbols = map.GetNumSymbols();
auto results = map.GetResults();

Ключевые операции

Замена переменных ( ReplaceDimsAndSymbols() ) : Размеры и/или символы карты могут быть заменены другими выражениями, что позволяет переназначать координатные пространства. Например:

// c2 and c3: SymbolicConstants
// sample_map: (d0, d1)[s0, s1] -> (d0 + s0, d1 * s1)

sample_map.ReplaceDimsAndSymbols(
      {d1, c2},  // New dimensions: Replace d0 with d1, d1 with c2
      {c3, d0},  // New symbols: Replace s0 with c3, s1 with d0
      2,         // New number of dimensions
      2)         // New number of symbols

// Result: (d1 + c3, c2 * d0)

Композиция ( Compose() ) : Можно комбинировать карты с совместимыми размерами. Это позволяет объединять преобразования в цепочки и формирует основной механизм для слияния нескольких операций HLO в единое вычисление индексации. Например:

map1: (d0, d1)[s0] -> (d0 + s0, d1 * 2)
map2: (d0)[s0] -> (d0 - 10, d0 + s0)
map1.compose(map2): (new_d0)[new_s0_map1, new_s0_map2] -> ((new_d0 - 10) + new_s0_map1, (new_d0 + new_s0_map2) * 2)

Оптимизация ( CompressDims() / CompressSymbols() ) : Упрощает карты, выявляя и удаляя неиспользуемые переменные, сохраняя при этом максимально лаконичную логику индексирования. Например:

map1: (d0, d1, d2)[s0] -> (d0 + d2, s0 * 5) // d1 is unused
map1.CompressDims(): (new_d0, new_d1)[new_s0] -> (new_d0 + new_d1, new_s0 * 5)

map2: (d0)[s0, s1, s2] -> {d0 + s2, s0 * 5} // s1 is unused
map2.CompressSymbols(): (new_d0)[new_s0, new_s1] -> (new_d0 + new_s1, new_s0 * 5)

SymbolicMap составляет математическую основу для IndexingMap , которая описывает, как элементы тензора отображаются друг на друга в семантике HLO .

IndexingMap состоит из символических карт с ограничениями, специфичными для предметной области. Он позволяет проводить анализ формы и разбиения на блоки, а также объединять цепочки операций (такие как последовательное изменение формы, транспонирование, широковещательная рассылка и т. д.) в оптимизированные вычисления индексирования. Подробнее см. конкретные примеры в разделе «Анализ индексирования» .