Adept | Машинное обучение | Открытое ПО

Более глубокое изучение вопроса с использованием OpenXLA для компиляции в SPIR-V показало что это напрямую сейчас всё еще не возможно. Только Intel занимается этим вопросом и у них есть какие-то экспериментальные наработки. Но в тоже время есть тензорный компилятор и среда исполнения IREE построенные на базе LLVM MLIR специально для компиляции и запуска SPIR-V. Но это полноценное end2end решение, т.е. ему на вход можно дать HLO и получить исполняемый файл. Вроде бы можно использовать и как библиотеку. Т.е. если использовать IREE то придется отказаться уже разработанной в Adept среды для запуска compute shaders в Vulkan. Так поступать не хочется. Поэтому следующим шагом буду изучать как напрямую использовать LLVM MLIR для понижения из HLO в набор SPIR-V с планом исполнения и промежуточных выделений памяти. А IREE можно использовать для сравнения производительности. #IR #OpenXLA #IREE #SPIR-V #MLIR #LLVM

Сделал детальное представление вычислительного графа и теперь его можно преобразовать в другое IR. Вот например для такого небольшого перцептрона: class MLP : public Module { public: MLP() : linear1(28 * 28, 1024), linear2(1024, 512), linear3(512, 256), linear4(256, 10) { ... } Tensor forward(Tensor input) { auto out = l1(input); out = relu(out); out = l2(out); out = relu(out); out = l3(out); out = relu(out); out = l4(out); return out; } ... }; Получается такое представление в StableHLO формате: module { func.func @main(%arg0: tensor<784x784xf32>, %arg1: tensor<1024x784xf32>, %arg2: tensor<1024xf32>, %arg3: tensor<512x1024xf32>, %arg4: tensor<512xf32>, %arg5: tensor<256x512xf32>, %arg6: tensor<256xf32>, %arg7: tensor<10x256xf32>, %arg8: tensor<10xf32>) -> tensor<784x10xf32> { %0 = stablehlo.transpose %arg1, permutation = [1, 0] : (tensor<1024x784xf32>) -> tensor<784x1024xf32> %1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0] : (tensor<784x784xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32> %2 = stablehlo.broadcast_in_dim %arg2, dims = [1] : (tensor<1024xf32>) -> tensor<784x1024xf32> %3 = stablehlo.add %1, %2 : tensor<784x1024xf32> %4 = stablehlo.constant dense<0> : tensor<784x1024xf32> %5 = stablehlo.maximum %3, %4 : tensor<784x1024xf32> %6 = stablehlo.transpose %arg3, permutation = [1, 0] : (tensor<512x1024xf32>) -> tensor<1024x512xf32> %7 = stablehlo.dot_general %5, %6, contracting_dims = [1] x [0] : (tensor<784x1024xf32>, tensor<1024x512xf32>) -> tensor<784x512xf32> %8 = stablehlo.broadcast_in_dim %arg4, dims = [1] : (tensor<512xf32>) -> tensor<784x512xf32> %9 = stablehlo.add %7, %8 : tensor<784x512xf32> %10 = stablehlo.constant dense<0> : tensor<784x512xf32> %11 = stablehlo.maximum %9, %10 : tensor<784x512xf32> %12 = stablehlo.transpose %arg5, permutation = [1, 0] : (tensor<256x512xf32>) -> tensor<512x256xf32> %13 = stablehlo.dot_general %11, %12, contracting_dims = [1] x [0] : (tensor<784x512xf32>, tensor<512x256xf32>) -> tensor<784x256xf32> %14 = stablehlo.broadcast_in_dim %arg6, dims = [1] : (tensor<256xf32>) -> tensor<784x256xf32> %15 = stablehlo.add %13, %14 : tensor<784x256xf32> %16 = stablehlo.constant dense<0> : tensor<784x256xf32> %17 = stablehlo.maximum %15, %16 : tensor<784x256xf32> %18 = stablehlo.transpose %arg7, permutation = [1, 0] : (tensor<10x256xf32>) -> tensor<256x10xf32> %19 = stablehlo.dot_general %17, %18, contracting_dims = [1] x [0] : (tensor<784x256xf32>, tensor<256x10xf32>) -> tensor<784x10xf32> %20 = stablehlo.broadcast_in_dim %arg8, dims = [1] : (tensor<10xf32>) -> tensor<784x10xf32> %21 = stablehlo.add %19, %20 : tensor<784x10xf32> return %21 : tensor<784x10xf32> } } Далее его надо преобразовать тензорным компиляторм. Эксперименты продолжаются.

После публикации релиза продолжаем работу: - Расширяем набор демонстрационных приложений для прикладного API. Работаем над реализацией моделей на основе архитектуры UNet и соответствующего конвейера для их обучения. - Продолжаем работу над реализацией поддержки формата GGUF. Сначала будет загрузка квантованных весов с распаковкой в FP32 и API для квантования с использованием ggml. Далее планируем делать поддержку квантованных данных в вычислительных ядрах. - Будем пробовать увеличить производительность реализации GEMM для Vulkan. - И я затеял глубокий рефакторинг системы построения вычислительного графа и авто-дифференцирования. Будем уходить от простого подхода с неявным построением графа к явному описанию с узлами, рёбрами и соответствующей дополнительной информацией. Это должно позволить осуществлять преобразования графа в другие промежуточные представления(IR) и производить его анализ и оптимизацию. На первом этапе хочется получить функциональный слой для всех операций и возможность генерации StableHLO IR для использования тензорного компилятора OpenXLA. Прикладной API это не затронет и даже упростит, так как уберем класс Variable и будем оперировать только тензорами. #plans #IR #GGUF #tensor #compiler #StableHLO